Training and testing scripts for the prediction model used in the "Interaction-Aware Sampling-Based MPC with Learned Local Goal Predictions" paper.
import sys
import os

sys.path.append('../')

# With debugger:
#os.chdir(os.getcwd() + "/src/")

import numpy as np
import argparse
import pylab as pl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pickle as pkl
import importlib
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import random
import progressbar
import scipy.io as sio
import json
from copy import deepcopy
import colorama
from colorama import Fore, Style
import json

if sys.version_info[0] < 3:
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from Performance import *
	from utils import *
	import Recorder as rec
else:
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.Performance import *
	from src.data_utils.utils import *
	from src.data_utils.Recorder import Recorder as rec
	from src.models.CV_model import ConstantVelocity


# Model directories
def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
											help='Path to directory that comprises the model (default="model_name").',
											type=str, default="VGDNN")
	parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=33)
	parser.add_argument('--n_samples', help='Number of samples', type=int, default=10)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
											type=str, default="datasets/ewap_dataset/seq_eth")
	parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool,
											default=True)
	parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool,
											default=True)
	parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float,
											default=0)
	parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float,
											default=5)
	parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool,
											default=False)
	parser.add_argument('--update_state', help='update_state', type=sup.str2bool,
											default=False)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
											default=False)
	parser.add_argument('--freeze_other_agents', help='FReeze other agents.', type=sup.str2bool,
											default=False)
	parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,
											default=False)
	parser.add_argument('--constant_velocity', help='Run CV comparison', type=sup.str2bool,
											default=False)
	args = parser.parse_args()

	return args


test_args = parse_args()

if test_args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

cwd = os.getcwd()

model_path = os.path.normpath(cwd + '/../') + '/trained_models/' + test_args.model_name + "/" + str(test_args.exp_num)

print("Loading data from: '{}'".format(model_path))
file = open(model_path + '/model_parameters.pkl', 'rb')
if sys.version_info[0] < 3:
	model_parameters = pkl.load(file)  # ,encoding='latin1')
else:
	model_parameters = pkl.load(file, encoding='latin1')
file.close()
args = model_parameters["args"]

with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

print("Training steps: ", args.total_training_steps)

# change some args because we are doing inference
# TODO: Why change these arguments, how to infer for multiple agents simultaneously?
truncated_backprop_length = args.truncated_backprop_length
args.truncated_backprop_length = 1
args.batch_size = 1
args.keep_prob = 1.0
#args.others_info = "none"#"relative"
# args.data_path = "../data/2_agents_swap/trajs/"
# args.scenario = "GA3C-CADRL-10-py27"
# args.dataset = args.scenario+'.pkl'
# args.data_path = '../data/'

training_scenario = args.scenario
args.scenario = test_args.scenario
args.real_world_data = test_args.real_world_data
args.dataset = '/' + args.scenario + '.pkl'
data_prep = dhlstm.DataHandlerLSTM(args)

# for comparison
#args.scenario = "comparison"

# Only used to create a map from png
# Load Map Parameters
map_params = os.path.join(args.data_path+args.scenario, 'map.json')
with open(map_params) as json_file:
	data = json.load(json_file)
map_args = {"file_name": data["file_name"],
						"resolution": data["resolution"],
						"map_size": np.array(data["map_size"]),
						"map_center": np.array(data["map_center"])}

data_prep.processData(**map_args)
if args.normalize_data:
	data_prep.compute_min_max_values()

# Import model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

model = NetworkModel(args)

# Lists for logging of the input / output data of the model
input_list = []
grid_list = []
goal_list = []
ped_grid_list = []
y_ground_truth_list = []
y_pred_list = []  # uses ground truth as input at every step
other_agents_list = []
all_predictions = []
all_traj_likelihood = []
trajectories = []
batch_y = []
batch_loss = []

config = tf.ConfigProto(
	device_count={'GPU': 0}
)


def clear_plot():
	axarr[0].clear()
	axarr[0].set_title('Training Sample')

if test_args.unit_testing:
	data_handler = dhlstm.DataHandlerLSTM(args)
	data_handler.unit_test_data_(map_args)

with tf.Session(config=config) as sess:
	model.warmstart_model(args, sess)
	try:
			model.warmstart_convnet(args, sess)
	except:
		print("")

	for exp_id in range(np.minimum(test_args.num_test_sequences,len(data_prep.trajectory_set)-1)):
		predictions = []
		traj_likelihood = []
		# sample a trajectory id for testing
		traj_id = random.randint(0, len(data_prep.trajectory_set) - 1)

		#print(args.others_info)

		batch_x, batch_vel, batch_pos,batch_goal, batch_grid, other_agents_info, batch_target,batch_end_pos, other_agents_pos, traj = data_prep.getTrajectoryAsBatch(
			traj_id,freeze = test_args.freeze_other_agents)  # trajectory_set random.randint(0, len(data_prep.dataset) - 1)
		
		#batch_x, batch_vel, batch_pos,batch_goal, batch_grid, other_agents_info, batch_target,batch_end_pos, other_agents_pos, traj = data_prep.getGroupOfTrajectoriesAsBatch(traj_id)

		trajectories.append(traj)
		x_input_series = np.zeros([0, (args.prev_horizon + 1) * args.input_dim])
		goal_input_series = np.zeros([0, 2])
		grid_input_series = np.zeros(
			[0, int(args.submap_size), int(args.submap_size)])
		if args.others_info == "relative":
			ped_grid_series = np.zeros([0, args.n_other_agents,args.pedestrian_vector_dim])
		elif "sequence" in args.others_info:
			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim*args.prediction_horizon])
		elif args.others_info == "prev_sequence":
			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim *(args.prev_horizon+1)])
		elif args.others_info == "sequence2":
			ped_grid_series = np.zeros([0, args.n_other_agents, args.prediction_horizon,args.pedestrian_vector_dim])
		else:
			ped_grid_series = np.zeros([0, args.pedestrian_vector_dim])
		y_ground_truth_series = np.zeros([0, args.prediction_horizon * 2])
		y_pred_series = np.zeros([0, args.n_mixtures * args.prediction_horizon * args.output_pred_state_dim])

		batch_y.append(batch_target)
		model.reset_test_cells(np.ones((args.batch_size)))
		cell_state_list= []
		cell_ped_list = []
		cell_concat_list = []
		if "grid" in args.model_name:
			batch_ped_grid_backup = np.zeros_like(batch_grid)
			data_prep.add_other_agents_to_grid(batch_ped_grid_backup, batch_x, [other_agents_pos])

		for step in range(batch_x.shape[1]):
			samples = []
			# Assemble feed dict for training
			if "future" in args.others_info:
				if step == 0:
					batch_y_pred = deepcopy(batch_vel)
					for pred_step in range(1, args.prediction_horizon):
						batch_y_pred[:, :, pred_step * 2:pred_step * 2 + 2] = batch_vel[:, :, :2]
				else:
					batch_y_pred = deepcopy(batch_vel)
					batch_y_pred[:, :, 2:] = y_model_pred[:, :, 2:]

			dict = {
							"batch_vel": batch_vel,
							"batch_grid": batch_grid,
							"batch_ped_grid": other_agents_info,
							"step": step,
							"other_agents_pos": [other_agents_pos]
			}

			feed_dict_ = model.feed_test_dic(
                **dict
            )


			# Append to logging series
			x_input_series = np.append(x_input_series, batch_x[:, step, :], axis=0)
			grid_input_series = np.append(grid_input_series, batch_grid[:, step, :, :], axis=0)
			goal_input_series = np.append(goal_input_series, batch_goal[:, step, :], axis=0)

			y_ground_truth_series = np.append(y_ground_truth_series, batch_target[:, step, :], axis=0)

			y_model_pred, likelihood = model.predict(sess, feed_dict_, True)

			# Backup cell states for later analysis
			cell_state_list.append(model.test_cell_state_current[0,:])
			#cell_ped_list.append(model.test_cell_state_current_lstm_ped[0, :])
			cell_concat_list.append(model.test_cell_state_current_lstm_concat[0, :])

			# Rotate predictions to global frame
			if args.rotated_grid:
				heading = math.atan2(batch_vel[0, step, 1], batch_vel[0, step, 0])
				rot_mat = np.array([[np.cos(heading), -np.sin(heading)], [np.sin(heading), np.cos(heading)]])
				for pred_step in range(args.prediction_horizon):
					y_model_pred[0][0, 2 * pred_step:2 * pred_step + 2] = np.dot(rot_mat, y_model_pred[0][0,2 * pred_step:2 * pred_step + 2])
				samples.append(y_model_pred[0])
			else:
				samples.append(y_model_pred[:,0,:])

			# If sample more than one trajectory from the model
			for sample_id in range(test_args.n_samples - 1):
				dict = {"batch_x": batch_x,
								"batch_vel": batch_vel,
								"batch_pos": batch_pos,
								"batch_grid": batch_grid,
								"batch_ped_grid": other_agents_info,
								"step": step,
								"batch_goal": batch_goal,
								"state_noise": 0.0,
								"grid_noise": 0.0
								}
				feed_dict_ = model.feed_test_dic(**dict)
				y_model_pred = model.predict(sess, feed_dict_, test_args.update_state)
				samples.append(y_model_pred[:,0,:])

			traj_likelihood.append(likelihood)
			predictions.append(samples)

		latent_analysis = False

		# Plot hidden states evolution over a trajectory
		if latent_analysis:

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			pl.show(block=False)
			x_state = np.arange(0, args.rnn_state_size, 1)
			x_ped = np.arange(0, args.rnn_state_size_lstm_ped * 2, 1)
			x_concat = np.arange(0, args.rnn_state_size_lstm_concat, 1)

			time = np.arange(0, len(cell_state_list), 1)
			R= np.array(cell_state_list)
			ax.clear()
			X, Y = np.meshgrid(x_state,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			save_img_to_file = model_path + "/figs"
			if not os.path.exists(save_img_to_file):
				os.makedirs(save_img_to_file)
			pl.savefig(save_img_to_file+"/cell_state" + str(exp_id) + '.png')
			pl.close()

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			time = np.arange(0, len(cell_ped_list), 1)
			R= np.array(cell_ped_list)
			#print(R)
			ax.clear()
			X, Y = np.meshgrid(x_ped,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			pl.savefig(save_img_to_file+"/cell_ped" + str(exp_id) + '.png')
			pl.close()

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			time = np.arange(0, len(cell_concat_list), 1)
			R= np.array(cell_concat_list)
			ax.clear()
			X, Y = np.meshgrid(x_concat,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			pl.savefig(save_img_to_file+"/cell_concat" + str(exp_id) + '.png')
			pl.close()

		all_predictions.append(predictions)
		all_traj_likelihood.append(traj_likelihood)
		input_list.append(x_input_series)
		goal_list.append(goal_input_series)
		grid_list.append(grid_input_series)
		y_ground_truth_list.append(y_ground_truth_series)
		other_agents_list.append(other_agents_pos)
		# update progress bar

sess.close()

# Save input information for comparison

dict["predictions"] = all_predictions
dict["batch_x"] = batch_x
#with open('comparison4.pkl', 'wb') as f:
#    pkl.dump(dict, f)


# Save data in Matlab format for post-processing
if not os.path.exists(args.model_path + '/../results/'):
	os.makedirs(args.model_path + '/../results/')
results= {
	"trajectories": trajectories,
	"predictions": all_predictions
}
scenario = args.scenario.split('/')[-1]
results_file = args.model_path + '/'+ scenario + "_results.mat"
sio.savemat(results_file,results)

if test_args.constant_velocity:
	# Make constant velocity predictions
	CVModel = ConstantVelocity(args)
	cv_predictions = CVModel.predict(input_list)
else:
	cv_predictions = None
"""

if test_args.record:
		recorder = rec(args, data_prep.agent_container.occupancy_grid)
		if ( "real_world" in test_args.scenario) and not test_args.unit_testing:
			print("Real data!!")
			recorder.plot_on_video(input_list, grid_list, all_predictions, y_ground_truth_list,
														 other_agents_list,
														 trajectories,all_traj_likelihood, test_args)
		else:
			#recorder.plot_on_image(input_list, grid_list, all_predictions, y_ground_truth_list, other_agents_list,
			#	                       trajectories,test_args)
			#recorder.animate_local(input_list, grid_list, ped_grid_list, all_predictions, y_ground_truth_list, other_agents_list,
			#							 trajectories,test_args)

			recorder.animate_global(input_list, grid_list, all_predictions, y_ground_truth_list,
														 other_agents_list,
														 trajectories, all_traj_likelihood, cv_predictions, test_args)

			print("Recorder is done!")
else:
	print("Performance tests")
	pred_error, pred_error_summary_lstm = compute_trajectory_prediction_mse(args, trajectories, all_predictions)
	pred_fde, pred_error_summary_lstm_fde = compute_trajectory_fde(args, trajectories, all_predictions)
	if test_args.constant_velocity: 
		cv_pred_error = compute_ade_cv(args, trajectories, cv_predictions)
		cv_fde_error = compute_fde_cv(args, trajectories, cv_predictions)
		print(
		Fore.LIGHTBLUE_EX + "\nConstant Velocity: MSE: {:01.2f}, FDE: {:01.2f}".format(np.mean(cv_pred_error), np.mean(cv_fde_error))+Style.RESET_ALL)
	diversity, diversity_summary = compute_2_wasserstein(args, all_predictions)
	args.scenario = training_scenario
	args.truncated_backprop_length = truncated_backprop_length
	write_results_summary(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde), np.mean(diversity_summary), args, test_args)
	print(
		Fore.LIGHTBLUE_EX + "\nMSE: {:01.2f}, FDE: {:01.2f}, DIVERSITY: {:01.2f}".format(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde),np.mean(diversity_summary))+Style.RESET_ALL)

"""
build/lib/compare_VGDNN.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import pylab as pl
import sys
if sys.version_info[0] < 3:
  import Trajectory as traj
  import OccupancyGrid as occ
  import AgentData as ad
  import Support as sup
else:
  import src.data_utils.Trajectory as traj
  import src.data_utils.OccupancyGrid as occ
  import src.data_utils.AgentData as ad
  import src.data_utils.Support as sup


class AgentContainer():
  """
  Container for all agents in one demonstration. 
  This data structure provides an interface to easily store and access data of various agents.
  It contains a dictionary of the agent data per agent id and the static occupancy grid.
  """
  def __init__(self):
    self.agent_data = {}
    self.occupancy_grid = occ.OccupancyGrid()
    
  def setOccupancyGrid(self, occupancy_grid):
    self.occupancy_grid = occupancy_grid
  
  def removeAgent(self, id):
    self.agent_data.pop(id)
  
  def addDataSample(self, id, timestamp, pose, vel, goal):
    """
    Add a data sample to the database consisting of agent data per ID.
    """
    # set up new agent if id has not been registered yet
    if id not in self.agent_data.keys():
      self.agent_data[id] = ad.AgentData(id=id)
    
    # add the data to the agent with the right id
    self.agent_data[id].addSample(timestamp, pose, vel, tuple(goal))
  
  
  # Getters and setters
  def getAgentIDs(self):
    """
    Return list of IDs of contained agents.
    """
    return [int(k) for k in self.agent_data.keys()]
  
  def getNumberOfAgents(self):
    """
    Get the number of agents in the container. 
    """
    return len(self.agent_data.keys())
  
  def getAgentTrajectories(self, id):
    """
    Get all trajectories of a certain agent.
    """
    return self.agent_data[id].trajectories
  
  def getAgentData(self, id):
    """
    Return the agent data of agent with ID=id
    """
    return self.agent_data[id]
  
  def getNumberOfTrajectoriesForAgent(self, id):
    """
    Return the number of trajectories (start -> goal) of agent with ID=id.
    """
    return len(self.agent_data[id].trajectories)
  
  def getTrajectorySetForTime(self, query_time):
    """
    Get all trajectories which are available at a certain time.
    """
    traj_set = {}
    for id in self.agent_data.keys():
      t = self.agent_data[id].getTrajectoryForTime(query_time)
      if t != None:
        traj_set[id] = t
    return traj_set
    
  def getAgentPositionsForTimeExclude(self, query_time, agent_id):
    """
    Get all agent positions for certain time.
    query_time: time where the entries should be queried
    agent_id: id to exclude
    """
    pos_array = np.zeros([0, 2])
    cnt = 0
    for id in self.agent_data.keys():
      if id != agent_id:
        t = self.agent_data[id].getTrajectoryForTime(query_time)
        if t != None and len(t) > 1:
          pos_array = np.append(pos_array, np.expand_dims(t.getPoseAtTime(query_time)[0:2], axis=0), axis=0)
          cnt += 1
    return pos_array
  
  def getAgentVelocitiesForTimeExclude(self, query_time, agent_id):
    """
    Get all agent velocities for certain time.
    query_time: time where the entries should be queried
    agent_id: id to exclude
    """
    vel_array = np.zeros([0, 2])
    cnt = 0
    for id in self.agent_data.keys():
      if id != agent_id:
        t = self.agent_data[id].getTrajectoryForTime(query_time)
        if t != None and len(t) > 1:
          vel_array = np.append(vel_array, np.expand_dims(t.getVelAtTime(query_time)[0:2], axis=0), axis=0)
          cnt += 1
    return vel_array
  
  def plot(self, ax, x_scale = 1, y_scale = 1):
    """
    Plots the trajectories and the static occupancy grid of all agents in this container.
    """
    colormap = pl.get_cmap('rainbow')
    c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=len(self.agent_data.keys()))
    scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
    
    for cnt, id in enumerate(self.agent_data.keys()):
      color_value = scalar_color_map.to_rgba(cnt)
      self.agent_data[id].plot(ax, color=color_value,x_scale = 1, y_scale = 1)
      
    #sup.plot_grid(ax, np.array([self.occupancy_grid.center[0], self.occupancy_grid.center[1]]), self.occupancy_grid.gridmap, self.occupancy_grid.resolution, self.occupancy_grid.map_size)
    sup.plot_grid(ax, np.array([0, 0]),self.occupancy_grid.gridmap, self.occupancy_grid.resolution, self.occupancy_grid.map_size)
    ax.set_xlim([-self.occupancy_grid.center[0], self.occupancy_grid.center[0]])
    ax.set_ylim([-self.occupancy_grid.center[1], self.occupancy_grid.center[1]])
    
    ax.set_aspect('equal')
    
  
build/lib/data_utils/AgentContainer.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import sys
if sys.version_info[0] < 3:
  import Trajectory as traj
else:
  import src.data_utils.Trajectory as traj

class AgentData():
  """
  Contains data for one agent, comprising all trajectories the agent traveled within the given map.
  Different trajectories are separated by their different goals.
  """
  
  def __init__(self, id, radius=0.25):
    self.id = id
    self.goals = []
    self.trajectories = []
    self.radius = radius
    self.cached_trajectory_idx = 0
    self.traj_idx = 0
    self.last_goal = 0
    
  def addSample(self, timestamp, pose, vel, goal):
    """
    Automatically appends the input state to the matching trajectory. 
    If this measurement is part of a new trajectory, one will be created.
    """
    # hack to catch situation of changing motion direction
    #if goal[0] < 0:
    #  goal = self.last_goal
    if goal not in self.goals:

        # adding goal identifier
        self.goals.append(goal)
        self.trajectories.append(traj.Trajectory(goal=np.array([[goal[0], goal[1]]])))
        if len(self.goals) > 1:
          self.traj_idx += 1
    else:
      if len(self.goals)>0:
        if not np.all(goal == self.last_goal):
          self.trajectories.append(traj.Trajectory(goal=np.array([[goal[0], goal[1]]])))
          self.traj_idx += 1

    # append data sample to trajectory (timestamp in [ns], pose and vel as np 1x3 arrays)
    self.trajectories[self.traj_idx].addData(timestamp, pose, vel, goal)
    self.last_goal = goal
    
  
  def getTrajectoryForTime(self, query_time):
    """
    Return the matching trajectory for the query time.
    """
    if self.trajectories[self.cached_trajectory_idx].contains(query_time):
      return self.trajectories[self.cached_trajectory_idx]
    elif query_time >= self.trajectories[0].getMinTime() and query_time <= self.trajectories[-1].getMaxTime():
      # Find trajectory that contains the query time
      for idx, t in enumerate(self.trajectories):
        if t.contains(query_time):
          self.cached_trajectory_idx = idx
          return t
    else:
      return None
  
  def plot(self, ax, color='b', x_scale = 1, y_scale = 1):
    for t in self.trajectories:
      ax.plot(t.pose_vec[:,0]*x_scale, t.pose_vec[:,1]*y_scale, color=color)

build/lib/data_utils/AgentData.py
eliatrevisan/IA-MPPI-LBM
#!/usr/bin/env python
import sys
import numpy as np
import matplotlib.pyplot as pl
import pickle as pkl
import tensorflow as tf

if sys.version_info[0] < 3:
	import yaml

	sys.path.append('../src/data_utils')
	sys.path.append('/opt/ros/kinetic/lib/python2.7/dist-packages')
	import Support as sup
	import OccupancyGrid as occ
	import MemoryBuffer
	import rospy
	from geometry_msgs.msg import Pose
	from geometry_msgs.msg import Twist
	from pedsim_msgs.msg import TrackedPersons
	from sensor_msgs.msg import LaserScan
	from nav_msgs.msg import OccupancyGrid
	from nav_msgs.msg import Path
	import rosbag
else:
	sys.path.append('../src/data_utils')
	import src.data_utils.Support as sup
	import src.data_utils.MemoryBuffer as MemoryBuffer
	import src.data_utils.OccupancyGrid as occ
import random
import os

from multiprocessing.pool import ThreadPool
import math
from time import sleep
from parser import *
import argparse

tfk = tf.keras


class DataHandler():

	def __init__(self, args):
		self.args = args
		self.control_command = np.empty([1, 2])
		self.laser_data = np.empty([1, 720])
		self.grid = np.empty([120, 120])
		self.pedestrian = []
		self.trajectory = []
		self.pedestrian_state = []
		self.robot_command_set = []
		self.robot_state = []
		self.robot_pred_cmd_set = []
		self.data = {}
		self.trajectory_set = []
		self.data_idx = 0
		self.dt = 0.1
		self.batch_size = self.args.batch_size
		self.tbpl = self.args.truncated_backprop_length
		self.cmd_vector_dim = self.args.cmd_vector_dim
		self.prev_horizon = self.args.prev_horizon
		self.prediction_horizon = self.args.prediction_horizon
		self.pred_cmd = np.empty([1, self.prediction_horizon * self.cmd_vector_dim])
		self.min_length_trajectory = self.tbpl + 1 + self.prediction_horizon * 2 + args.prev_horizon
		# Batches
		self.batch_cmd = np.zeros((self.batch_size, self.tbpl, self.cmd_vector_dim))
		self.batch_pred_cmd = np.zeros((self.batch_size, self.tbpl, self.prediction_horizon * self.cmd_vector_dim))
		self.batch_robot_state = np.zeros((self.batch_size, self.tbpl, 3))
		self.batch_grid = np.zeros((self.batch_size, self.tbpl, int(self.args.submap_width / self.args.submap_resolution),
		                            int(self.args.submap_height / self.args.submap_resolution)))
		self.batch_x = np.zeros((self.batch_size, self.tbpl, self.args.input_dim*(self.prev_horizon+1)))
		self.batch_vel = np.zeros((self.batch_size, self.tbpl, self.args.input_state_dim*(self.prev_horizon+1)))
		self.batch_pos = np.zeros((self.batch_size, self.tbpl, self.args.input_state_dim*(self.prev_horizon+1)))
		self.batch_goal = np.zeros((self.batch_size, self.tbpl, 2))
		self.other_agents_info = np.zeros((self.batch_size, self.tbpl, self.args.pedestrian_vector_dim))
		self.other_agents_pos = np.zeros((self.batch_size, self.tbpl, self.args.pedestrian_vector_state_dim))
		self.batch_enc = np.zeros((self.batch_size, self.prev_horizon, self.args.input_dim))
		self.batch_dec = np.zeros((self.batch_size, self.prediction_horizon, self.args.output_dim))
		self.batch_target = np.zeros((self.batch_size, self.tbpl, self.args.output_dim * self.prediction_horizon))
		self.batch_sequences = []
		# indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.sequence_reset = np.ones([self.batch_size])
		self.sequence_idx = np.zeros([self.batch_size])+self.args.prev_horizon

		# Replay buffer for oline training
		self.buffer = MemoryBuffer.MemoryBuffer(self.args.min_buffer_size, self.args.max_buffer_size)

		self.max_traj = self.args.max_trajectories

		self.global_fig = pl.figure("occupancy grid ")
		self.ax_grid = pl.subplot()
		"""
		print("Extracting the occupancy grid ...")
		# Occupancy grid data
		self.occupancy_grid = occ.OccupancyGrid()
		temp = self.args.data_path + '/map.npy'
		# Only used to create a map from png
		map_args = {"file_name": self.args.file_name,
								"resolution": self.args.resolution,
								"map_size": np.array([self.args.map_size_x,self.args.map_size_y]), }
		if not os.path.isfile(self.args.data_path + '/map.npy'):
				print("Creating map from png ...")
				sup.create_map_from_png(data_path=self.args.data_path, **map_args)

		map_data = np.load(os.path.join(self.args.data_path, 'map.npy'), encoding='latin1', allow_pickle=True)
		self.occupancy_grid.gridmap = map_data.item(0)['Map']  # occupancy values of cells
		self.occupancy_grid.resolution = map_data.item(0)['Resolution']  # map resolution in [m / cell]
		self.occupancy_grid.map_size = map_data.item(0)['Size']  # map size in [m]
		self.occupancy_grid.center = self.occupancy_grid.map_size / 2.0

		# FIlter all -1 values present on MAP to zero.
		for ii in range(self.occupancy_grid.gridmap.shape[0]):
				for jj in range(self.occupancy_grid.gridmap.shape[1]):
						if self.occupancy_grid.gridmap[ii, jj] == -1:
								self.occupancy_grid.gridmap[ii, jj] = 0.0
		"""
		if sys.version_info[0] < 3:
			config = yaml.load(open(self.args.topics_config))  # read the yaml file and get the required parameters

			self.robot_topic = config['robot_state_topic']
			self.pedestrian_topic = config['pedestrian_topic']
			self.commmand_topic = config['command_topic']
			self.laser_topic = config['laser_topic']
			self.grid_topic = config['grid_topic']
			self.predict_topic = config['predict_topic']
			self.rate = config['rate']
			rospy.Subscriber(self.robot_topic, Pose, self.RobotStateCB)  # subscribe to various topics
			rospy.Subscriber(self.commmand_topic, Twist, self.ControlCommandCB)
			rospy.Subscriber(self.laser_topic, LaserScan, self.LaserDataCB)
			rospy.Subscriber(self.pedestrian_topic, TrackedPersons, self.TrackedPersonsCB)
			rospy.Subscriber(self.grid_topic, OccupancyGrid, self.GridProcessingCB)
			rospy.Subscriber(self.predict_topic, Path, self.PredictedCommandCB)

	def process_simulation_data_(self):
		# load data from csv for offline training

		# read the required csv files
		if os.path.exists(self.args.data_path):

			if os.path.exists(os.path.join(self.args.data_path, self.args.bag_file)):
				bag = rosbag.Bag(os.path.join(self.args.data_path, self.args.bag_file))

				# variables for storing robot control commands
				linear_vel = np.empty([1])
				angular_vel = np.empty([1])

				# get the control command data
				data_idx = 0
				for topic, msg, t in bag.read_messages(topics=['jackal_velocity_controller/cmd_vel']):
					linear_vel[:] = msg.linear.x
					angular_vel[:] = msg.angular.z
					control_command = np.concatenate([linear_vel, angular_vel])  # [1,2]
					self.robot_command_set.append(control_command.copy())
					data_idx += 1

				# convert from list to numpy array
				self.robot_command_set = np.array(self.robot_command_set)

				for topic, msg, t in bag.read_messages(topics=['/predicted_cmd']):
					self.args.prediction_cmd_horizon = len(msg.poses)
					cmd_pred = np.zeros(self.args.prediction_cmd_horizon * 2)
					for n in range(self.args.prediction_cmd_horizon):
						linear_vel = msg.poses[n].pose.position.x
						angular_vel = msg.poses[n].pose.position.y

						if math.isnan(angular_vel):
							angular_vel = 0

						if math.isnan(linear_vel):
							linear_vel = 0
						cmd_pred[2 * n] = linear_vel
						cmd_pred[2 * n + 1] = angular_vel
					self.robot_pred_cmd_set.append(cmd_pred.copy())

				# convert from list to numpy array
				self.robot_pred_cmd_set = np.asanyarray(self.robot_pred_cmd_set)

				# variables storing robot pose data
				position = np.empty([2])
				orientation = np.empty([1])
				self.robot_state = []

				for topic, msg, t in bag.read_messages(topics=['/robot_state']):
					position[:] = [msg.position.x, msg.position.y]
					orientation[:] = msg.orientation.z
					robot_pose = np.concatenate([position, orientation])  # [1,3]
					self.robot_state.append(robot_pose)

				# convert from list to numpy array
				self.robot_state = np.asanyarray(self.robot_state)

				# check the number of times the robot reaches the goal state, hence providing the number of trajectories.
				self.robot_trajectory_number = 0
				for i in range(1, self.robot_state.shape[0]):
					dist = np.linalg.norm(self.robot_state[i, :2] - self.robot_state[i - 1, :2])
					if dist > self.args.distance_threshold:  # trajectory completion condition
						self.robot_trajectory_number += 1

				print("total trajectories: {}".format(self.robot_trajectory_number))

				for topic, msg, t in bag.read_messages(topics=['/pedsim_visualizer/tracked_persons']):
					# number of pedestrians in recorded data
					ped_number = len(msg.tracks)
					pedestrian = []

					for ped_idx in range(ped_number):
						ped_data = {
							"position": [msg.tracks[ped_idx].pose.pose.position.x, msg.tracks[ped_idx].pose.pose.position.y],
							"velocity": [msg.tracks[ped_idx].twist.twist.linear.x, msg.tracks[ped_idx].twist.twist.linear.y]
							}
						pedestrian.append(ped_data)
					self.pedestrian_state.append(pedestrian)

				# convert from list to numpy array
				# self.pedestrian_state = np.asanyarray(self.pedestrian_state)

				# obtain the training data with the necessary data assigned for each trajectory
				traj = []
				for j in range(np.min(
						[self.robot_state.shape[0], self.robot_pred_cmd_set.shape[0], len(self.pedestrian_state),
						 self.robot_command_set.shape[0]])):  # iterate over trajectory steps

					data = {"robot_state": self.robot_state[j, :],
					        "pedestrian_state": self.pedestrian_state[j],
					        "control_command": self.robot_command_set[j, :],
					        "pred_command": self.robot_pred_cmd_set[j, :]
					        }

					traj.append(data)

					# break the loop when ever a goal state is encountered, signalling the end of data storage
					# for current trajectory
					if j > 1:
						dist = np.linalg.norm(self.robot_state[j, :2] - self.robot_state[j - 1, :2])
						if dist > self.args.distance_threshold:  # trajectory completion condition
							self.trajectory_set.append(traj)
							traj = []

				bag.close()
				random.shuffle(self.trajectory_set)  # shuffle the final dataset
			else:
				print("Could not find bag file...")

	def RobotStateCB(self, msg):

		pos_x = msg.position.x  # get the current state of the robot
		pos_y = msg.position.y
		orientation = msg.orientation.z
		self.robot_state = np.array([[pos_x, pos_y, orientation]])

	def ControlCommandCB(self, msg):

		linear_vel = msg.linear.x  # get the output commands from LMPCC
		angular_vel = msg.angular.z
		self.control_command = np.array([[linear_vel, angular_vel]])

	def LaserDataCB(self, msg):

		laser_data = []  # store the LiDAR data
		laser_data[:] = msg.ranges

		for i in range(len(laser_data)):
			if np.isinf(laser_data[i]):
				laser_data[i] = 30.0
		self.laser_data = np.array([laser_data])

	def TrackedPersonsCB(self, msg):
		self.pedestrian_state = []
		pos = np.zeros((1, 2))
		vel = np.zeros((1, 2))
		for person in range(len(msg.tracks)):  # for each detected pedestrian, get the required data
			pose_x = msg.tracks[person].pose.pose.position.x
			pose_y = msg.tracks[person].pose.pose.position.y
			vel_x = msg.tracks[person].twist.twist.linear.x
			vel_y = msg.tracks[person].twist.twist.linear.y
			pos = np.vstack((pos, [pose_x, pose_y]))
			vel = np.vstack((vel, [vel_x, vel_y]))
			ped_data = {"position": pos[1:],
			            "velocity": vel[1:]
			            }
		self.pedestrian = ped_data

	def GridProcessingCB(self, msg):
		data = []
		data[:] = msg.data
		for i in range(len(data)):  # binarize the static grid data
			if data[i] != 0:
				data[i] = 1
			else:
				data[i] = 0
		data = np.array([data])
		data = np.reshape(data, (msg.info.width, msg.info.height))
		self.grid = self.grid_rotate(data)

	def PredictedCommandCB(self, msg):

		pred_cmd = np.zeros(self.args.prediction_cmd_horizon * 2)
		horizon = len(msg.poses)

		for n in range(horizon):
			linear_vel = msg.poses[n].pose.position.x
			angular_vel = msg.poses[n].pose.position.y

			if math.isnan(angular_vel):
				angular_vel = 0
			if math.isnan(linear_vel):
				linear_vel = 0
			pred_cmd[2 * n] = linear_vel
			pred_cmd[2 * n + 1] = angular_vel

		self.pred_cmd = np.array([pred_cmd])

	def collect_data(self):

		if self.control_command.size != 0 and self.pred_cmd.size != 0:
			self.data = {"robot_state": self.robot_state,
			             "pedestrian_state": self.pedestrian,
			             "control_command": self.control_command,
			             "predicted_cmd": self.pred_cmd
			             }

			self.trajectory.append(self.data)

			if len(self.trajectory) > 1:
				dist = np.linalg.norm(self.trajectory[-1]["robot_state"][:, :2] - self.trajectory[-2]["robot_state"][:, :2])
				if dist > self.args.distance_threshold:
					self.trajectory_set.append(self.trajectory)
					self.trajectory = []

		print("Dataset length: " + str(len(self.trajectory_set)))
		print("Trajectory length" + str(len(self.trajectory)))

	def getSeqToSeqBatch(self):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are filled and can be used or need to be updated.
		new_epoch = False
		t_idx = self.prev_horizon
		for ii in range(0, self.batch_size):
			traj = self.trajectory_set[self.data_idx]
			if self.sequence_idx[ii] + self.prediction_horizon + 1 >= len(traj):  # change trajectory
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set))
				if self.data_idx == 0:
					new_epoch = True
				self.sequence_idx[ii] = self.prev_horizon
				self.sequence_reset[ii] = 1
				t_idx = self.prev_horizon
			else:
				# start_index = 0 + self.prev_horizon
				# end_index = len(traj) - self.prediction_horizon
				# for t_idx in range(start_index,end_index):
				idx = 0
				for i in range(t_idx - self.prev_horizon, t_idx + 1, 1):
					robot_pos = traj[i]["robot_state"][:2]
					robot_vel = traj[i]["control_command"][:]
					ped_pos = traj[i]["pedestrian_state"]["position"]
					ped_vel = traj[i]["pedestrian_state"]["velocity"]
					self.batch_enc[ii, idx, :] = np.concatenate((robot_pos, robot_vel, ped_pos, ped_vel))
					idx += 1
				idx = 0
				for i in range(t_idx, t_idx + self.prediction_horizon, 1):
					self.batch_dec[ii, idx, :] = traj[i - 1]["pedestrian_state"]["velocity"]
					self.batch_dec_target[ii, idx, :] = traj[i]["pedestrian_state"]["velocity"]
					idx += 1
				t_idx += 1
				self.sequence_idx[ii] += 1
				self.sequence_reset[ii] = 0

		return self.batch_enc, self.batch_dec, self.batch_dec_target, new_epoch

	def getBatch(self):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are filled and can be used or need to be updated.
		"""
				Get the next batch of training data.
		"""
		# If the dataset if smaller that the batch_size, repeat trajectories
		traj_id = 0
		while len(self.trajectory_set) < self.batch_size:
			self.trajectory_set.append(self.trajectory_set[traj_id])
			traj_id += 1
		# Update sequences
		# If batch sequences are empty and need to be filled
		if len(self.batch_sequences) == 0:
			for b in range(self.batch_size):
				trajectory = self.trajectory_set[self.data_idx]
				self.data_idx += 1
				self.batch_sequences.append(trajectory)

		# If batch sequences are filled and can be used or need to be updated.
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			# if self.sequence_idx[ii] + self.prediction_horizon + self.tbpl + 1 >= len(traj):  # change trajectory
			if self.sequence_idx[ii] + self.prediction_horizon + self.tbpl >= len(traj):  # change trajectory
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set))
				if self.data_idx == 0:
					new_epoch = True
				self.sequence_idx[ii] = self.args.prev_horizon
				self.sequence_reset[ii] = 1
			else:
				self.sequence_reset[ii] = 0
		# fill the batch
		other_agents_pos = []

		for ii in range(self.batch_size):
			traj = self.batch_sequences[ii]
			start_idx = int(self.sequence_idx[ii])
			for tbp_step in range(self.tbpl):
				for prev_step in range(self.args.prev_horizon, -1, -1):
					robot_pos = traj[start_idx + tbp_step - prev_step]["robot_state"][:2]
					robot_vel = traj[start_idx + tbp_step - prev_step]["control_command"][:]
					ped_pos = traj[start_idx + tbp_step - prev_step]["pedestrian_state"]["position"]
					ped_vel = traj[start_idx + tbp_step - prev_step]["pedestrian_state"]["velocity"]

					self.batch_vel[ii, tbp_step, prev_step * self.args.input_state_dim:(prev_step + 1) * self.args.input_state_dim] = ped_vel
					self.batch_pos[ii, tbp_step, prev_step * self.args.input_state_dim:(prev_step + 1) * self.args.input_state_dim] = ped_pos

					self.batch_x[ii, tbp_step, prev_step * self.args.input_dim:(prev_step + 1) * self.args.input_dim] = np.array(
						[ped_pos[0], ped_pos[1], ped_vel[0], ped_vel[1]])

				# self.batch_goal[ii, tbp_step, :] = traj[start_idx + tbp_step]["pedestrian_goal_position"]
				self.batch_goal[ii, tbp_step, :] = traj[start_idx + tbp_step]["goal_position"]

				heading = math.atan2(ped_vel[1], ped_vel[0])

				# expands vector because the first axis represents an individual agent
				other_poses = np.expand_dims(robot_pos, axis=0)
				self.other_agents_pos[ii, tbp_step, :] = other_poses
				other_pos_local_frame = sup.positions_in_local_frame(ped_pos, heading, other_poses)
				radial_pedestrian_grid = sup.compute_radial_distance_vector(self.args.pedestrian_vector_dim,
				                                                            other_pos_local_frame,
				                                                            max_range=self.args.max_range_ped_grid,
				                                                            min_angle=0, max_angle=2 * np.pi,
				                                                            normalize=True)
				rel_pos = robot_pos - ped_pos
				rel_vel = robot_vel - ped_vel

				if (self.args.relative_info):
					if self.args.pedestrian_vector_dim == 3:
						self.other_agents_info[ii, tbp_step, :] = np.concatenate((np.array([np.linalg.norm(rel_pos)]), rel_vel))
					else:
						self.other_agents_info[ii, tbp_step, :] = np.concatenate((rel_pos, rel_vel))
				else:
					self.other_agents_info[ii, tbp_step, :] = radial_pedestrian_grid

				self.batch_grid[ii, tbp_step, :, :] = np.zeros((int(self.args.submap_width / self.args.submap_resolution),
				                                                int(self.args.submap_height / self.args.submap_resolution)))

				if self.args.agents_on_grid:
					self.add_other_agents_to_grid(self.batch_grid, ii, tbp_step, ped_pos, self.other_agents_pos)

				if self.args.rotated_grid:
					grid = sup.rotate_grid_around_center(grid, -heading * 180 / math.pi)  # rotation in degrees

				# Output values
				if self.args.predict_positions:
					data = "position"
				else:
					data = "velocity"
				for pred_step in range(self.prediction_horizon):
					if self.args.normalize_data:
						self.batch_target[ii, tbp_step, pred_step * self.args.output_dim] = \
							(traj[tbp_step + 1 + pred_step]["pedestrian_state"][data][0] - self.min_x) / (
									self.max_x - self.min_x)
						self.batch_target[ii, tbp_step, pred_step * self.args.output_dim + 1] = \
							(traj[tbp_step + 1 + pred_step]["pedestrian_state"][data][1] - self.min_y) / (
									self.max_y - self.min_y)
					else:
						self.batch_target[ii, tbp_step,
						pred_step * self.args.output_dim:pred_step * self.args.output_dim + self.args.output_dim] = \
							traj[start_idx + tbp_step + 1 + pred_step]["pedestrian_state"][data]

			self.sequence_idx[ii] += self.tbpl

		by = self.batch_target.copy()
		bx = self.batch_x.copy()

		# Rotate velocities into agent frame at each time step (x-axis is aligned with agents direction of motion)
		if self.args.rotated_grid:
			for batch_idx in range(self.batch_x.shape[0]):
				for tbp_step in range(self.batch_x.shape[1]):
					heading = math.atan2(bx[batch_idx, tbp_step, 3], bx[batch_idx, tbp_step, 2])
					rot_mat = np.array([[np.cos(-heading), -np.sin(-heading)], [np.sin(-heading), np.cos(-heading)]])
					bx[batch_idx, tbp_step, 2:] = np.dot(rot_mat, bx[batch_idx, tbp_step, 2:])
					for pred_step in range(by.shape[2] / 2):
						by[batch_idx, tbp_step, 2 * pred_step:2 * pred_step + 2] = np.dot(rot_mat,
						                                                                  by[batch_idx, tbp_step,
						                                                                  2 * pred_step:2 * pred_step + 2])

		return bx, self.batch_vel, self.batch_pos, self.batch_goal, self.batch_grid, self.other_agents_info, by, self.other_agents_pos, new_epoch


	def add_other_agents_to_grid(self, batch_grid, seq_index, t_idx, ped_pos, other_agents):
		sigma = 3  # assumed position error
		rows = batch_grid.shape[2]
		cols = batch_grid.shape[3]
		# Todo expand this for more agents
		other_agents_pos = other_agents[seq_index][t_idx]
		for i in range(1):
			center_x = int(
				(other_agents_pos[0] - ped_pos[0] + rows / 2 * self.args.submap_resolution) / self.args.submap_resolution)
			center_y = int(
				(other_agents_pos[1] - ped_pos[1] + cols / 2 * self.args.submap_resolution) / self.args.submap_resolution)
			for idx in range(max(center_x - 5, 0), min(center_x + 5, rows)):
				for idy in range(max(center_y - 5, 0), min(center_y + 5, cols)):
					z = 1.0 * np.square(idx - center_x) / np.square(sigma) + 1.0 * np.square(
						idy - center_y) / np.square(
						sigma)
					pdf = max(0.0, 100.0 / 2.0 / np.pi / np.square(sigma) * np.exp(-z / 2.0))
					batch_grid[seq_index, t_idx, idx, idy] += pdf
					if batch_grid[seq_index, t_idx, idx, idy] > 1.0:
						batch_grid[seq_index, t_idx, idx, idy] = 1.0


	def getBatchTBP(self):
		"""
		Get the next batch of training data.
		"""
		self.batch_enc = np.zeros((self.batch_size, self.tbpl, self.args.input_dim * self.args.prev_horizon))
		self.batch_dec = np.zeros((self.batch_size, self.tbpl, self.args.output_dim))
		self.batch_dec_target = np.zeros((self.batch_size, self.tbpl, self.args.output_dim * self.args.prediction_horizon))
		# Update sequences
		# If batch sequences are empty and need to be filled
		trajectory = []
		if len(self.batch_sequences) == 0:
			for b in range(0, min(self.batch_size, len(self.trajectory_set) - len(self.trajectory_set) % self.batch_size)):
				trajectory = self.trajectory_set[self.data_idx]
				self.data_idx += 1
				self.batch_sequences.append(trajectory)
		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			if self.sequence_idx[ii] + self.tbpl + self.prediction_horizon + 1 >= len(traj):
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set))
				trajectory = self.trajectory_set[self.data_idx]
				if self.data_idx == 0:
					new_epoch = True
				self.batch_sequences[ii] = trajectory
				self.sequence_idx[ii] = self.prev_horizon
				self.sequence_reset[ii] = 1
			else:
				self.sequence_reset[ii] = 0

		# Fill the batch
		for ii in range(0, min(self.batch_size, len(self.trajectory_set) - len(
				self.trajectory_set) % self.batch_size)):  # min condition is used for small datasets
			traj = self.batch_sequences[ii]
			self.fillBatch(ii, int(self.sequence_idx[ii]), self.batch_enc, self.batch_dec, self.batch_dec_target, traj)
			self.sequence_idx[ii] += self.tbpl

		return self.batch_enc, self.batch_dec, self.batch_dec_target, new_epoch


	def fillBatch(self, batch_idx, start_idx, batch_enc, batch_dec, batch_dec_target, traj):
		"""
		Fill the data batches of batch_idx with data for all truncated backpropagation steps.
		"""
		other_agents_pos = []

		for tbp_step in range(self.tbpl):

			# Input values
			idx = 0
			for i in range(-self.prev_horizon, 0, 1):
				robot_pos = traj[start_idx + tbp_step + i]["robot_state"][0, :2]
				robot_vel = traj[start_idx + tbp_step + i]["control_command"][0, :]
				ped_pos = traj[start_idx + tbp_step + i]["pedestrian_state"]["position"]
				ped_vel = traj[start_idx + tbp_step + i]["pedestrian_state"]["velocity"]
				batch_enc[batch_idx, tbp_step, idx:idx + self.args.input_dim] = np.concatenate((robot_pos,
				                                                                                robot_vel,
				                                                                                ped_pos,
				                                                                                ped_vel))
				idx += self.args.input_dim

			ped_vel = traj[start_idx + tbp_step]["pedestrian_state"]["velocity"]

			radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim,
			                                                            other_pos_local_frame,
			                                                            max_range=self.max_range_ped_grid, min_angle=0,
			                                                            max_angle=2 * np.pi,
			                                                            normalize=True)
			pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			batch_dec[batch_idx, tbp_step, :] = ped_vel
			idx = 0
			for i in range(0, self.prediction_horizon, 1):
				ped_vel = traj[start_idx + tbp_step + i]["pedestrian_state"]["velocity"]
				batch_dec_target[batch_idx, tbp_step, idx:idx + self.args.output_dim] = ped_vel
				idx += self.args.output_dim


	def getTrajectoryAsBatchTBP(self):
		"""
		Get the next batch of training data.
		"""
		id = np.random.randint(0, len(self.trajectory_set))
		traj = self.trajectory_set[id]
		start_index = 0 + self.prev_horizon
		end_index = len(traj) - self.prediction_horizon
		batch_enc_seq = []
		batch_dec_seq = []
		for t_idx in range(start_index, end_index):
			enc_seq = np.zeros((1, 8 * self.prev_horizon))
			dec_seq = np.zeros((1, 2))
			idx = 0
			for i in range(t_idx - self.prev_horizon, t_idx, 1):
				robot_pos = traj[i]["robot_state"][0, :2]
				robot_vel = traj[i]["control_command"][0, :]
				ped_pos = traj[i]["pedestrian_state"]["position"]
				ped_vel = traj[i]["pedestrian_state"]["velocity"]
				input_data = np.concatenate((robot_pos, robot_vel, ped_pos, ped_vel))
				enc_seq[0, idx:idx + 8] = input_data
				idx += 8
			idx = 0
			batch_enc_seq.append(enc_seq)

			dec_seq[0, :] = traj[t_idx]["pedestrian_state"]["velocity"]

			batch_dec_seq.append(dec_seq)

		return batch_enc_seq, batch_dec_seq, traj


	def getTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		traj = self.trajectory_set[trajectory_idx]

		sequence_length = min(max_sequence_length, len(traj) - self.prediction_horizon)
		batch_x = np.zeros([1, sequence_length, self.args.input_dim*(self.prev_horizon+1)])  # data fed for training
		batch_pos = np.zeros([1, sequence_length, self.args.input_state_dim*(self.prev_horizon+1)])  # data fed for training
		batch_vel = np.zeros([1, sequence_length, self.args.input_state_dim*(self.prev_horizon+1)])  # data fed for training
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_target = np.zeros([1, sequence_length, self.args.output_dim * self.prediction_horizon])

		other_agents_info = np.zeros([1, sequence_length, self.args.pedestrian_vector_dim])
		batch_grid = np.zeros((1, sequence_length, int(self.args.submap_width / self.args.submap_resolution),
		                       int(self.args.submap_height / self.args.submap_resolution)))
		other_agents_pos = np.zeros((1, sequence_length, self.args.pedestrian_vector_state_dim))
		# fill the batch
		start_idx = self.args.prev_horizon
		for tbp_step in range(sequence_length):
			for prev_step in range(self.args.prev_horizon, -1, -1):
				robot_pos = traj[start_idx + tbp_step - prev_step]["robot_state"][:2]
				robot_vel = traj[start_idx + tbp_step - prev_step]["control_command"][:]
				ped_pos = traj[start_idx + tbp_step - prev_step]["pedestrian_state"]["position"]
				ped_vel = traj[start_idx + tbp_step - prev_step]["pedestrian_state"]["velocity"]

				batch_vel[0, tbp_step,
				prev_step * self.args.input_state_dim:(prev_step + 1) * self.args.input_state_dim] = ped_vel
				batch_pos[0, tbp_step,
				prev_step * self.args.input_state_dim:(prev_step + 1) * self.args.input_state_dim] = ped_pos

				batch_x[0, tbp_step, prev_step * self.args.input_dim:(prev_step + 1) * self.args.input_dim] = np.array(
					[ped_pos[0], ped_pos[1], ped_vel[0], ped_vel[1]])

			# batch_goal[0, tbp_step, :] = traj[tbp_step]["pedestrian_goal_position"]
			batch_goal[0, tbp_step, :] = traj[tbp_step]["goal_position"]

			heading = math.atan2(robot_vel[1], robot_vel[0])

			# expands vector because the first axis represents an individual agent
			other_poses = robot_pos
			other_agents_pos[0, tbp_step, :] = other_poses
			other_pos_local_frame = sup.positions_in_local_frame(ped_pos, heading, np.expand_dims(other_poses, axis=0))
			radial_pedestrian_grid = sup.compute_radial_distance_vector(self.args.pedestrian_vector_dim,
			                                                            other_pos_local_frame,
			                                                            max_range=self.args.max_range_ped_grid,
			                                                            min_angle=0, max_angle=2 * np.pi,
			                                                            normalize=True)
			rel_pos = robot_pos - ped_pos
			rel_vel = robot_vel - ped_vel

			if (self.args.relative_info):
				if self.args.pedestrian_vector_dim == 3:
					other_agents_info[0, tbp_step, :] = np.concatenate((np.array([np.linalg.norm(rel_pos)]), rel_vel))
				else:
					other_agents_info[0, tbp_step, :] = np.concatenate((rel_pos, rel_vel))
			else:
				other_agents_info[0, tbp_step, :] = radial_pedestrian_grid

			# self.add_other_agents_to_grid(batch_grid,0, tbp_step, ped_pos, other_agents_pos)

			# Output values
			for pred_step in range(self.prediction_horizon):
				if self.args.normalize_data:
					batch_target[0, tbp_step, pred_step * self.args.output_dim] = \
						(traj[tbp_step + 1 + pred_step]["pedestrian_state"]["velocity"][0] - self.min_x) / (self.max_x - self.min_x)
					batch_target[0, tbp_step, pred_step * self.args.output_dim + 1] = \
						(traj[tbp_step + 1 + pred_step]["pedestrian_state"]["velocity"][1] - self.min_y) / (self.max_y - self.min_y)
				else:
					batch_target[0, tbp_step,
					pred_step * self.args.output_dim:pred_step * self.args.output_dim + self.args.output_dim] = \
						traj[tbp_step + 1 + pred_step]["pedestrian_state"]["velocity"]

		return batch_x, batch_vel, batch_pos, batch_goal, batch_grid, other_agents_info, batch_target, other_agents_pos, traj


	def getTrajectoryAsBatchSeqToSeq(self):
		"""
		Get the next batch of training data.
		"""
		id = np.random.randint(0, len(self.trajectory_set))
		traj = self.trajectory_set[id]
		start_index = 0 + self.prev_horizon
		end_index = len(traj) - self.prediction_horizon
		batch_enc_seq = []
		batch_dec_seq = []
		for t_idx in range(start_index, end_index):
			enc_seq = np.zeros((self.prev_horizon, 8))
			dec_seq = np.zeros((self.prediction_horizon, 2))
			idx = 0
			for i in range(t_idx - self.prev_horizon, t_idx, 1):
				robot_pos = traj[i]["robot_state"][:2]
				robot_vel = traj[i]["control_command"][:]
				ped_pos = traj[i]["pedestrian_state"]["position"]
				ped_vel = traj[i]["pedestrian_state"]["velocity"]
				input_data = np.concatenate((robot_pos, robot_vel, ped_pos, ped_vel))
				enc_seq[idx, :] = input_data
				idx += 1
			idx = 0
			batch_enc_seq.append(enc_seq)
			for i in range(t_idx, t_idx + self.prediction_horizon, 1):
				dec_seq[idx, :] = traj[i]["pedestrian_state"]["velocity"]
				idx += 1
			batch_dec_seq.append(dec_seq)

		return batch_enc_seq, batch_dec_seq, traj


	def calc_scale(self, keep_ratio=True):
		self.sx = 1 / (self.max_x - self.min_x)
		self.sy = 1 / (self.max_y - self.min_y)
		if keep_ratio:
			if self.sx > self.sy:
				self.sx = self.sy
			else:
				self.sy = self.sx


	def normalize(self, data, shift=True, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_x * shift) * self.sx
			data_copy[1] = (data[1] - self.min_y * shift) * self.sy
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_x * shift) * self.sx
			data_copy[:, 1] = (data[:, 1] - self.min_y * shift) * self.sy
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_x * shift) * self.sx
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_y * shift) * self.sy
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_x * shift) * self.sx
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_y * shift) * self.sy
		else:
			return False
		return data_copy


	def denormalize(self, data, shift=True, inPlace=False):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx + self.min_x * shift
			data_copy[1] = data[1] / self.sy + self.min_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx + self.min_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy + self.min_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx + self.min_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy + self.min_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx + self.min_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy + self.min_y * shift
		else:
			return False

		return data_copy


	def save_data(self):
		file = open(self.args.data_path, 'wb')
		print("Saving data to: '{}'".format(self.args.data_path))
		# self.dataset.pop(0)
		pkl.dump(self.trajectory_set, file)

		file.close()


	def load_data(self):
		print("Loading data from: '{}'".format(self.args.data_path + self.args.dataset))
		self.file = open(self.args.data_path + self.args.dataset, 'rb')
		if sys.version_info[0] < 3:
			tmp_self = pkl.load(self.file)
		else:
			tmp_self = pkl.load(self.file, encoding='latin1')
		self.trajectory_set = tmp_self
		trajectory_set = []
		# Get normalization constants
		self.pedestrian_dataset = []
		self.robot_dataset = []
		self.max_x = -100
		self.min_x = 100
		self.max_y = -100
		self.min_y = 100
		for traj_id in range(len(self.trajectory_set)):
			traj = self.trajectory_set[traj_id]
			self.robot_vel = np.zeros((len(traj), 2))
			self.pedestrian_vel = np.zeros((len(traj), 2))
			sub_traj = []
			# Sub-sample trajectories to make dt = 0.4
			for t_id in range(0, len(traj), 4):
				self.robot_vel[t_id, :] = traj[t_id]["control_command"]
				self.pedestrian_vel[t_id, :] = traj[t_id]["pedestrian_state"][""]
				self.max_x = max(self.max_x, traj[t_id]["pedestrian_state"]["position"][0])
				self.min_x = min(self.min_x, traj[t_id]["pedestrian_state"]["position"][0])
				self.max_y = max(self.max_y, traj[t_id]["pedestrian_state"]["position"][1])
				self.min_y = min(self.min_y, traj[t_id]["pedestrian_state"]["position"][1])
				sub_traj.append(traj[t_id])
			# Append zero velocities to goal position
			if traj:
				goal_data = traj[-1].copy()
				goal_data["pedestrian_state"]["velocity"] = np.array([0, 0])
				for t_id in range(10):
					sub_traj.append(goal_data)
				if len(sub_traj) > self.min_length_trajectory:
					trajectory_set.append(sub_traj)
				self.robot_dataset.append(self.robot_vel)
				self.pedestrian_dataset.append(self.pedestrian_vel)
		self.robot_trajectory_set = np.asarray(self.robot_dataset)
		self.pedestrian_dataset = np.asarray(self.pedestrian_dataset)
		self.trajectory_set = trajectory_set
		self.calc_scale()
		# random.shuffle(self.trajectory_set)


	def load_map(self):
		try:
			file_path = self.args.data_path + '/map.npy'
			grid_map = np.load(file_path)[()]
			print("Map loaded")
			return grid_map
		except IOError:
			print("ERROR: Cannot load map.")
			return False


	def start_node(self):
		rospy.init_node('OnlineDataHandler', anonymous=True)
		# start node to collect data online
		while not rospy.is_shutdown():  # while the node is running, get the complete data, input and output batches
			self.collect_data()
			if len(self.trajectory_set) > self.args.max_trajectories:
				break
			sleep(1.0 / self.rate)
		self.save_data()


if __name__ == '__main__':
	data_prep = DataHandler()
	data_prep.start_node()

build/lib/data_utils/DataHandler.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import os
import random
import sys
import math
import cv2
import pickle as pkl
import json
from pykalman import KalmanFilter
if sys.version_info[0] < 3:
	import Support as sup
	from Trajectory import *
	import AgentContainer as ped_cont
else:
	import src.data_utils.Support as sup
	from src.data_utils.Trajectory import *
	import src.data_utils.AgentContainer as ped_cont
from copy import deepcopy
import matplotlib.pyplot as pl
import matplotlib.animation as animation
from time import sleep
import random
from scipy.stats import multivariate_normal
#import matplotlib.pyplot as plt

from matplotlib.patches import Ellipse

class DataHandlerLSTM():
	"""
	Data handler for training an LSTM pedestrian prediction model
	"""
	def __init__(self, args):

		self.data_path = args.data_path
		self.batch_size = args.batch_size
		self.tbpl = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.output_sequence_length = args.prediction_horizon
		self.prediction_horizon = args.prediction_horizon
		self.input_dim = args.input_dim
		self.input_state_dim = args.input_state_dim
		self.output_state_dim = args.output_dim
		#self.submap_width = args.submap_width
		#self.submap_height = args.submap_height
		self.submap_size = args.submap_size
		self.submap_width_real = args.submap_span_real
		self.submap_height_real = args.submap_span_real
		#self.submap_resolution = args.submap_resolution
		self.centered_grid = args.centered_grid
		self.rotated_grid = args.rotated_grid
		self.multi_pedestrian = True
		self.pedestrian_radius = args.pedestrian_radius
		self.min_length_trajectory = self.tbpl + 1 + self.output_sequence_length + args.prev_horizon
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.max_range_ped_grid = args.max_range_ped_grid
		self.dt = args.dt
		self.n_mixture = args.n_mixtures
		self.output_pred_state_dim = args.output_pred_state_dim
		self.args = args
		self.real_world_data = args.real_world_data
		# Normalization constants
		self.norm_const_x = 1.0
		self.norm_const_y = 1.0
		self.norm_const_heading = 1.0
		self.norm_const_vx = 1.0
		self.norm_const_vy = 1.0
		self.norm_const_omega = 1.0
		self.min_pos_x = 1000
		self.min_pos_y = 1000
		self.max_pos_x = -1000
		self.max_pos_y = -1000
		self.min_vel_x = 1000
		self.min_vel_y = 1000
		self.max_vel_x = -1000
		self.max_vel_y = -1000
		self.train_set = args.train_set
		self.avg_traj_length = 0

		# Default Distribution Parameters
		self.amplitude_ = 1.0
		self.factor_ = 5.0
		self.cutoff_ = 1.0
		self.covar_ = 0.25
		self.var = multivariate_normal(mean=0, cov=1)

		# Data structure containing all the information about agents
		self.agent_container = ped_cont.AgentContainer()
		self.test_container = ped_cont.AgentContainer()

		self.agent_traj_idx = None  # indicates which trajectory index will be the next one per agent
		self.trajectory_set = []
		self.test_trajectory_set = []
		self.test_ids = []
		self.data_idx = 0
		self.val_data_idx = 0

		# Patch for updating the grid around a pedestrian position
		self.pedestrian_patch = None

		# Training variables
		self.sequence_reset = np.ones([self.batch_size])  # indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.sequence_idx = np.zeros([self.batch_size])+self.args.prev_horizon

		# Batch with info about other agents
		if "future" in self.args.others_info:
			self.batch_vel = np.zeros(
				[self.batch_size, self.tbpl, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			self.batch_vel = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		self.batch_x = np.zeros([self.batch_size, self.tbpl, self.input_dim*(self.prev_horizon+1)])  # data fed for training
		self.batch_grid = np.zeros([self.batch_size, self.tbpl,
																int(self.submap_size),
																int(self.submap_size)])
		if self.args.others_info == "relative":
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim*self.args.n_other_agents])
			self.val_pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.pedestrian_vector_dim*self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		elif self.args.others_info == "ped_grid":
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl,
										int(self.submap_size),
										int(self.submap_size)])
		else:
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])
			self.val_pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])
		self.batch_goal = np.zeros([self.batch_size, self.tbpl, 2])
		self.batch_y = np.zeros([self.batch_size, self.tbpl, self.output_state_dim * self.output_sequence_length])
		self.batch_sequences = []  # Sequence stack (tuples of trajectory + grid) where parts of length self.tbpl are taken from until a sequence is finished. As soon as one is finished, a new sequence is put in the stack.
		self.batch_ids = []
		self.batch_pos = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		self.batch_pos_target = np.zeros([self.batch_size, self.tbpl, 2 * self.output_sequence_length])

		# Validation variables
		self.val_sequence_reset = np.ones([self.batch_size])  # indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.val_sequence_idx = np.zeros([self.batch_size]) + self.args.prev_horizon
		# Batch with current query-agent velocities
		self.val_batch_vel = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		# Batch with query-agent current velocities and positions
		self.val_batch_x = np.zeros([self.batch_size, self.tbpl, self.input_dim * (self.prev_horizon + 1)])  # data fed for training
		# Query agent static environment information
		self.val_batch_grid = np.zeros([self.batch_size, self.tbpl,
									int(self.submap_size),
									int(self.submap_size)])

		# Batch of query-agent goal positions
		self.val_batch_goal = np.zeros([self.batch_size, self.tbpl, 2])
		# Batch with quer-agent next velocities
		self.val_batch_y = np.zeros([self.batch_size, self.tbpl, self.output_state_dim * self.output_sequence_length])
		# Sequence stack (tuples of trajectory + grid) where parts of length self.tbpl are taken from until a sequence is finished. As soon as one is finished, a new sequence is put in the stack.
		self.val_batch_sequences = []
		self.val_batch_ids = []
		# Batch of query-agent current position
		self.val_batch_pos = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		# Batch with quer-agent next positions
		self.val_batch_pos_target = np.zeros([self.batch_size, self.tbpl, 2 * self.output_sequence_length])

	def addAgentTrajectoriesToSet(self,agent_container,trajectory_set, id):
		"""
		Goes through all trajectories of agent and adds them to the member set if they fulfill the criteria.
		For all the time steps within the trajectory it also computes the positions of the other agents at that
		timestep in order to make training more efficient.
		"""
		for traj_idx, traj in enumerate(agent_container.getAgentTrajectories(id)):
			traj_with_collision = False
			if len(traj) > self.min_length_trajectory:
				#if traj.getMinTime() < 100:
				traj.updateInterpolators()
				# Find other agent's trajectories which overlap with each time step
				for time_idx in range(traj.time_vec.shape[0]):
					if self.multi_pedestrian:
						query_time = traj.time_vec[time_idx]
						other_agents_positions = agent_container.getAgentPositionsForTimeExclude(query_time, id)
						other_agents_velocities = agent_container.getAgentVelocitiesForTimeExclude(query_time, id)
			
						# Remove ego agent
						traj.other_agents_positions.append(other_agents_positions)
						traj.other_agents_velocities.append(other_agents_velocities)
				# TODO: MAYBE REMOVE IN THE FUTURE
				#if id != -1:
					#if not traj_with_collision:
				trajectory_set.append((id, traj))
	
	def setOtherAgentData(self, agent_container, trajectory_set, id):
		for traj_idx, traj in enumerate(agent_container.getAgentTrajectories(id)):
			traj_with_collision = False
			if len(traj) > self.min_length_trajectory:
				#traj.updateInterpolators()
				for time_idx in range(traj.time_vec.shape[0]):
					if self.multi_pedestrian:
						query_time = traj.time_vec[time_idx]
						other_agents_positions = agent_container.getAgentPositionsForTimeExclude(query_time, id)
						other_agents_velocities = agent_container.getAgentVelocitiesForTimeExclude(query_time, id)
						# Remove ego agent
						traj.other_agents_positions.append(other_agents_positions)
						traj.other_agents_velocities.append(other_agents_velocities)
				trajectory_set[traj_idx][1].other_agents_positions = traj.other_agents_positions
				trajectory_set[traj_idx][1].other_agents_velocities = traj.other_agents_velocities
				

	def processData(self, **kwargs):
		"""
		Processes the simulation or real-world data, depending on the usage.
		"""
		data_pickle = self.args.data_path + self.args.scenario + "/data" + str(self.args.prediction_horizon) + "_" + str(
			self.args.truncated_backprop_length)+ "_" + str(
			self.args.prev_horizon) + "_roboat.pickle"
		#data_pickle = self.args.data_path + self.args.scenario + "/data" + str(self.args.prediction_horizon) + "_" + str(
		#	self.args.truncated_backprop_length)+ "_" + str(
		#	self.args.prev_horizon) + "experiment3agents-norm.pickle"
		if os.path.isfile(data_pickle):
			self.loadTrajectoryData(data_pickle)
		else:
			if "roboat" in data_pickle:
				print("Processing roboat data.")
				self._process_roboat_data_(**kwargs)
			else:
				print("No roboat data present")

			self.saveTrajectoryData(data_pickle)

	def load_map(self,**kwargs):
		print("Extracting the occupancy grid ...")
		# Occupancy grid data

		with open(os.path.join(self.data_path+self.args.scenario, 'map.json')) as json_file:
			map_data = json.load(json_file)

		grid_map = pl.imread(self.data_path+self.args.scenario + '/map.png')
		
		if self.args.scenario == 'open_crossing' or self.args.scenario == "test_open_crossing":
			grid_map = np.round(grid_map)
		else:
			grid_map = np.round(grid_map[:,:,0])

		#map_data = np.load(os.path.join(self.data_path+self.args.scenario, 'map.npy'), encoding='latin1', allow_pickle=True)
		self.agent_container.occupancy_grid.gridmap = grid_map #map_data.item(0)['Map']  # occupancy values of cells
		self.agent_container.occupancy_grid.resolution = map_data['resolution'] #map_data.item(0)['Resolution']  # map resolution in [m / cell]
		self.agent_container.occupancy_grid.map_size = map_data['map_size'] #map_data.item(0)['Size']  # map size in [m]
		self.agent_container.occupancy_grid.center =  map_data['map_center'] #self.agent_container.occupancy_grid.map_size / 2.0

		#FIlter all -1 values present on MAP to zero.
		for ii in range(self.agent_container.occupancy_grid.gridmap.shape[0]):
			for jj in range(self.agent_container.occupancy_grid.gridmap.shape[1]):
				if self.agent_container.occupancy_grid.gridmap[ii, jj] == -1:
					self.agent_container.occupancy_grid.gridmap[ii, jj] = 0.0

	
	def _process_roboat_data_(self, **kwargs):
		"""
		Import the data from the log file stored in the directory of data_path.
		This method brings all the data into a suitable format for training.
		"""
		self.load_map(**kwargs)
		# Pedestrian data
		# [id, timestep (s), timestep (ns), pos x, pos y, yaw, vel x, vel y, omega, goal x, goal y]
		print("Extracting data from csv file")
		pedestrian_data = np.genfromtxt(os.path.join(self.data_path+self.args.scenario, 'total_log.csv'), delimiter=",")[1:, :]
		print("Data extracted")

		# Iterate through the data and fill the register
		for sample_idx in range(pedestrian_data.shape[0]):
			if pedestrian_data[sample_idx, 0] != -1:
				id = pedestrian_data[sample_idx, 0]
				timestamp = np.round(pedestrian_data[sample_idx, 1],1)# + pedestrian_data[sample_idx, 2] * 1e-9  # time in seconds28.35
				pose = np.zeros([1,3])
				vel = np.zeros([1,3])
				pose[:,0:2] = np.true_divide(pedestrian_data[sample_idx, 3:5], np.array([self.norm_const_x, self.norm_const_y]))
				vel[:,0:2] = np.true_divide(pedestrian_data[sample_idx, 5:7], np.array([self.norm_const_vx, self.norm_const_vy]))
				goal = np.true_divide(pedestrian_data[sample_idx, 7:], np.array([self.norm_const_x, self.norm_const_y]))

				self.agent_container.addDataSample(id, timestamp, pose, vel, goal)

		# Set the initial indices for agent trajectories (which trajectory will be returned when queried)
		self.agent_traj_idx = [0] * self.agent_container.getNumberOfAgents()

		# Subsample trajectories (longer discretization time) from dt=0.1 to dt=0.4
		print("Subsampling data..")
		for id in self.agent_container.getAgentIDs():
			for traj in self.agent_container.getAgentTrajectories(id):
				traj.subsample(int(self.args.dt*10))
		print("Sampling done.")

		for id in self.agent_container.getAgentIDs():
			for traj in self.agent_container.getAgentTrajectories(id):
				for t_id in range(1, traj.pose_vec.shape[0]):
					traj.vel_vec[t_id, :2] = (traj.pose_vec[t_id, :2] - traj.pose_vec[t_id-1, :2]) / self.dt

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id in self.agent_container.getAgentIDs():
			for traj_idx, traj in enumerate(self.agent_container.getAgentTrajectories(id)):
				if len(traj) > self.min_length_trajectory:
					traj.updateInterpolators()

		# Put all the trajectories in the trajectory set and randomize
		for id in self.agent_container.getAgentIDs():
			print("Processing agent {} / {}".format(id, self.agent_container.getNumberOfAgents()))
			# Adds trajectory if bigger than a minimum length and maximum size
			self.addAgentTrajectoriesToSet(self.agent_container,self.trajectory_set,id)
		
		self.compute_min_max_values()

		

	def compute_min_max_values(self):
		for traj_id in range(len(self.trajectory_set)):
			for t_id in range(1, self.trajectory_set[traj_id][1].pose_vec.shape[0]):
				self.trajectory_set[traj_id][1].vel_vec[t_id, :2] = (self.trajectory_set[traj_id][1].pose_vec[t_id, :2] -
																	 self.trajectory_set[traj_id][1].pose_vec[t_id - 1,
																	 :2]) / self.dt
				self.min_pos_x = min(self.min_pos_x,self.trajectory_set[traj_id][1].pose_vec[t_id,0])
				self.min_pos_y = min(self.min_pos_y, self.trajectory_set[traj_id][1].pose_vec[t_id, 1])
				self.max_pos_x = max(self.max_pos_x, self.trajectory_set[traj_id][1].pose_vec[t_id, 0])
				self.max_pos_y = max(self.max_pos_y, self.trajectory_set[traj_id][1].pose_vec[t_id, 1])
				self.min_vel_x = min(self.min_vel_x,self.trajectory_set[traj_id][1].vel_vec[t_id,0])
				self.min_vel_y = min(self.min_vel_y, self.trajectory_set[traj_id][1].vel_vec[t_id, 1])
				self.max_vel_x = max(self.max_vel_x, self.trajectory_set[traj_id][1].vel_vec[t_id, 0])
				self.max_vel_y = max(self.max_vel_y, self.trajectory_set[traj_id][1].vel_vec[t_id, 1])

		self.calc_scale()

		self.args.min_pos_x = self.min_pos_x
		self.args.min_pos_y = self.min_pos_y
		self.args.max_pos_x = self.max_pos_x
		self.args.max_pos_y = self.max_pos_y
		self.args.min_vel_x = self.min_vel_x
		self.args.min_vel_y = self.min_vel_y
		self.args.max_vel_x = self.max_vel_x
		self.args.max_vel_y = self.max_vel_y
		self.args.sx_vel = self.sx_vel
		self.args.sy_vel = self.sy_vel
		self.args.sx_pos = self.sx_pos
		self.args.sy_pos = self.sy_pos

	def unit_test_data_(self,map_args):
			"""
			Generates predefined trajectories to analyse the interaction
			"""
			if not os.path.isfile(self.data_path + self.args.scenario + '/unit_map.npy'):
				print("Creating unit map from png ...")
				sup.create_map_from_png(data_path=self.args.data_path + self.args.scenario, **map_args)

			map_data = np.load(os.path.join(self.data_path + self.args.scenario, 'map.npy'), encoding='latin1',
							   allow_pickle=True)
			self.agent_container.occupancy_grid.gridmap = map_data.item(0)['Map']  # occupancy values of cells
			self.agent_container.occupancy_grid.resolution = map_data.item(0)['Resolution']  # map resolution in [m / cell]
			self.agent_container.occupancy_grid.map_size = map_data.item(0)['Size']  # map size in [m]
			self.agent_container.occupancy_grid.center = self.agent_container.occupancy_grid.map_size / 2.0

			# FIlter all -1 values present on MAP to zero.
			for ii in range(self.agent_container.occupancy_grid.gridmap.shape[0]):
				for jj in range(self.agent_container.occupancy_grid.gridmap.shape[1]):
					if self.agent_container.occupancy_grid.gridmap[ii, jj] == -1:
						self.agent_container.occupancy_grid.gridmap[ii, jj] = 0.0

			# Iterate through the data and fill the register
			pose_agent_1 = np.array([[-15.0,0.0,0.0]])
			pose_agent_2 = np.array([[-13.0, 0.0,0.0]])
			vel = np.zeros([1, 3])
			v0 = 1.0
			goal = np.array([15, 0])
			timestamp = 0
			dt =0.4
			for sample_idx in range(int(320/dt/10)):
				id = 0
				timestamp += dt #self.args.dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 1
				pose_agent_2[:, 0] += v0 * dt

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)

			# Second test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 3
				timestamp += dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

			# Third test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 4
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0
				goal = np.array([15,0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 5
				pose_agent_2[:, 0] -= v0 * dt
				vel[:, 0] = -v0

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, -goal)

			# Forth test
			pose_agent_1 = np.array([[15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 6
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] -= v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = -v0
				goal = np.array([-15,0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

			# Fith test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, -1.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 7
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 8
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, -1])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)
			# Forth test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, 1.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 9
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 10
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, -1])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)
			# Forth test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 11
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 12
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)

			# Put all the trajectories in the trajectory set and randomize
			pool_list = []
			for id in self.test_container.getAgentIDs():
				print("Processing agent {} / {}".format(id, self.test_container.getNumberOfAgents()))
				# Adds trajectory if bigger than a minimum length and maximum size
				self.addAgentTrajectoriesToSet(self.test_container,self.test_trajectory_set,id)

			self.min_pos_x = -1.0
			self.min_pos_y = -1.0
			self.max_pos_x = 15.0
			self.max_pos_y = 15.0
			self.min_vel_x = -1.0
			self.min_vel_y = -1.0
			self.max_vel_x = 1.0
			self.max_vel_y = 1.0

			self.calc_scale()


	def calc_scale(self, keep_ratio=True):
		self.sx_vel = 1 / (self.max_vel_x - self.min_vel_x)
		self.sy_vel = 1 / (self.max_vel_y - self.min_vel_y)
		if keep_ratio:
			if self.sx_vel > self.sy_vel:
				self.sx_vel = self.sy_vel
			else:
				self.sy_vel = self.sx_vel

		self.sx_pos = 1 / (self.max_pos_x - self.min_pos_x)
		self.sy_pos = 1 / (self.max_pos_y - self.min_pos_y)
		if keep_ratio:
			if self.sx_pos > self.sy_pos:
				self.sx_pos = self.sy_pos
			else:
				self.sy_pos = self.sx_pos

	def normalize_vel(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[1] = (data[1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, 1] = (data[:, 1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_vel_y * shift) * self.sy_vel
		else:
			return False
		return data_copy

	def normalize_pos(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[1] = (data[1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, 1] = (data[:, 1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_pos_y * shift) * self.sy_pos
		else:
			return False
		return data_copy

	def denormalize_pos(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx_pos + self.min_pos_x * shift
			data_copy[1] = data[1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy_pos + self.min_pos_y * shift
		else:
			return False

		return data_copy

	def denormalize_vel(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx_vel + self.min_vel_x * shift
			data_copy[1] = data[1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy_vel + self.min_vel_y * shift
		else:
			return False

		return data_copy

	def filter_data(self, data):
		# Kalman filter considering a constant velocity model
		Transition_Matrix = [[1, 0, self.dt, 0], [0, 1, 0, self.dt], [0, 0, 1, 0], [0, 0, 0, 1]]
		Observation_Matrix = [[1, 0, 0, 0], [0, 1, 0, 0]]

		xinit = data[0, 0]
		yinit = data[0, 1]
		vxinit = (data[1, 0] - data[0, 0]) / self.dt
		vyinit = (data[1, 1] - data[0, 1]) / self.dt
		initstate = [xinit, yinit, vxinit, vyinit]
		initcovariance = 1.0e-3 * np.eye(4) #  1.0e-3 * np.eye(4)
		transistionCov = 1.0e-4 * np.eye(4) # 1.0e-4 * np.eye(4)
		observationCov = 2.0e-3 * np.eye(2) # 1.0e-2 * np.eye(2)   -6
		kf = KalmanFilter(transition_matrices=Transition_Matrix,
						  observation_matrices=Observation_Matrix,
						  initial_state_mean=initstate,
						  initial_state_covariance=initcovariance,
						  transition_covariance=transistionCov,
						  observation_covariance=observationCov)

		(filtered_state_means, filtered_state_covariances) = kf.filter(data)

		return filtered_state_means

	def saveTrajectoryData(self, save_path):
		print("Saving data to: '{}'".format(save_path))
		if not os.path.isdir(self.args.data_path + self.args.scenario):
			os.makedirs(self.args.data_path + self.args.scenario)

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id, traj in self.trajectory_set:
			traj.updateInterpolators()

		#if "test" not in self.args.scenario:
		random.shuffle(self.trajectory_set)

		data = {
			"trajectories" : self.trajectory_set,
			"agent_container" : self.agent_container,
			"min_pos_x" : self.min_pos_x,
			"min_pos_y" : self.min_pos_y,
			"max_pos_x" : self.max_pos_x,
			"max_pos_y" : self.max_pos_y,
			"min_vel_x" : self.min_vel_x,
			"min_vel_y" : self.min_vel_y,
			"max_vel_x" : self.max_vel_x,
			"max_vel_y" : self.max_vel_y
		}
		pkl.dump(data, open(save_path, 'wb'),protocol=2)

	def loadTrajectoryData(self, load_path):
		print("Loading data from: '{}'".format(load_path))
		self.file = open(load_path, 'rb')
		if sys.version_info[0] < 3:
			tmp_self = pkl.loads(self.file)#,encoding='latin1')
		else:
			tmp_self = pkl.load(self.file , encoding='latin1')
		self.trajectory_set = tmp_self["trajectories"]
		self.agent_container = tmp_self["agent_container"]

		try:
			self.min_pos_x = tmp_self["min_pos_x"]
			self.min_pos_y = tmp_self["min_pos_y"]
			self.max_pos_x = tmp_self["max_pos_x"]
			self.max_pos_y = tmp_self["max_pos_y"]
			self.min_vel_x = tmp_self["min_vel_x"]
			self.min_vel_y = tmp_self["min_vel_y"]
			self.max_vel_x = tmp_self["max_vel_x"]
			self.max_vel_y = tmp_self["max_vel_y"]
		except:
			print("Delete dataset to compute limits")
			self.max_x = 1
			self.max_y = 1
			self.min_x = -1
			self.min_y = -1

		self.calc_scale()

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id, traj in self.trajectory_set:
			traj.updateInterpolators()

	def getAgentTrajectory(self, agent_id):
		"""
		Return the next agent trajectory in the queue for the agent with id agent_id.
		"""
		trajectory = self.agent_container.agent_data[agent_id].trajectories[self.agent_traj_idx[agent_id]]
		self.agent_traj_idx[agent_id] = (self.agent_traj_idx[agent_id] + 1)  % self.agent_container.getNumberOfTrajectoriesForAgent(agent_id)
		return trajectory

	def getRandomAgentTrajectory(self, agent_id):
		"""
		Return a totally random trajectory for the agent with id agent_id.
		"""
		random_traj_idx = np.random.randint(0, len(self.agent_container.agent_data[agent_id].trajectories))
		return self.agent_container.agent_data[agent_id].trajectories[random_traj_idx]

	def fillBatch(self, agent_id, batch_idx, start_idx, truncated_backprop_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, trajectory,batch_pos_target, centered_grid=False, testing = False):
		"""
		Fill the data batches of batch_idx with data for all truncated backpropagation steps.
		"""
		other_agents_pos = []

		for tbp_step in range(truncated_backprop_length):

			# Input values
			query_time = trajectory.time_vec[start_idx + tbp_step]
			for prev_step in range(self.prev_horizon,-1,-1):
				current_pos = np.array([trajectory.pose_vec[start_idx + tbp_step - prev_step, 0], trajectory.pose_vec[
											start_idx + tbp_step - prev_step, 1]])
				#current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step, 0], trajectory.vel_vec[
				#							start_idx + tbp_step - prev_step - prev_step, 1]])
				current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step - prev_step, 0], trajectory.vel_vec[
											start_idx + tbp_step - prev_step, 1]])

				if self.args.normalize_data:
					self.normalize_pos(current_pos)
					self.normalize_vel(current_vel)

				batch_x[batch_idx, tbp_step, prev_step*self.input_dim:(prev_step+1)*self.input_dim] = np.array([current_pos[0],
																										current_pos[1],
																										current_vel[0],
																										current_vel[1]])
				#batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([np.linalg.norm(current_vel),
				#                                                                                                              np.arctan2(current_vel[1],current_vel[0])])
				batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_vel[0],
																											current_vel[1]])
				batch_pos[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_pos[0],
																											current_pos[1]])
				
			heading = math.atan2(current_vel[1], current_vel[0])
			if centered_grid:
				grid_center = current_pos
				grid = self.agent_container.occupancy_grid.getSubmapByCoords(grid_center[0], grid_center[1], self.submap_width_real, self.submap_height_real)
				grid = cv2.resize(grid, (60,60)) # Resize grid to (60,60) to match autoencoder input size
			if self.rotated_grid:
				grid = sup.rotate_grid_around_center(grid, heading * 180 / math.pi)  # rotation in degrees

			batch_grid[batch_idx, tbp_step, :, :] = grid	

			# Find positions of other pedestrians at the current timestep and order them by dstance to query agent
			other_positions = trajectory.other_agents_positions[start_idx + tbp_step]
			n_other_agents = other_positions.shape[0]
			other_velocities = trajectory.other_agents_velocities[start_idx + tbp_step]

			# Compute distance to other agents and order them
			other_poses_ordered = np.zeros((other_positions.shape[0], 6))
			other_poses_ordered[:, :2] = other_positions
			other_poses_ordered[:, 2:4] = other_velocities
			other_agents_pos.append(other_poses_ordered)

			current_pos = np.array(
				[trajectory.pose_vec[start_idx + tbp_step, 0], trajectory.pose_vec[start_idx + tbp_step, 1]])
			current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step, 0], trajectory.vel_vec[start_idx + tbp_step, 1]])
			heading = math.atan2(current_vel[1], current_vel[0])

			for ag_id in range(n_other_agents):
				other_poses_ordered[ag_id,4] = np.linalg.norm(other_poses_ordered[ag_id,:2] - current_pos)
				# ordered ids
				other_poses_ordered[ag_id,5] = ag_id
			other_poses_ordered= other_poses_ordered[other_poses_ordered[:, 4].argsort()]

			# other_poses_ordered = [ x, y, vx, vy, norm distance, id]

			if self.args.others_info == "relative":
				for ag_id in range(min(n_other_agents,self.args.n_other_agents)):
					rel_pos = np.array([other_poses_ordered[ag_id,0] - current_pos[0],other_poses_ordered[ag_id, 1] - current_pos[1]])*\
								  multivariate_normal.pdf(np.linalg.norm(np.array([other_poses_ordered[ag_id,:2] - current_pos])),
																			  mean=0.0,
																			  cov=self.args.relative_covariance)

					rel_vel = np.array([other_poses_ordered[ag_id,2] - current_vel[0],other_poses_ordered[ag_id, 3] - current_vel[1]])
					
					pedestrian_grid[batch_idx, tbp_step, ag_id*4:ag_id*4+4] = np.concatenate([rel_pos, rel_vel])
					#pedestrian_grid[batch_idx, tbp_step, ag_id, 4] = np.linalg.norm(rel_pos)
					#pedestrian_grid[batch_idx, tbp_step, ag_id, 5] = np.arctan2(rel_pos[1], rel_pos[0])
			

			elif self.args.others_info == "angular_grid":
				other_pos_local_frame = sup.positions_in_local_frame(current_pos, heading, other_positions)
				radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																															max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																															normalize=True)
				pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			# Output values
			for pred_step in range(self.output_sequence_length):
				vx = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 0]
				vy = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 1]
				px = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 0]
				py = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 1]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step] = vx #- current_vel[0]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = vy #- current_vel[1]
				if self.args.normalize_data:
					self.normalize_vel(batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step:self.output_state_dim*pred_step+2])
				batch_pos_target[batch_idx, tbp_step, self.output_state_dim*pred_step] = px - trajectory.pose_vec[start_idx + tbp_step, 0]
				batch_pos_target[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = py - trajectory.pose_vec[start_idx + tbp_step, 1]

			current_pos = trajectory.pose_vec[start_idx + tbp_step, :2]
			#delta_pos = trajectory.pose_vec[start_idx + tbp_step + self.prediction_horizon,:2] - current_pos
			delta_pos = trajectory.goal[0] - current_pos
			#batch_goal[batch_idx, tbp_step, 0] = np.linalg.norm(delta_pos)
			#batch_goal[batch_idx, tbp_step, 1] = np.arctan2(delta_pos[1],
			#                                                delta_pos[0])
			# Used for planning
			batch_goal[batch_idx, tbp_step, 0] = delta_pos[0]/np.linalg.norm(delta_pos) #trajectory.goal[0,0]
			batch_goal[batch_idx, tbp_step, 1] = delta_pos[1]/np.linalg.norm(delta_pos) #trajectory.goal[0,1]

		return other_agents_pos

	def getBatch(self):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled
		trajectory=[]
		if len(self.batch_sequences) == 0:
			for b in range(0,min(self.batch_size,int(len(self.trajectory_set)*self.train_set))):
				id, trajectory = self.trajectory_set[self.data_idx]
				self.data_idx += 1
				self.batch_sequences.append(trajectory)
				self.batch_ids.append(id)
		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			if self.sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj): # Make sure trajs are long enough
				id, trajectory = self.trajectory_set[self.data_idx]
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set)*self.train_set)
				if self.data_idx == 0:
					new_epoch = True
				self.batch_sequences[ii] = trajectory
				self.batch_ids[ii] = id
				self.sequence_idx[ii] = self.args.prev_horizon
				self.sequence_reset[ii] = 1
			else:
				self.sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,min(self.batch_size,len(self.trajectory_set)-len(self.trajectory_set)%self.batch_size)):
			traj = self.batch_sequences[ii]
			agent_id = self.batch_ids[ii]
			other_agents_pos.append(self.fillBatch(agent_id, ii, int(self.sequence_idx[ii]), self.tbpl, self.batch_x, self.batch_vel, self.batch_pos,self.batch_grid, self.pedestrian_grid, self.batch_goal, self.batch_y, traj,self.batch_pos_target, centered_grid=self.centered_grid))
			self.sequence_idx[ii] += self.tbpl

		if self.args.rotated_grid:
			_, self.batch_y = sup.rotate_batch_to_local_frame(self.batch_y,self.batch_x)
			self.batch_x, self.batch_pos_target = sup.rotate_batch_to_local_frame(self.batch_pos_target, self.batch_x)

		return deepcopy(self.batch_x), deepcopy(self.batch_vel), deepcopy(self.batch_pos), deepcopy(self.batch_goal), \
						deepcopy(self.batch_grid), deepcopy(self.pedestrian_grid), deepcopy(self.batch_y), deepcopy(self.batch_pos_target), \
						deepcopy(other_agents_pos), new_epoch

	def getTestBatch(self):
		"""
		Get the next batch for model validation
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled

		self.val_data_idx = max(self.val_data_idx,int(len(self.trajectory_set)*self.train_set))
		trajectory=[]
		if len(self.val_batch_sequences) == 0:
			for b in range(0,self.batch_size):
				id, trajectory = self.trajectory_set[self.val_data_idx]
				self.val_data_idx += 1
				self.val_batch_sequences.append(trajectory)
				self.val_batch_ids.append(id)
				if self.val_data_idx == len(self.trajectory_set):
					self.val_data_idx =int (len(self.trajectory_set) * self.train_set)

		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		for ii, traj in enumerate(self.val_batch_sequences):
			if self.val_sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
				id, trajectory = self.trajectory_set[self.val_data_idx]
				self.val_data_idx = (self.val_data_idx + 1) % int(len(self.trajectory_set)*(1-self.train_set)-1) + int(len(self.trajectory_set)*self.train_set)
				self.val_batch_sequences[ii] = trajectory
				self.val_batch_ids[ii] = id
				self.val_sequence_idx[ii] = self.args.prev_horizon
				self.val_sequence_reset[ii] = 1
			else:
				self.val_sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,min(self.batch_size,len(self.trajectory_set)-len(self.trajectory_set)%self.batch_size)):
			traj = self.val_batch_sequences[ii]
			agent_id = self.val_batch_ids[ii]
			other_agents_pos.append(self.fillBatch(agent_id, ii, int(self.val_sequence_idx[ii]), self.tbpl, self.val_batch_x,
									self.val_batch_vel, self.val_batch_pos,self.val_batch_grid, self.val_pedestrian_grid, self.val_batch_goal,
									self.val_batch_y, traj,self.val_batch_pos_target, centered_grid=self.centered_grid))
			self.val_sequence_idx[ii] += self.tbpl

		if self.args.rotated_grid:
			_, self.val_batch_y = sup.rotate_batch_to_local_frame(self.val_batch_y,self.val_batch_x)
			self.val_batch_x, self.val_batch_pos_target = sup.rotate_batch_to_local_frame(self.val_batch_pos_target, self.val_batch_x)

		# Create dictionary to feed into the model
		dict = {"batch_x": self.val_batch_x,
				"batch_vel": self.val_batch_vel,
				"batch_pos": self.val_batch_pos,
				"batch_goal": self.val_batch_goal,
				"batch_grid": self.val_batch_grid,
				"batch_ped_grid": self.val_pedestrian_grid,
				"batch_y": self.val_batch_y,
				"batch_div": self.val_batch_y,
				"batch_pos_target": self.val_batch_pos_target,
				"state_noise": 0.0,
				"grid_noise": 0.0,
				"ped_noise":0.0,
				"other_agents_pos": other_agents_pos
				}

		return dict

	def getTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000, freeze = False):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		id = self.trajectory_set[trajectory_idx][0]
		traj = self.trajectory_set[trajectory_idx][1]
		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])
		batch_pos_target = np.zeros([1, sequence_length, 2*self.args.prediction_horizon])
		if "future" in self.args.others_info:
			batch_vel = np.zeros(
				[1, sequence_length, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													int(self.submap_size),
													int(self.submap_size)])
		
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * (self.output_sequence_length +3)]) #longer prev horizon, tests
		#batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		if self.args.others_info == "relative":
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		elif self.args.others_info == "ped_grid":
			pedestrian_grid = np.zeros([1, sequence_length,
								   int(self.submap_size),
								   int(self.submap_size)])
		else:
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
		other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj,batch_pos_target, centered_grid=self.centered_grid)

		if freeze:
			pedestrian_grid = pedestrian_grid*0.0

		return batch_x, batch_vel, batch_pos, batch_goal, batch_grid, pedestrian_grid, batch_y, batch_pos_target, other_agents_pos, traj

	def getGroupOfTrajectoriesAsBatch(self, trajectory_idx, max_sequence_length=1000, n_trajs=1):
		"""
		Get a group of trajectories out o the trajectory set in the same format as for the standard training data
		(e.g. for group prediction).
		"""

		batch_x_list = []
		batch_vel_list = []
		batch_pos_list = []
		batch_goal_list = []
		batch_grid_list = []
		pedestrian_grid_list = []
		batch_y_list = []
		batch_pos_final_list = []
		other_agents_pos_list = []
		traj_list = []

		agent_id = self.trajectory_set[trajectory_idx][0]
		traj = self.trajectory_set[trajectory_idx][1]
		grid = self.agent_container.occupancy_grid.getSubmapByCoords(traj.pose_vec[0,0] * self.norm_const_x,
																																 traj.pose_vec[0,1] * self.norm_const_y,
																																 self.submap_size, self.submap_size)
		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])
		batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
		batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													 int(self.submap_size),
													 int(self.submap_size)])
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])

		other_agents_pos_list.append(self.fillBatch(agent_id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel, batch_pos,
										  batch_grid, pedestrian_grid, batch_goal, batch_y, traj, batch_pos_final,
										  centered_grid=self.centered_grid))

		batch_x_list.append(batch_x)
		batch_vel_list.append(batch_vel)
		batch_pos_list.append(batch_pos)
		batch_goal_list.append(batch_goal)
		batch_grid_list.append(batch_grid)
		pedestrian_grid_list.append(pedestrian_grid)
		batch_y_list.append(batch_y)
		batch_pos_final_list.append(batch_pos_final)
		traj_list.append(traj)

		for id in self.agent_container.agent_data.keys():
			if id != agent_id:
				t = self.agent_container.agent_data[id].getTrajectoryForTime(traj.time_vec[0])
				sequence_length = min(max_sequence_length,
									  t.pose_vec.shape[0] - self.output_sequence_length - self.prev_horizon)
				batch_x = np.zeros([1, sequence_length, (self.prev_horizon + 1) * self.input_dim])
				batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
				batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
				batch_grid = np.zeros([1, sequence_length,
									   int(self.submap_size),
									   int(self.submap_size)])
				batch_goal = np.zeros([1, sequence_length, 2])
				batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
				batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
				pedestrian_grid = np.zeros(
					[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
				other_agents_pos_list.append(self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, t,batch_pos_final, centered_grid=self.centered_grid))
				batch_x_list.append(batch_x)
				batch_vel_list.append(batch_vel)
				batch_pos_list.append(batch_pos)
				batch_goal_list.append(batch_goal)
				batch_grid_list.append(batch_grid)
				pedestrian_grid_list.append(pedestrian_grid)
				batch_y_list.append(batch_y)
				batch_pos_final_list.append(batch_pos_final)
				traj_list.append(t)

		return batch_x_list, batch_vel_list, batch_pos_list,batch_goal_list, batch_grid_list, pedestrian_grid_list, batch_y_list, batch_pos_final_list, other_agents_pos_list, traj_list

	def getTestTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000,freeze=False):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		print(trajectory_idx)
		id = self.test_trajectory_set[trajectory_idx][0]
		traj = self.test_trajectory_set[trajectory_idx][1]

		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])  # data fed for training
		batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
		if "future" in self.args.others_info:
			batch_vel = np.zeros(
				[1, sequence_length, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													 int(self.submap_size),
													 int(self.submap_size)])
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		if self.args.others_info == "relative":
			pedestrian_grid = np.zeros([1,sequence_length, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		else:
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
		other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj,batch_pos_final, centered_grid=self.centered_grid,testing=True)

		if freeze:
			pedestrian_grid = pedestrian_grid*0.0

		return batch_x, batch_vel, batch_pos,batch_goal, batch_grid, pedestrian_grid, batch_y, other_agents_pos, traj

	def getAgentDataAtTimestep(self, query_time):
		"""
		Get data from all agents which are active at a certain time.
		"""
		agent_dict = {}  # key is agent id
		for id in self.agent_container.agent_data.keys():
			traj_for_time = self.agent_container.agent_data[id].getTrajectoryForTime(query_time)
			if traj_for_time != None:
				pose = traj_for_time.getPoseAtTime(query_time)
				vel = traj_for_time.getVelAtTime(query_time)
				heading = math.atan2(vel[1], vel[0])
				grid = self.agent_container.occupancy_grid.getSubmapByCoords(pose[0], pose[1], self.submap_size, self.submap_size)
				rotated_grid = sup.rotate_grid_around_center(grid, -heading * 180 / math.pi)
				radial_pedestrian_grid = np.zeros([self.pedestrian_vector_dim])  # Initialize pedestriangrid with 0, will be evaluated later
				goal = traj_for_time.goal
				future_pos_gt = np.zeros([0, 2])
				future_vel_gt = np.zeros([0, 2])
				for future_step in range(1, self.output_sequence_length+1):
					t_future = query_time + future_step * self.dt
					if t_future <= traj_for_time.getMaxTime():
						pose_step = traj_for_time.getPoseAtTime(t_future)
						vel_step = traj_for_time.getVelAtTime(t_future)
						future_pos_gt = np.append(future_pos_gt, np.expand_dims(pose_step[:2], axis=0), axis=0)
						future_vel_gt = np.append(future_vel_gt, np.expand_dims(vel_step[:2], axis=0), axis=0)

				agent_dict[id] = [pose, vel, rotated_grid, radial_pedestrian_grid, goal, future_pos_gt, future_vel_gt]

		# For each pedestrian, find other agents position for this time and compute pedestrian grid
		for id in agent_dict.keys():
			other_agents_pos = np.zeros([0,2])
			for other_id in agent_dict.keys():
				if other_id != id:
					other_agents_pos = np.append(other_agents_pos, np.expand_dims(agent_dict[other_id][0][:2], axis=0), axis=0)
			pos = agent_dict[id][0][:2]
			vel = agent_dict[id][1][:2]
			heading = math.atan2(vel[1], vel[0])
			other_pos_local_frame = sup.positions_in_local_frame(pos, heading, other_agents_pos)
			agent_dict[id][3] = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																													max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																													normalize=True)

		return agent_dict

	def plot_global_scenario(self,batch_grid,batch_x,batch_agent_real_traj,batch_goal,other_agents_pos,batch_agent_predicted_traj,t,seq_index):
		"""
			inputs:
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_x: batch of initial states [x,y, vx, vy] on the global frame
				batch_agent_real_traj: agent ground truth velocity in global frame also know as batch_y
				batch_goal: batch with the trajectory goal positons
				other_agents_pos: other agent position in the global frame
				batch_agent_predicted_traj: agent predicted velocities in global frame
			"""
		global_grid = self.agent_container.occupancy_grid
		fig = pl.figure("Global Trajectory Predictions")
		ax_in = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)

		ax_in.clear()
		#plot scenario grid
		sup.plot_grid(ax_in, np.array([0.0, 0.0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
		ax_in.set_xlim([-global_grid.center[0], global_grid.center[0]])
		ax_in.set_ylim([-global_grid.center[1], global_grid.center[1]])
		#ax_in.set_aspect("equal")

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index, t, 1]*self.norm_const_y
		ax_in.plot(x_init_global_frame, y_init_global_frame, color='g', marker='o', label='Agent initial pos')
		ax_in.plot(batch_goal[seq_index,t, 0]*self.norm_const_x, batch_goal[seq_index,t, 1]*self.norm_const_y,
							 color='purple', marker='o', label='Agent goal')

		#plot predicted trajectory global frame
		vel_pred = np.zeros((self.args.prediction_horizon, 2))
		for mix_idx in range(self.args.n_mixtures):
			# plot predicted trajectory global frame
			for pred_step in range(self.args.prediction_horizon):
				idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
				idy = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures
				mu_x = batch_agent_predicted_traj[seq_index, t, idx]
				mu_y = batch_agent_predicted_traj[seq_index, t, idy]
				vel_pred[pred_step, :] = [mu_x, mu_y]

			traj_pred = sup.path_from_vel(initial_pos=np.array([batch_x[seq_index,t, 0]*self.norm_const_x, batch_x[seq_index,t, 1]*self.norm_const_y]),
																		pred_vel=vel_pred, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
		x_pred = traj_pred[:, 0] #* self.norm_const_x
		y_pred = traj_pred[:, 1] #* self.norm_const_y
		ax_in.plot(x_pred, y_pred, color='r',label='Predicted trajectory')

		#plot real trajectory global frame
		color_value = scalar_color_map.to_rgba(r)
		real_traj = np.zeros((self.args.prediction_horizon, 2))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = batch_agent_real_traj[seq_index, t, idx]
			mu_y = batch_agent_real_traj[seq_index, t, idy]
			real_traj[i, :] = [mu_x, mu_y]
		real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([batch_x[seq_index, t, 0]*self.norm_const_x, batch_x[seq_index, t, 1]*self.norm_const_y]),
												pred_vel=real_traj, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
		x = real_traj_global_frame[:, 0] #* self.norm_const_x
		y = real_traj_global_frame[:, 1] #* self.norm_const_y
		ax_in.plot(x, y, marker='x',color='b',label='Real trajectory')

		#Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			try:
				ax_in.plot(other_pos[0, 0], other_pos[0, 1], color='r', marker='o', label='Other agents')
				for jj in range(1,other_pos.shape[0]):
					ax_in.plot(other_pos[jj, 0], other_pos[jj, 1], color='r', marker='o') #other agent pos positon
			except IndexError:
				print("Oops!  That was no valid number.  Try again...")
		ax_in.legend()
		fig.canvas.draw()
		pl.show(block=False)

	def plot_local_scenario(self,batch_grid,batch_x,batch_agent_real_traj,batch_goal,other_agents_pos,batch_agent_predicted_traj,t,seq_index):
		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth trajectory in global frame
				batch_agent_predicted_traj: agent predicted trajectory in global frame
			"""
		#pl.close('all')
		global_grid = self.agent_container.occupancy_grid

		fig = pl.figure("Local Trajectory Predictions")
		ax_out = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)
		color_value = scalar_color_map.to_rgba(r)
		#Convert initial position to local frame
		#ini_pos_local_frame = sup.rotate_batch_x_to_local_frame(batch_x)
		ax_out.clear()

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index, t, 1]*self.norm_const_y

		# plot initial and goal pose
		ax_out.plot(0, 0, color='g', marker='o', label='Agent initial pos')

		x_lim = [-self.submap_size/2, self.submap_size/2]
		y_lim = [-self.submap_size/2, self.submap_size/2]
		ax_out.set_xlim(x_lim)
		ax_out.set_ylim(y_lim)

		# plot predicted trajectory local frame
		vel_pred = np.zeros((self.args.prediction_horizon, 2))
		for mix_idx in range(self.args.n_mixtures):
			# plot predicted trajectory global frame
			for pred_step in range(self.args.prediction_horizon):
				idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
				idy = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures
				mu_x = batch_agent_predicted_traj[seq_index, t, idx]
				mu_y = batch_agent_predicted_traj[seq_index, t, idy]
				vel_pred[pred_step, :] = [mu_x, mu_y]

			traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),
																		pred_vel=vel_pred, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
			x_pred = traj_pred[:, 0] #* self.norm_const_x
			y_pred = traj_pred[:, 1] #* self.norm_const_y
			ax_out.plot(x_pred, y_pred, color='r',label='Predicted trajectory')

		# plot real trajectory local frame does it make sense? isnt batch_y in the robot frame
		real_traj = np.zeros((self.args.prediction_horizon, 2))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = batch_agent_real_traj[seq_index, t, idx]
			mu_y = batch_agent_real_traj[seq_index, t, idy]
			real_traj[i, :] = [mu_x, mu_y]
		real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]),
																								 pred_vel=real_traj,
																							 dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)

		# Query Agent Local frame plot
		if real_traj[0, 0] > 0.1:
			sup.plot_grid(ax_out, np.array([3.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))
		elif real_traj[0, 0] < -0.1:
			sup.plot_grid(ax_out, np.array([-3.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))
		else:
			sup.plot_grid(ax_out, np.array([0.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))

		x = real_traj_global_frame[:, 0]
		y = real_traj_global_frame[:, 1]
		ax_out.plot(x, y, marker='x', color='b', label='Agent real trajectory')
		#ax_out.plot(batch_x[seq_index,:, 0] * self.norm_const_x-x_init_global_frame*np.ones_like(batch_x[seq_index,:, 0]),
		#            batch_x[seq_index, :, 1] * self.norm_const_y-y_init_global_frame*np.ones_like(batch_x[seq_index,:, 0]), color='k')

		# Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			heading = math.atan2(batch_x[seq_index, t, 3], batch_x[seq_index, t, 2])

			try:
				ax_out.plot(other_pos[0, 0]-x_init_global_frame, other_pos[0, 1]-y_init_global_frame, color='r', marker='o',label='Other agents')
				for jj in range(1,other_pos.shape[0]):
					ax_out.plot(other_pos[jj, 0]-x_init_global_frame, other_pos[jj, 1]-y_init_global_frame, color='r', marker='o')#,markersize=23)  # other agent pos positon
					#ax_out.annotate(str(jj), xy=(other_pos[jj, 0]-x_init_global_frame , other_pos[jj, 1]-y_init_global_frame))
			except IndexError:
				print("Oops!  That was no valid number.  Try again...")

		circle = pl.Circle((0, 0), self.max_range_ped_grid, color='blue',
											 fill=False)
		ax_out.add_artist(circle)
		ax_out.set_aspect("equal")
		ax_out.legend()
		fig.canvas.draw()
				#sleep(0.5)  # Time in seconds.
		pl. show(block=False)

	def normalize_angle_positive(self,angle):
		""" Normalizes the angle to be 0 to 2*pi
				It takes and returns radians. """
		return angle % (2.0 * np.pi)

	def normalize_angle(self,angle):
		""" Normalizes the angle to be -pi to +pi
				It takes and returns radians."""
		a = self.normalize_angle_positive(angle)
		if a > np.pi:
			a -= 2.0 * np.pi
		return a

	def shortest_angular_distance(self,from_angle, to_angle):
		""" Given 2 angles, this returns the shortest angular
				difference.  The inputs and ouputs are of course radians.

				The result would always be -pi <= result <= pi. Adding the result
				to "from" will always get you an equivelent angle to "to".
		"""
		return self.normalize_angle(to_angle - from_angle)

	def get_radius(self,cutoff, A,var):

		return np.sqrt(-2 * var * np.log(cutoff / A))

	def getGaussian(self,x,y,x0,y0,varx,vary,skew):

		dx = x - x0
		dy = y - y0

		h = np.linalg.norm(np.array([dx,dy]))
		angle = np.arctan2(dy, dx)
		mx = np.cos(angle - skew) * h
		my = np.sin(angle - skew) * h
		f1 = np.square(mx) / (2.0 * varx)
		f2 = np.square(my) / (2.0 * vary)

		return self.amplitude_ * np.exp(-(f1 + f2))

	def plot_angular_grid(self, batch_ped_grid,t,seq_index):
		if self.multi_pedestrian:
			fig = pl.figure("Angular grid")
			ax_ped_grid = pl.subplot()
			ax_ped_grid.clear()
			n_sectors = batch_ped_grid.shape[2]
			grid = batch_ped_grid[seq_index,t,:]
			grid = np.expand_dims(grid, axis=1)
			grid_flipped = np.zeros_like(grid)
			grid_flipped[0:int(self.pedestrian_vector_dim/2)] = grid[-int(self.pedestrian_vector_dim/2):]
			grid_flipped[int(self.pedestrian_vector_dim / 2):] = grid[0:int(self.pedestrian_vector_dim / 2)]
			sup.plot_radial_distance_vector(ax_ped_grid, grid_flipped, max_range=1.0, min_angle=0.0,
																			max_angle=2 * np.pi)
			ax_ped_grid.plot(30, 30, color='r', marker='o', markersize=4)
			ax_ped_grid.arrow(0, 0, 1, 0, head_width=0.1,head_length=self.max_range_ped_grid) #agent poiting direction
		# x- and y-range only need to be [-1, 1] since the pedestrian grid is normalized
			ax_ped_grid.set_xlim([-self.max_range_ped_grid-1, self.max_range_ped_grid+1])
			ax_ped_grid.set_ylim([-self.max_range_ped_grid-1, self.max_range_ped_grid+1])

			fig.canvas.draw()
			# sleep(0.5)  # Time in seconds.
			pl.show(block=False)
			sleep(0.5)  # Time in seconds.

	def query_agent_velocity_direction(self, batch_grid, batch_x, batch_agent_real_traj,other_agents_pos,t,seq_index):
		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth velocity in global frame
				batch_agent_predicted_traj: agent predicted trajectory in global frame
			"""
		#pl.close('all')
		global_grid = self.agent_container.occupancy_grid
		fig = pl.figure("Query agent direction")
		ax_in = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)

		number_samples =1
		ax_in.clear()
		#clear plots
		#plot scenario grid
		sup.plot_grid(ax_in, np.array([0.0, 0.0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
		ax_in.set_xlim([-global_grid.center[0], global_grid.center[0]])
		ax_in.set_ylim([-global_grid.center[1], global_grid.center[1]])

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index,t, 1]*self.norm_const_y
		ax_in.plot(x_init_global_frame, y_init_global_frame, color='g', marker='o', label='start')

		#plot real trajectory global frame
		color_value = scalar_color_map.to_rgba(r)

		#Plot velocity direction
		ax_in.arrow(x_init_global_frame, y_init_global_frame, batch_x[seq_index,t, 2]*self.norm_const_vx, batch_x[seq_index,t, 3]*self.norm_const_vy
				, head_width=0.1, head_length=self.max_range_ped_grid+0.5,alpha=0.4,label='Agent velocity direction') # angular grid detection distance
		circle = pl.Circle((x_init_global_frame,y_init_global_frame),self.max_range_ped_grid,color='blue',fill=False,label='Detection range')
		ax_in.add_artist(circle)
		#Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			for jj in range(other_pos.shape[0]):
				ax_in.plot(other_pos[jj, 0], other_pos[jj, 1], color='r', marker='o') #other agent pos positon

		ax_in.set_aspect("equal")
		fig.canvas.draw()
		#sleep(0.5)  # Time in seconds.
		pl. show(block=False)

build/lib/data_utils/DataHandlerLSTM.py
eliatrevisan/IA-MPPI-LBM
#!/usr/bin/env python
import sys
import numpy as np
# from scipy.linalg import hankel

if sys.version_info[0] < 3:
    sys.path.append('../src/data_utils')
    import Support as sup
    from DataHandler import *
else:
    import src.data_utils.Support as sup
    from src.data_utils.DataHandler import *
    
class DataHandlerLSTM_Keras(DataHandler):
    def __init__(self, args):
        super().__init__(args)
        # self.tbpl = self.prev_horizon+1
    
    def getTestTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000):
        """
        Get a trajectory out of the trajectory set in the same format as for the standard training data
        (e.g. for validation purposes).
        """
        print(trajectory_idx)
        id = self.test_trajectory_set[trajectory_idx][0]
        traj = self.test_trajectory_set[trajectory_idx][1]

        sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
        batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])  # data fed for training
        batch_vel = np.zeros([1, sequence_length, (self.prev_horizon + 1) * self.input_state_dim])
        batch_grid = np.zeros([1, sequence_length, int(np.ceil(self.submap_width / self.args.submap_resolution)), int(np.ceil(self.submap_height / self.args.submap_resolution))])
        batch_goal = np.zeros([1, sequence_length, 2])
        batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
        batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
        pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
        other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj, centered_grid=self.centered_grid,testing=True)

        vel_seq = 
        other_agents_info_seq = 

        # return batch_x, batch_vel, batch_pos,batch_goal, batch_grid, pedestrian_grid, batch_y, other_agents_pos, traj

        batch = {
            'x': batch_x,
            'vel': batch_vel,
            'pos': batch_pos,
            'goal': batch_goal,
            'grid': batch_grid,
            'other_agents_info': other_agents_info,
            'y': batch_target,
            'other_agents_pos': other_agents_pos,
            'other_agents_vel': other_agents_vel,
            'traj': traj,
            'vel_seq': batch_vel_seq,
            'other_agents_info_seq': other_agents_info_seq
        }

        return batch
		
    
    def fillBatch(self, agent_id, batch_idx, start_idx, truncated_backprop_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, trajectory, centered_grid=False,testing=False):
		"""
		Fill the data batches of batch_idx with data for all truncated backpropagation steps.
		"""
		other_agents_pos = []

		for tbp_step in range(truncated_backprop_length):

			# Input values
			query_time = trajectory.time_vec[start_idx + tbp_step]
			for prev_step in range(self.prev_horizon,-1,-1):
				current_pos = np.array([trajectory.pose_vec[start_idx + tbp_step - prev_step, 0], trajectory.pose_vec[
					                        start_idx + tbp_step - prev_step, 1]])
				current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step - prev_step, 0] * self.norm_const_vx, trajectory.vel_vec[
					                        start_idx + tbp_step - prev_step, 1] * self.norm_const_vy])

				if self.args.normalize_data:
					self.normalize_pos(current_pos)
					self.normalize_vel(current_vel)

				batch_x[batch_idx, tbp_step, prev_step*self.input_dim:(prev_step+1)*self.input_dim] = np.array([current_pos[0],
																										current_pos[1],
																										current_vel[0],
																										current_vel[1]])
				batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_vel[0],
																											current_vel[1]])

			heading = math.atan2(current_vel[1], current_vel[0])
			if not testing:
				if centered_grid:
					grid_center = current_pos
					grid = self.agent_container.occupancy_grid.getSubmapByCoords(grid_center[0],
																																			 grid_center[1],
																																			 self.submap_width, self.submap_height)
				else:
					grid = self.batch_sequences[batch_idx][1]
					grid_center = trajectory.pose_vec[0, 0:2]

				if self.rotated_grid:
					grid = sup.rotate_grid_around_center(grid, heading * 180 / math.pi)  # rotation in degrees
				else:
					heading = 0
					grid = sup.rotate_grid_around_center(grid, heading * 180 / math.pi)  # rotation in degrees

				batch_grid[batch_idx, tbp_step, :, :] = grid

			batch_goal[batch_idx, tbp_step, :] = trajectory.goal
			# Find positions of other pedestrians at the current timestep

			other_poses = trajectory.other_agents_positions[start_idx + tbp_step]
			n_other_agents = other_poses.shape[0]
			other_velocities = trajectory.other_agents_velocities[start_idx + tbp_step]
			other_agents_pos.append(other_poses)
			other_pos_local_frame = sup.positions_in_local_frame(current_pos, heading, other_poses)
			try:
				if self.args.relative_info:
					radial_pedestrian_grid = np.zeros([self.pedestrian_vector_dim])
					for ag_id in range(n_other_agents):
						radial_pedestrian_grid[ag_id*4:ag_id*4+4] = np.array([other_poses[ag_id,0] - current_pos[0],
						                                                      other_poses[ag_id, 1] - current_pos[1],
						                                                      other_velocities[ag_id, 0] - current_vel[0],
						                                                      other_velocities[ag_id, 1] - current_vel[1]])
					for ag_id in range(n_other_agents,int(self.args.pedestrian_vector_dim/4)):
						radial_pedestrian_grid[ag_id*4:ag_id*4+4] = np.array([100,100,0,0])

				else:
					radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																																max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																																normalize=True)
			except:
				radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
				                                                            max_range=self.max_range_ped_grid, min_angle=0,
				                                                            max_angle=2 * np.pi,
				                                                            normalize=True)

			pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			try:
				rel_pos = np.array([other_poses[0,0] - current_pos[0],other_poses[0, 1] - current_pos[1]])
				rel_vel = np.array([other_velocities[0, 0] - current_vel[0],
							              other_velocities[0, 1] - current_vel[1]])
			except:
				rel_pos = np.array([-100,-100])
				rel_vel = np.array([0,0])

			if (self.args.pedestrian_vector_dim == 4):
				pedestrian_grid[batch_idx, tbp_step, :] = np.concatenate((rel_pos, rel_vel), axis=0)
			elif (self.args.pedestrian_vector_dim == 3):
				pedestrian_grid[batch_idx, tbp_step, :] = np.concatenate((np.linalg.norm(rel_pos), rel_vel), axis=0)
			else:
				pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			# Output values
			for pred_step in range(self.output_sequence_length):
				vx = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 0]
				vy = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 1]
				px = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 0]
				py = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 1]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step] = vx
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = vy
				if self.args.normalize_data:
					self.normalize_vel(batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step:self.output_state_dim*pred_step+2])
				batch_pos[batch_idx, tbp_step, self.output_state_dim*pred_step] = px
				batch_pos[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = py


		return other_agents_pos
build/lib/data_utils/DataHandlerLSTM_Keras.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import os
import random
import sys
import math
import cv2
import pickle as pkl
import json
from pykalman import KalmanFilter
if sys.version_info[0] < 3:
	import Support as sup
	from Trajectory import *
	import AgentContainer as ped_cont
else:
	import src.data_utils.Support as sup
	from src.data_utils.Trajectory import *
	import src.data_utils.AgentContainer as ped_cont
from copy import deepcopy
import matplotlib.pyplot as pl
import matplotlib.animation as animation
from time import sleep
import random
from scipy.stats import multivariate_normal
#import matplotlib.pyplot as plt

from matplotlib.patches import Ellipse

class DataHandlerRoboat():
	"""
	Data handler for training an LSTM pedestrian prediction model
	"""
	def __init__(self, args):

		self.data_path = args.data_path
		self.batch_size = args.batch_size
		self.tbpl = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.output_sequence_length = args.prediction_horizon
		self.prediction_horizon = args.prediction_horizon
		self.input_dim = args.input_dim
		self.input_state_dim = args.input_state_dim
		self.output_state_dim = args.output_dim
		self.submap_size = args.submap_size
		self.submap_width_real = args.submap_span_real
		self.submap_height_real = args.submap_span_real
		self.centered_grid = args.centered_grid
		self.rotated_grid = args.rotated_grid
		self.multi_pedestrian = True
		self.pedestrian_radius = args.pedestrian_radius
		self.min_length_trajectory = self.tbpl + 1 + self.output_sequence_length + args.prev_horizon
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.max_range_ped_grid = args.max_range_ped_grid
		self.dt = args.dt
		self.n_mixture = args.n_mixtures
		self.output_pred_state_dim = args.output_pred_state_dim
		self.args = args
		self.real_world_data = args.real_world_data
		# Normalization constants
		self.norm_const_x = 1.0
		self.norm_const_y = 1.0
		self.norm_const_heading = 1.0
		self.norm_const_vx = 1.0
		self.norm_const_vy = 1.0
		self.norm_const_omega = 1.0
		self.min_pos_x = 1000
		self.min_pos_y = 1000
		self.max_pos_x = -1000
		self.max_pos_y = -1000
		self.min_vel_x = 1000
		self.min_vel_y = 1000
		self.max_vel_x = -1000
		self.max_vel_y = -1000
		self.train_set = args.train_set
		self.avg_traj_length = 0

		# Default Distribution Parameters
		self.amplitude_ = 1.0
		self.factor_ = 5.0
		self.cutoff_ = 1.0
		self.covar_ = 0.25
		self.var = multivariate_normal(mean=0, cov=1)

		# Data structure containing all the information about agents
		self.agent_container = ped_cont.AgentContainer()
		self.test_container = ped_cont.AgentContainer()

		self.agent_traj_idx = None  # indicates which trajectory index will be the next one per agent
		self.trajectory_set = []
		self.test_trajectory_set = []
		self.test_ids = []
		# HARDCODED, 4 datasets
		#self.data_idx = [0,0,0,0]
		#self.val_data_idx = [0,0,0,0]
		self.data_idx = 0
		self.val_data_idx = 0

		self.data_idx_log = []
		self.data_val_log = []
		self.datahandlers = []
		self.traj_set_idx = []
		self.val_traj_idx = []

		# Patch for updating the grid around a pedestrian position
		self.pedestrian_patch = None

		# Training variables
		self.sequence_reset = np.ones([self.batch_size])  # indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.sequence_idx = np.zeros([self.batch_size])+self.args.prev_horizon

		# Batch with info about other agents
		if "future" in self.args.others_info:
			self.batch_vel = np.zeros(
				[self.batch_size, self.tbpl, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			self.batch_vel = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		self.batch_x = np.zeros([self.batch_size, self.tbpl, self.input_dim*(self.prev_horizon+1)])  # data fed for training
		self.batch_grid = np.zeros([self.batch_size, self.tbpl,
																int(self.submap_size),
																int(self.submap_size)])
		if self.args.others_info == "relative":
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim*self.args.n_other_agents])
			self.val_pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.pedestrian_vector_dim*self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		elif self.args.others_info == "ped_grid":
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl,
										int(self.submap_size),
										int(self.submap_size)])
		else:
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])
			self.val_pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])
		self.batch_goal = np.zeros([self.batch_size, self.tbpl, 2])
		self.batch_y = np.zeros([self.batch_size, self.tbpl, self.output_state_dim * self.output_sequence_length])
		self.batch_sequences = []  # Sequence stack (tuples of trajectory + grid) where parts of length self.tbpl are taken from until a sequence is finished. As soon as one is finished, a new sequence is put in the stack.
		self.batch_ids = []
		self.batch_pos = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		self.batch_pos_target = np.zeros([self.batch_size, self.tbpl, 2 * self.output_sequence_length])

		# Validation variables
		self.val_sequence_reset = np.ones([self.batch_size])  # indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.val_sequence_idx = np.zeros([self.batch_size]) + self.args.prev_horizon
		# Batch with current query-agent velocities
		self.val_batch_vel = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		# Batch with query-agent current velocities and positions
		self.val_batch_x = np.zeros([self.batch_size, self.tbpl, self.input_dim * (self.prev_horizon + 1)])  # data fed for training
		# Query agent static environment information
		self.val_batch_grid = np.zeros([self.batch_size, self.tbpl,
									int(self.submap_size),
									int(self.submap_size)])

		# Batch of query-agent goal positions
		self.val_batch_goal = np.zeros([self.batch_size, self.tbpl, 2])
		# Batch with quer-agent next velocities
		self.val_batch_y = np.zeros([self.batch_size, self.tbpl, self.output_state_dim * self.output_sequence_length])
		# Sequence stack (tuples of trajectory + grid) where parts of length self.tbpl are taken from until a sequence is finished. As soon as one is finished, a new sequence is put in the stack.
		self.val_batch_sequences = []
		self.val_batch_ids = []
		# Batch of query-agent current position
		self.val_batch_pos = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		# Batch with quer-agent next positions
		self.val_batch_pos_target = np.zeros([self.batch_size, self.tbpl, 2 * self.output_sequence_length])

	def addAgentTrajectoriesToSet(self,agent_container,trajectory_set, id):
		"""
		Goes through all trajectories of agent and adds them to the member set if they fulfill the criteria.
		For all the time steps within the trajectory it also computes the positions of the other agents at that
		timestep in order to make training more efficient.
		"""
		for traj_idx, traj in enumerate(agent_container.getAgentTrajectories(id)):
			traj_with_collision = False
			if len(traj) > self.min_length_trajectory:
				#if traj.getMinTime() < 100:
				traj.updateInterpolators()
				# Find other agent's trajectories which overlap with each time step
				for time_idx in range(traj.time_vec.shape[0]):
					if self.multi_pedestrian:
						query_time = traj.time_vec[time_idx]
						other_agents_positions = agent_container.getAgentPositionsForTimeExclude(query_time, id)
						other_agents_velocities = agent_container.getAgentVelocitiesForTimeExclude(query_time, id)
			
						# Remove ego agent
						traj.other_agents_positions.append(other_agents_positions)
						traj.other_agents_velocities.append(other_agents_velocities)
				# TODO: MAYBE REMOVE IN THE FUTURE
				#if id != -1:
					#if not traj_with_collision:
				trajectory_set.append((id, traj))
	
	def setOtherAgentData(self, agent_container, trajectory_set, id):
		for traj_idx, traj in enumerate(agent_container.getAgentTrajectories(id)):
			traj_with_collision = False
			if len(traj) > self.min_length_trajectory:
				#traj.updateInterpolators()
				for time_idx in range(traj.time_vec.shape[0]):
					if self.multi_pedestrian:
						query_time = traj.time_vec[time_idx]
						other_agents_positions = agent_container.getAgentPositionsForTimeExclude(query_time, id)
						other_agents_velocities = agent_container.getAgentVelocitiesForTimeExclude(query_time, id)
						# Remove ego agent
						traj.other_agents_positions.append(other_agents_positions)
						traj.other_agents_velocities.append(other_agents_velocities)
				trajectory_set[traj_idx][1].other_agents_positions = traj.other_agents_positions
				trajectory_set[traj_idx][1].other_agents_velocities = traj.other_agents_velocities
				

	def processData(self, **kwargs):
		"""
		Processes the simulation or real-world data, depending on the usage.
		"""
		data_pickle = self.args.data_path + self.args.scenario + "/data" + str(self.args.prediction_horizon) + "_" + str(
			self.args.truncated_backprop_length)+ "_" + str(
			self.args.prev_horizon) + "_roboat.pickle"
		#data_pickle = self.args.data_path + self.args.scenario + "/data" + str(self.args.prediction_horizon) + "_" + str(
		#	self.args.truncated_backprop_length)+ "_" + str(
		#	self.args.prev_horizon) + "experiment3agents-norm.pickle"
		if os.path.isfile(data_pickle):
			self.loadTrajectoryData(data_pickle)
		else:
			if "roboat" in data_pickle:
				print("Processing roboat data.")
				self._process_roboat_data_(**kwargs)
			else:
				print("No roboat data present")

			self.saveTrajectoryData(data_pickle)

	def load_map(self,**kwargs):
		print("Extracting the occupancy grid ...")
		# Occupancy grid data

		with open(os.path.join(self.data_path+self.args.scenario, 'map.json')) as json_file:
			map_data = json.load(json_file)

		grid_map = pl.imread(self.data_path+self.args.scenario + '/map.png')
		
		if self.args.scenario == 'open_crossing' or self.args.scenario == "test_open_crossing":
			grid_map = np.round(grid_map)
		else:
			grid_map = np.round(grid_map[:,:,0])

		#map_data = np.load(os.path.join(self.data_path+self.args.scenario, 'map.npy'), encoding='latin1', allow_pickle=True)
		self.agent_container.occupancy_grid.gridmap = grid_map #map_data.item(0)['Map']  # occupancy values of cells
		self.agent_container.occupancy_grid.resolution = map_data['resolution'] #map_data.item(0)['Resolution']  # map resolution in [m / cell]
		self.agent_container.occupancy_grid.map_size = map_data['map_size'] #map_data.item(0)['Size']  # map size in [m]
		self.agent_container.occupancy_grid.center =  map_data['map_center'] #self.agent_container.occupancy_grid.map_size / 2.0

		#FIlter all -1 values present on MAP to zero.
		for ii in range(self.agent_container.occupancy_grid.gridmap.shape[0]):
			for jj in range(self.agent_container.occupancy_grid.gridmap.shape[1]):
				if self.agent_container.occupancy_grid.gridmap[ii, jj] == -1:
					self.agent_container.occupancy_grid.gridmap[ii, jj] = 0.0

	
	def _process_roboat_data_(self, **kwargs):
		"""
		Import the data from the log file stored in the directory of data_path.
		This method brings all the data into a suitable format for training.
		"""
		self.load_map(**kwargs)
		# Pedestrian data
		# [id, timestep (s), timestep (ns), pos x, pos y, yaw, vel x, vel y, omega, goal x, goal y]
		print("Extracting data from csv file")
		pedestrian_data = np.genfromtxt(os.path.join(self.data_path+self.args.scenario, 'total_log.csv'), delimiter=",")[1:, :]
		print("Data extracted")

		# Iterate through the data and fill the register
		for sample_idx in range(pedestrian_data.shape[0]):
			if pedestrian_data[sample_idx, 0] != -1:
				id = pedestrian_data[sample_idx, 0]
				timestamp = np.round(pedestrian_data[sample_idx, 1],1)# + pedestrian_data[sample_idx, 2] * 1e-9  # time in seconds28.35
				pose = np.zeros([1,3])
				vel = np.zeros([1,3])
				pose[:,0:2] = np.true_divide(pedestrian_data[sample_idx, 3:5], np.array([self.norm_const_x, self.norm_const_y]))
				vel[:,0:2] = np.true_divide(pedestrian_data[sample_idx, 5:7], np.array([self.norm_const_vx, self.norm_const_vy]))
				goal = np.true_divide(pedestrian_data[sample_idx, 7:], np.array([self.norm_const_x, self.norm_const_y]))
				self.agent_container.addDataSample(id, timestamp, pose, vel, goal)

		# Set the initial indices for agent trajectories (which trajectory will be returned when queried)
		self.agent_traj_idx = [0] * self.agent_container.getNumberOfAgents()

		# Subsample trajectories (longer discretization time) from dt=0.1 to dt=0.4
		print("Subsampling data..")
		for id in self.agent_container.getAgentIDs():
			for traj in self.agent_container.getAgentTrajectories(id):
				traj.subsample(int(self.args.dt*10))
		print("Sampling done.")

		for id in self.agent_container.getAgentIDs():
			for traj in self.agent_container.getAgentTrajectories(id):
				for t_id in range(1, traj.pose_vec.shape[0]):
					traj.vel_vec[t_id, :2] = (traj.pose_vec[t_id, :2] - traj.pose_vec[t_id-1, :2]) / self.dt

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id in self.agent_container.getAgentIDs():
			for traj_idx, traj in enumerate(self.agent_container.getAgentTrajectories(id)):
				if len(traj) > self.min_length_trajectory:
					traj.updateInterpolators()

		# Put all the trajectories in the trajectory set and randomize
		for id in self.agent_container.getAgentIDs():
			print("Processing agent {} / {}".format(id, self.agent_container.getNumberOfAgents()))
			# Adds trajectory if bigger than a minimum length and maximum size
			self.addAgentTrajectoriesToSet(self.agent_container,self.trajectory_set,id)
		
		self.compute_min_max_values()

		

	def compute_min_max_values(self):
		for traj_id in range(len(self.trajectory_set)):
			for t_id in range(1, self.trajectory_set[traj_id][1].pose_vec.shape[0]):
				self.trajectory_set[traj_id][1].vel_vec[t_id, :2] = (self.trajectory_set[traj_id][1].pose_vec[t_id, :2] -
																	 self.trajectory_set[traj_id][1].pose_vec[t_id - 1,
																	 :2]) / self.dt
				self.min_pos_x = min(self.min_pos_x,self.trajectory_set[traj_id][1].pose_vec[t_id,0])
				self.min_pos_y = min(self.min_pos_y, self.trajectory_set[traj_id][1].pose_vec[t_id, 1])
				self.max_pos_x = max(self.max_pos_x, self.trajectory_set[traj_id][1].pose_vec[t_id, 0])
				self.max_pos_y = max(self.max_pos_y, self.trajectory_set[traj_id][1].pose_vec[t_id, 1])
				self.min_vel_x = min(self.min_vel_x,self.trajectory_set[traj_id][1].vel_vec[t_id,0])
				self.min_vel_y = min(self.min_vel_y, self.trajectory_set[traj_id][1].vel_vec[t_id, 1])
				self.max_vel_x = max(self.max_vel_x, self.trajectory_set[traj_id][1].vel_vec[t_id, 0])
				self.max_vel_y = max(self.max_vel_y, self.trajectory_set[traj_id][1].vel_vec[t_id, 1])

		self.calc_scale()

		self.args.min_pos_x = self.min_pos_x
		self.args.min_pos_y = self.min_pos_y
		self.args.max_pos_x = self.max_pos_x
		self.args.max_pos_y = self.max_pos_y
		self.args.min_vel_x = self.min_vel_x
		self.args.min_vel_y = self.min_vel_y
		self.args.max_vel_x = self.max_vel_x
		self.args.max_vel_y = self.max_vel_y
		self.args.sx_vel = self.sx_vel
		self.args.sy_vel = self.sy_vel
		self.args.sx_pos = self.sx_pos
		self.args.sy_pos = self.sy_pos

	def unit_test_data_(self,map_args):
			"""
			Generates predefined trajectories to analyse the interaction
			"""
			if not os.path.isfile(self.data_path + self.args.scenario + '/unit_map.npy'):
				print("Creating unit map from png ...")
				sup.create_map_from_png(data_path=self.args.data_path + self.args.scenario, **map_args)

			map_data = np.load(os.path.join(self.data_path + self.args.scenario, 'map.npy'), encoding='latin1',
							   allow_pickle=True)
			self.agent_container.occupancy_grid.gridmap = map_data.item(0)['Map']  # occupancy values of cells
			self.agent_container.occupancy_grid.resolution = map_data.item(0)['Resolution']  # map resolution in [m / cell]
			self.agent_container.occupancy_grid.map_size = map_data.item(0)['Size']  # map size in [m]
			self.agent_container.occupancy_grid.center = self.agent_container.occupancy_grid.map_size / 2.0

			# FIlter all -1 values present on MAP to zero.
			for ii in range(self.agent_container.occupancy_grid.gridmap.shape[0]):
				for jj in range(self.agent_container.occupancy_grid.gridmap.shape[1]):
					if self.agent_container.occupancy_grid.gridmap[ii, jj] == -1:
						self.agent_container.occupancy_grid.gridmap[ii, jj] = 0.0

			# Iterate through the data and fill the register
			pose_agent_1 = np.array([[-15.0,0.0,0.0]])
			pose_agent_2 = np.array([[-13.0, 0.0,0.0]])
			vel = np.zeros([1, 3])
			v0 = 1.0
			goal = np.array([15, 0])
			timestamp = 0
			dt =0.4
			for sample_idx in range(int(320/dt/10)):
				id = 0
				timestamp += dt #self.args.dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 1
				pose_agent_2[:, 0] += v0 * dt

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)

			# Second test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 3
				timestamp += dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

			# Third test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 4
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0
				goal = np.array([15,0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 5
				pose_agent_2[:, 0] -= v0 * dt
				vel[:, 0] = -v0

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, -goal)

			# Forth test
			pose_agent_1 = np.array([[15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 6
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] -= v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = -v0
				goal = np.array([-15,0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

			# Fith test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, -1.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 7
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 8
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, -1])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)
			# Forth test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, 1.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 9
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 10
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, -1])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)
			# Forth test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 11
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 12
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)

			# Put all the trajectories in the trajectory set and randomize
			pool_list = []
			for id in self.test_container.getAgentIDs():
				print("Processing agent {} / {}".format(id, self.test_container.getNumberOfAgents()))
				# Adds trajectory if bigger than a minimum length and maximum size
				self.addAgentTrajectoriesToSet(self.test_container,self.test_trajectory_set,id)

			self.min_pos_x = -1.0
			self.min_pos_y = -1.0
			self.max_pos_x = 15.0
			self.max_pos_y = 15.0
			self.min_vel_x = -1.0
			self.min_vel_y = -1.0
			self.max_vel_x = 1.0
			self.max_vel_y = 1.0

			self.calc_scale()


	def calc_scale(self, keep_ratio=True):
		self.sx_vel = 1 / (self.max_vel_x - self.min_vel_x)
		self.sy_vel = 1 / (self.max_vel_y - self.min_vel_y)
		if keep_ratio:
			if self.sx_vel > self.sy_vel:
				self.sx_vel = self.sy_vel
			else:
				self.sy_vel = self.sx_vel

		self.sx_pos = 1 / (self.max_pos_x - self.min_pos_x)
		self.sy_pos = 1 / (self.max_pos_y - self.min_pos_y)
		if keep_ratio:
			if self.sx_pos > self.sy_pos:
				self.sx_pos = self.sy_pos
			else:
				self.sy_pos = self.sx_pos

	def normalize_vel(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[1] = (data[1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, 1] = (data[:, 1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_vel_y * shift) * self.sy_vel
		else:
			return False
		return data_copy

	def normalize_pos(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[1] = (data[1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, 1] = (data[:, 1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_pos_y * shift) * self.sy_pos
		else:
			return False
		return data_copy

	def denormalize_pos(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx_pos + self.min_pos_x * shift
			data_copy[1] = data[1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy_pos + self.min_pos_y * shift
		else:
			return False

		return data_copy

	def denormalize_vel(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx_vel + self.min_vel_x * shift
			data_copy[1] = data[1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy_vel + self.min_vel_y * shift
		else:
			return False

		return data_copy

	def filter_data(self, data):
		# Kalman filter considering a constant velocity model
		Transition_Matrix = [[1, 0, self.dt, 0], [0, 1, 0, self.dt], [0, 0, 1, 0], [0, 0, 0, 1]]
		Observation_Matrix = [[1, 0, 0, 0], [0, 1, 0, 0]]

		xinit = data[0, 0]
		yinit = data[0, 1]
		vxinit = (data[1, 0] - data[0, 0]) / self.dt
		vyinit = (data[1, 1] - data[0, 1]) / self.dt
		initstate = [xinit, yinit, vxinit, vyinit]
		initcovariance = 1.0e-3 * np.eye(4) #  1.0e-3 * np.eye(4)
		transistionCov = 1.0e-4 * np.eye(4) # 1.0e-4 * np.eye(4)
		observationCov = 2.0e-3 * np.eye(2) # 1.0e-2 * np.eye(2)   -6
		kf = KalmanFilter(transition_matrices=Transition_Matrix,
						  observation_matrices=Observation_Matrix,
						  initial_state_mean=initstate,
						  initial_state_covariance=initcovariance,
						  transition_covariance=transistionCov,
						  observation_covariance=observationCov)

		(filtered_state_means, filtered_state_covariances) = kf.filter(data)

		return filtered_state_means

	def saveTrajectoryData(self, save_path):
		print("Saving data to: '{}'".format(save_path))
		if not os.path.isdir(self.args.data_path + self.args.scenario):
			os.makedirs(self.args.data_path + self.args.scenario)

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id, traj in self.trajectory_set:
			traj.updateInterpolators()

		#if "test" not in self.args.scenario:
		random.shuffle(self.trajectory_set)

		data = {
			"trajectories" : self.trajectory_set,
			"agent_container" : self.agent_container,
			"min_pos_x" : self.min_pos_x,
			"min_pos_y" : self.min_pos_y,
			"max_pos_x" : self.max_pos_x,
			"max_pos_y" : self.max_pos_y,
			"min_vel_x" : self.min_vel_x,
			"min_vel_y" : self.min_vel_y,
			"max_vel_x" : self.max_vel_x,
			"max_vel_y" : self.max_vel_y
		}
		pkl.dump(data, open(save_path, 'wb'),protocol=2)

	def loadTrajectoryData(self, load_path):
		print("Loading data from: '{}'".format(load_path))
		self.file = open(load_path, 'rb')
		if sys.version_info[0] < 3:
			tmp_self = pkl.loads(self.file)#,encoding='latin1')
		else:
			tmp_self = pkl.load(self.file , encoding='latin1')
		self.trajectory_set = tmp_self["trajectories"]
		self.agent_container = tmp_self["agent_container"]

		try:
			self.min_pos_x = tmp_self["min_pos_x"]
			self.min_pos_y = tmp_self["min_pos_y"]
			self.max_pos_x = tmp_self["max_pos_x"]
			self.max_pos_y = tmp_self["max_pos_y"]
			self.min_vel_x = tmp_self["min_vel_x"]
			self.min_vel_y = tmp_self["min_vel_y"]
			self.max_vel_x = tmp_self["max_vel_x"]
			self.max_vel_y = tmp_self["max_vel_y"]
		except:
			print("Delete dataset to compute limits")
			self.max_x = 1
			self.max_y = 1
			self.min_x = -1
			self.min_y = -1

		self.calc_scale()

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id, traj in self.trajectory_set:
			traj.updateInterpolators()

	def getAgentTrajectory(self, agent_id):
		"""
		Return the next agent trajectory in the queue for the agent with id agent_id.
		"""
		trajectory = self.agent_container.agent_data[agent_id].trajectories[self.agent_traj_idx[agent_id]]
		self.agent_traj_idx[agent_id] = (self.agent_traj_idx[agent_id] + 1)  % self.agent_container.getNumberOfTrajectoriesForAgent(agent_id)
		return trajectory

	def getRandomAgentTrajectory(self, agent_id):
		"""
		Return a totally random trajectory for the agent with id agent_id.
		"""
		random_traj_idx = np.random.randint(0, len(self.agent_container.agent_data[agent_id].trajectories))
		return self.agent_container.agent_data[agent_id].trajectories[random_traj_idx]

	def fillBatch(self, dh, agent_id, batch_idx, start_idx, truncated_backprop_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, trajectory,batch_pos_target, centered_grid=False, testing = False):
		"""
		Fill the data batches of batch_idx with data for all truncated backpropagation steps.
		"""
		other_agents_pos = []

		for tbp_step in range(truncated_backprop_length):

			# Input values
			query_time = trajectory.time_vec[start_idx + tbp_step]
			for prev_step in range(self.prev_horizon,-1,-1):
				current_pos = np.array([trajectory.pose_vec[start_idx + tbp_step - prev_step, 0], trajectory.pose_vec[
											start_idx + tbp_step - prev_step, 1]])
				#current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step, 0], trajectory.vel_vec[
				#							start_idx + tbp_step - prev_step - prev_step, 1]])
				current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step - prev_step, 0], trajectory.vel_vec[
											start_idx + tbp_step - prev_step, 1]])

				if self.args.normalize_data:
					self.normalize_pos(current_pos)
					self.normalize_vel(current_vel)

				batch_x[batch_idx, tbp_step, prev_step*self.input_dim:(prev_step+1)*self.input_dim] = np.array([current_pos[0],
																										current_pos[1],
																										current_vel[0],
																										current_vel[1]])
				#batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([np.linalg.norm(current_vel),
				#                                                                                                              np.arctan2(current_vel[1],current_vel[0])])
				batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_vel[0],
																											current_vel[1]])
				batch_pos[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_pos[0],
																											current_pos[1]])
				
			heading = math.atan2(current_vel[1], current_vel[0])
			if centered_grid:
				grid_center = current_pos
				grid = dh.agent_container.occupancy_grid.getSubmapByCoords(grid_center[0], grid_center[1], self.submap_width_real, self.submap_height_real)
				grid = cv2.resize(grid, (60,60)) # Resize grid to (60,60) to match autoencoder input size
			if self.rotated_grid:
				grid = sup.rotate_grid_around_center(grid, heading * 180 / math.pi)  # rotation in degrees

			batch_grid[batch_idx, tbp_step, :, :] = grid	

			# Find positions of other pedestrians at the current timestep and order them by dstance to query agent
			other_positions = trajectory.other_agents_positions[start_idx + tbp_step]
			n_other_agents = other_positions.shape[0]
			other_velocities = trajectory.other_agents_velocities[start_idx + tbp_step]

			# Compute distance to other agents and order them
			other_poses_ordered = np.zeros((other_positions.shape[0], 6))
			other_poses_ordered[:, :2] = other_positions
			other_poses_ordered[:, 2:4] = other_velocities
			other_agents_pos.append(other_poses_ordered)

			current_pos = np.array(
				[trajectory.pose_vec[start_idx + tbp_step, 0], trajectory.pose_vec[start_idx + tbp_step, 1]])
			current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step, 0], trajectory.vel_vec[start_idx + tbp_step, 1]])
			heading = math.atan2(current_vel[1], current_vel[0])

			for ag_id in range(n_other_agents):
				other_poses_ordered[ag_id,4] = np.linalg.norm(other_poses_ordered[ag_id,:2] - current_pos)
				# ordered ids
				other_poses_ordered[ag_id,5] = ag_id
			other_poses_ordered= other_poses_ordered[other_poses_ordered[:, 4].argsort()]

			# other_poses_ordered = [ x, y, vx, vy, norm distance, id]

			if self.args.others_info == "relative":
				for ag_id in range(min(n_other_agents,self.args.n_other_agents)):
					rel_pos = np.array([other_poses_ordered[ag_id,0] - current_pos[0],other_poses_ordered[ag_id, 1] - current_pos[1]])*\
								  multivariate_normal.pdf(np.linalg.norm(np.array([other_poses_ordered[ag_id,:2] - current_pos])),
																			  mean=0.0,
																			  cov=self.args.relative_covariance)

					rel_vel = np.array([other_poses_ordered[ag_id,2] - current_vel[0],other_poses_ordered[ag_id, 3] - current_vel[1]])
					
					pedestrian_grid[batch_idx, tbp_step, ag_id*4:ag_id*4+4] = np.concatenate([rel_pos, rel_vel])
					#pedestrian_grid[batch_idx, tbp_step, ag_id, 4] = np.linalg.norm(rel_pos)
					#pedestrian_grid[batch_idx, tbp_step, ag_id, 5] = np.arctan2(rel_pos[1], rel_pos[0])
			

			elif self.args.others_info == "angular_grid":
				other_pos_local_frame = sup.positions_in_local_frame(current_pos, heading, other_positions)
				radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																															max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																															normalize=True)
				pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			# Output values
			for pred_step in range(self.output_sequence_length):
				vx = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 0]
				vy = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 1]
				px = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 0]
				py = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 1]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step] = vx #- current_vel[0]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = vy #- current_vel[1]
				if self.args.normalize_data:
					self.normalize_vel(batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step:self.output_state_dim*pred_step+2])
				batch_pos_target[batch_idx, tbp_step, self.output_state_dim*pred_step] = px - trajectory.pose_vec[start_idx + tbp_step, 0]
				batch_pos_target[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = py - trajectory.pose_vec[start_idx + tbp_step, 1]

			current_pos = trajectory.pose_vec[start_idx + tbp_step, :2]
			#delta_pos = trajectory.pose_vec[start_idx + tbp_step + self.prediction_horizon,:2] - current_pos
			delta_pos = trajectory.goal[0] - current_pos
			#batch_goal[batch_idx, tbp_step, 0] = np.linalg.norm(delta_pos)
			#batch_goal[batch_idx, tbp_step, 1] = np.arctan2(delta_pos[1],
			#                                                delta_pos[0])
			# Used for planning
			batch_goal[batch_idx, tbp_step, 0] = delta_pos[0]/np.linalg.norm(delta_pos) #trajectory.goal[0,0]
			batch_goal[batch_idx, tbp_step, 1] = delta_pos[1]/np.linalg.norm(delta_pos) #trajectory.goal[0,1]

		return other_agents_pos

	def CombineTrajectorySets(self, dpreps):
		print("Combining data into Roboat DataHandler...")
		n = len(dpreps)
		for dh in range(n):
			for t in range(len(dpreps[dh].trajectory_set)):
				self.trajectory_set.append( dpreps[dh].trajectory_set[t])
				self.datahandlers.append(dh)
		# Shuffle
		c = list(zip(self.trajectory_set, self.datahandlers))
		random.shuffle(c)
		self.trajectory_set, self.datahandlers = zip(*c)
		print("Sucessfully combined datasets!")

	def RemoveStops(self):
		print("Remove stopping trajectories...")
		for t in range(0, len(self.trajectory_set)):
			id, traj = self.trajectory_set[t]
			idx = 1
			while idx < len(traj):	# In other words, check for stopping, but only in the last 10 time steps 
				if np.linalg.norm(traj.pose_vec[-1-idx,:] - traj.pose_vec[-1,:]) < 1.0:
					idx += 1
				else:
					if idx == 1:
						break
					else: 
						traj.pose_vec = traj.pose_vec[:-1-(idx-1)]
						traj.other_agents_positions = traj.other_agents_positions[:-1-(idx-1)]
						traj.other_agents_velocities = traj.other_agents_velocities[:-1-(idx-1)]
						traj.time_vec = traj.time_vec[:-1-(idx-1)]
						traj.vel_vec = traj.vel_vec[:-1-(idx-1)]
						
						#if len(traj.pose_vec) < (self.tbpl + self.output_sequence_length + self.prev_horizon):
							# Remove the trajectory if its too short as a result of deleting time steps
							#print(self.tbpl + self.output_sequence_length + self.prev_horizon, len(traj.pose_vec))
							#self.trajectory_set = self.trajectory_set[:t] + self.trajectory_set[t+1:]

						break
		self.trajectory_set = tuple(item for item in self.trajectory_set if len(item[1].pose_vec) >= (self.tbpl + self.output_sequence_length + self.prev_horizon))
		print("Done!")
	
	def getRoboatBatch(self, dpreps):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled
		trajectory=[]
		if len(self.batch_sequences) == 0:
			for b in range(0,min(self.batch_size,int(len(self.trajectory_set)*self.train_set))):
				id, trajectory = self.trajectory_set[self.data_idx]
				self.traj_set_idx.append(self.data_idx)
				self.data_idx += 1
				self.batch_sequences.append(trajectory)
				self.batch_ids.append(id)
				 
		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			if self.sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj): # Make sure trajs are long enough
				id, trajectory = self.trajectory_set[self.data_idx]
				self.traj_set_idx[ii] = self.data_idx
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set)*self.train_set)
				if self.data_idx == 0:
					new_epoch = True
				self.batch_sequences[ii] = trajectory
				self.batch_ids[ii] = id
				self.sequence_idx[ii] = self.args.prev_horizon
				self.sequence_reset[ii] = 1
			else:
				self.sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		dhs = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,min(self.batch_size,len(self.trajectory_set)-len(self.trajectory_set)%self.batch_size)):
			traj = self.batch_sequences[ii]
			agent_id = self.batch_ids[ii]
			dh = self.datahandlers[self.traj_set_idx[ii]]
			other_agents_pos.append(self.fillBatch(dpreps[dh], agent_id, ii, int(self.sequence_idx[ii]), self.tbpl, self.batch_x, self.batch_vel, self.batch_pos,self.batch_grid, self.pedestrian_grid, self.batch_goal, self.batch_y, traj,self.batch_pos_target, centered_grid=self.centered_grid))
			self.sequence_idx[ii] += self.tbpl
			dhs.append(dh)
		
		
		if self.args.rotated_grid:
			_, self.batch_y = sup.rotate_batch_to_local_frame(self.batch_y,self.batch_x)
			self.batch_x, self.batch_pos_target = sup.rotate_batch_to_local_frame(self.batch_pos_target, self.batch_x)

		return deepcopy(self.batch_x), deepcopy(self.batch_vel), deepcopy(self.batch_pos), deepcopy(self.batch_goal), \
						deepcopy(self.batch_grid), deepcopy(self.pedestrian_grid), deepcopy(self.batch_y), deepcopy(self.batch_pos_target), \
						deepcopy(other_agents_pos), new_epoch, [self.batch_ids, dh]

	def getRoboatValidationBatch(self, dpreps):
		"""
		Get the next batch for model validation
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled

		self.val_data_idx = max(self.val_data_idx,int(len(self.trajectory_set)*self.train_set))
		trajectory=[]
		if len(self.val_batch_sequences) == 0:
			for b in range(0,self.batch_size):
				id, trajectory = self.trajectory_set[self.val_data_idx]
				self.val_traj_idx.append(self.val_data_idx)
				self.val_data_idx += 1
				self.val_batch_sequences.append(trajectory)
				self.val_batch_ids.append(id)
				if self.val_data_idx == len(self.trajectory_set):
					self.val_data_idx =int (len(self.trajectory_set) * self.train_set)

		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		for ii, traj in enumerate(self.val_batch_sequences):
			if self.val_sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
				id, trajectory = self.trajectory_set[self.val_data_idx]
				self.val_traj_idx[ii] = self.val_data_idx
				self.val_data_idx = (self.val_data_idx + 1) % int(len(self.trajectory_set)*(1-self.train_set)-1) + int(len(self.trajectory_set)*self.train_set)
				self.val_batch_sequences[ii] = trajectory
				self.val_batch_ids[ii] = id
				self.val_sequence_idx[ii] = self.args.prev_horizon
				self.val_sequence_reset[ii] = 1
			else:
				self.val_sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		dhs = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,min(self.batch_size,len(self.trajectory_set)-len(self.trajectory_set)%self.batch_size)):
			traj = self.val_batch_sequences[ii]
			agent_id = self.val_batch_ids[ii]
			dh = self.datahandlers[self.val_traj_idx[ii]]
			other_agents_pos.append(self.fillBatch(dpreps[dh], agent_id, ii, int(self.val_sequence_idx[ii]), self.tbpl, self.val_batch_x,
									self.val_batch_vel, self.val_batch_pos,self.val_batch_grid, self.val_pedestrian_grid, self.val_batch_goal,
									self.val_batch_y, traj,self.val_batch_pos_target, centered_grid=self.centered_grid))
			self.val_sequence_idx[ii] += self.tbpl
			dhs.append(dh)

		if self.args.rotated_grid:
			_, self.val_batch_y = sup.rotate_batch_to_local_frame(self.val_batch_y,self.val_batch_x)
			self.val_batch_x, self.val_batch_pos_target = sup.rotate_batch_to_local_frame(self.val_batch_pos_target, self.val_batch_x)

		# Create dictionary to feed into the model
		dict = {"batch_x": self.val_batch_x,
				"batch_vel": self.val_batch_vel,
				"batch_pos": self.val_batch_pos,
				"batch_goal": self.val_batch_goal,
				"batch_grid": self.val_batch_grid,
				"batch_ped_grid": self.val_pedestrian_grid,
				"batch_y": self.val_batch_y,
				"batch_div": self.val_batch_y,
				"batch_pos_target": self.val_batch_pos_target,
				"state_noise": 0.0,
				"grid_noise": 0.0,
				"ped_noise":0.0,
				"other_agents_pos": other_agents_pos
				}

		return dict, [self.val_batch_ids, dhs]	


	def getBatch(self, dpreps):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled
		trajectory=[]
		
		# Number of datahandlers
		n_dhs = len(dpreps)

		# Determine the (random) order of dataset to get traj from. Re
		dhs = np.random.choice(n_dhs, self.batch_size)
		
		# Select traj from random dataset and fill the trajectory_set
		if len(self.batch_sequences) == 0:
			for b in range(0,self.batch_size):
				# Choose dataset to pick traj from
				id, trajectory = dpreps[dhs[b]].trajectory_set[self.data_idx[dhs[b]]]
				self.data_idx_log.append(self.data_idx[dhs[b]])
				self.data_idx[dhs[b]] += 1
				self.batch_sequences.append(trajectory)
				self.batch_ids.append(id)

		# Right now, new epoch when one of the datasets has reached end of trajectory_set.
		# TODO: make sure new epoch when ALL data is taken

		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			if self.sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj): 
				id, trajectory =  dpreps[dhs[ii]].trajectory_set[self.data_idx[dhs[ii]]]
				self.data_idx_log[ii] = self.data_idx[dhs[ii]]
				self.data_idx[dhs[ii]] = (self.data_idx[dhs[ii]] + 1) % int(len(dpreps[dhs[ii]].trajectory_set)*self.train_set)
				if self.data_idx[dhs[ii]] == 0:
					# Make sure they're all set to zero when one dataset was fully passed through network
					for id in range(n_dhs):
						self.data_idx[id] = 0
					new_epoch = True
				self.batch_sequences[ii] = trajectory
				self.batch_ids[ii] = id
				self.sequence_idx[ii] = self.args.prev_horizon
				self.sequence_reset[ii] = 1
				
			else:
				self.sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,self.batch_size):
			traj = self.batch_sequences[ii]
			agent_id = self.batch_ids[ii]
			other_agents_pos.append(self.fillBatch(dpreps[dhs[ii]], agent_id, ii, int(self.sequence_idx[ii]), self.tbpl, self.batch_x, self.batch_vel, self.batch_pos,self.batch_grid, self.pedestrian_grid, self.batch_goal, self.batch_y, traj,self.batch_pos_target, centered_grid=self.centered_grid))
			self.sequence_idx[ii] += self.tbpl

		if self.args.rotated_grid:
			_, self.batch_y = sup.rotate_batch_to_local_frame(self.batch_y,self.batch_x)
			self.batch_x, self.batch_pos_target = sup.rotate_batch_to_local_frame(self.batch_pos_target, self.batch_x)

		return deepcopy(self.batch_x), deepcopy(self.batch_vel), deepcopy(self.batch_pos), deepcopy(self.batch_goal), \
						deepcopy(self.batch_grid), deepcopy(self.pedestrian_grid), deepcopy(self.batch_y), deepcopy(self.batch_pos_target), \
						deepcopy(other_agents_pos), new_epoch, [dhs, self.data_idx_log]

	def getTestBatch(self, dpreps):
		"""
		Get the next batch for model validation
		"""
		# Number of datahandlers
		n_dhs = len(dpreps)

		# Determine the (random) order of dataset to get traj from. Re
		dhs = np.random.choice(n_dhs, self.batch_size)

		## HARDCODED, amount of trajectories of smallest dataset.
		# Should later make sure ALL data is used
		self.total_len_small = int(len(dpreps[0].trajectory_set)*self.train_set)

		# Update sequences
		# If batch sequences are empty and need to be filled

		#self.val_data_idx = max(self.val_data_idx, self.total_len_small)
		
		#self.data_len = []
		#for i in range(n_dhs):
		#	self.data_len.append(int(len(dpreps[i].trajectory_set)*self.train_set))

		for i in range(n_dhs):
			self.val_data_idx[i] = max( self.val_data_idx[i], int(len(dpreps[i].trajectory_set)*self.train_set) ) 
		
		
		trajectory=[]
		if len(self.val_batch_sequences) == 0:
			for b in range(0,self.batch_size):
				id, trajectory = dpreps[dhs[b]].trajectory_set[self.val_data_idx[dhs[b]]]
				self.data_val_log.append(self.val_data_idx[dhs[b]])
				self.val_data_idx[dhs[b]] += 1
				self.val_batch_sequences.append(trajectory)
				self.val_batch_ids.append(id)
				if self.val_data_idx[dhs[b]] == len(dpreps[dhs[b]].trajectory_set):
					self.val_data_idx[dhs[b]] = int (len(dpreps[dhs[b]].trajectory_set) * self.train_set)

		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		for ii, traj in enumerate(self.val_batch_sequences):
			if self.val_sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
				id, trajectory = dpreps[dhs[ii]].trajectory_set[self.val_data_idx[dhs[ii]]]
				self.data_val_log[ii] = self.val_data_idx[dhs[ii]]
				self.val_data_idx[dhs[ii]] = (self.val_data_idx[dhs[ii]] + 1) % int(len(dpreps[dhs[ii]].trajectory_set)*(1-self.train_set)-1) + int(len(dpreps[dhs[ii]].trajectory_set)*self.train_set) 
				self.val_batch_sequences[ii] = trajectory
				self.val_batch_ids[ii] = id
				self.val_sequence_idx[ii] = self.args.prev_horizon
				self.val_sequence_reset[ii] = 1

			else:
				self.val_sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0, self.batch_size):
			traj = self.val_batch_sequences[ii]
			agent_id = self.val_batch_ids[ii]
			other_agents_pos.append(self.fillBatch(dpreps[dhs[ii]], agent_id, ii, int(self.val_sequence_idx[ii]), self.tbpl, self.val_batch_x,
									self.val_batch_vel, self.val_batch_pos,self.val_batch_grid, self.val_pedestrian_grid, self.val_batch_goal,
									self.val_batch_y, traj,self.val_batch_pos_target, centered_grid=self.centered_grid))
			self.val_sequence_idx[ii] += self.tbpl

		if self.args.rotated_grid:
			_, self.val_batch_y = sup.rotate_batch_to_local_frame(self.val_batch_y,self.val_batch_x)
			self.val_batch_x, self.val_batch_pos_target = sup.rotate_batch_to_local_frame(self.val_batch_pos_target, self.val_batch_x)

		# Create dictionary to feed into the model
		dict = {"batch_x": self.val_batch_x,
				"batch_vel": self.val_batch_vel,
				"batch_pos": self.val_batch_pos,
				"batch_goal": self.val_batch_goal,
				"batch_grid": self.val_batch_grid,
				"batch_ped_grid": self.val_pedestrian_grid,
				"batch_y": self.val_batch_y,
				"batch_div": self.val_batch_y,
				"batch_pos_target": self.val_batch_pos_target,
				"state_noise": 0.0,
				"grid_noise": 0.0,
				"ped_noise":0.0,
				"other_agents_pos": other_agents_pos
				}

		return dict, [dhs, self.data_val_log]

	def getPredefinedBatch(self, dpreps, log):

		dset = log[0]
		traj_idx = log[1]

		self.val_batch_sequences = []

		trajectory=[]
		if len(self.val_batch_sequences) == 0:
			for b in range(0,self.batch_size):
				id, trajectory = dpreps[dset[b]].trajectory_set[traj_idx[dset[b]]]
				self.val_batch_sequences.append(trajectory)
				self.val_batch_ids.append(id)
		"""
		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		for ii, traj in enumerate(self.val_batch_sequences):
			if self.val_sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
				id, trajectory = dpreps[dhs[ii]].trajectory_set[self.val_data_idx[dhs[ii]]]
				self.val_data_idx[dhs[ii]] = (self.val_data_idx[dhs[ii]] + 1) % int(len(dpreps[dhs[ii]].trajectory_set)*(1-self.train_set)-1) + int(len(dpreps[dhs[ii]].trajectory_set)*self.train_set) 
				self.val_batch_sequences[ii] = trajectory
				self.val_batch_ids[ii] = id
				self.val_sequence_idx[ii] = self.args.prev_horizon
				self.val_sequence_reset[ii] = 1
			else:
				self.val_sequence_reset[ii] = 0
		"""
		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0, self.batch_size):
			traj = self.val_batch_sequences[ii]
			agent_id = self.val_batch_ids[ii]
			other_agents_pos.append(self.fillBatch(dpreps[dset[ii]], agent_id, ii, int(self.val_sequence_idx[ii]), self.tbpl, self.val_batch_x,
									self.val_batch_vel, self.val_batch_pos,self.val_batch_grid, self.val_pedestrian_grid, self.val_batch_goal,
									self.val_batch_y, traj,self.val_batch_pos_target, centered_grid=self.centered_grid))
			self.val_sequence_idx[ii] += self.tbpl

		# Create dictionary
		dict = {"batch_x": self.val_batch_x,
				"batch_vel": self.val_batch_vel,
				"batch_pos": self.val_batch_pos,
				"batch_goal": self.val_batch_goal,
				"batch_grid": self.val_batch_grid,
				"batch_ped_grid": self.val_pedestrian_grid,
				"batch_y": self.val_batch_y,
				"batch_div": self.val_batch_y,
				"batch_pos_target": self.val_batch_pos_target,
				"state_noise": 0.0,
				"grid_noise": 0.0,
				"ped_noise":0.0,
				"other_agents_pos": other_agents_pos
				}

		return dict


	def getTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000, freeze = False):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		id = self.trajectory_set[trajectory_idx][0]
		traj = self.trajectory_set[trajectory_idx][1]
		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])
		batch_pos_target = np.zeros([1, sequence_length, 2*self.args.prediction_horizon])
		if "future" in self.args.others_info:
			batch_vel = np.zeros(
				[1, sequence_length, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													int(self.submap_size),
													int(self.submap_size)])
		
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * (self.output_sequence_length +3)]) #longer prev horizon, tests
		#batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		if self.args.others_info == "relative":
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		elif self.args.others_info == "ped_grid":
			pedestrian_grid = np.zeros([1, sequence_length,
								   int(self.submap_size),
								   int(self.submap_size)])
		else:
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
		other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj,batch_pos_target, centered_grid=self.centered_grid)

		if freeze:
			pedestrian_grid = pedestrian_grid*0.0

		return batch_x, batch_vel, batch_pos, batch_goal, batch_grid, pedestrian_grid, batch_y, batch_pos_target, other_agents_pos, traj

	def getGroupOfTrajectoriesAsBatch(self, trajectory_idx, max_sequence_length=1000, n_trajs=1):
		"""
		Get a group of trajectories out o the trajectory set in the same format as for the standard training data
		(e.g. for group prediction).
		"""

		batch_x_list = []
		batch_vel_list = []
		batch_pos_list = []
		batch_goal_list = []
		batch_grid_list = []
		pedestrian_grid_list = []
		batch_y_list = []
		batch_pos_final_list = []
		other_agents_pos_list = []
		traj_list = []

		agent_id = self.trajectory_set[trajectory_idx][0]
		traj = self.trajectory_set[trajectory_idx][1]
		grid = self.agent_container.occupancy_grid.getSubmapByCoords(traj.pose_vec[0,0] * self.norm_const_x,
																																 traj.pose_vec[0,1] * self.norm_const_y,
																																 self.submap_size, self.submap_size)
		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])
		batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
		batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													 int(self.submap_size),
													 int(self.submap_size)])
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])

		other_agents_pos_list.append(self.fillBatch(agent_id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel, batch_pos,
										  batch_grid, pedestrian_grid, batch_goal, batch_y, traj, batch_pos_final,
										  centered_grid=self.centered_grid))

		batch_x_list.append(batch_x)
		batch_vel_list.append(batch_vel)
		batch_pos_list.append(batch_pos)
		batch_goal_list.append(batch_goal)
		batch_grid_list.append(batch_grid)
		pedestrian_grid_list.append(pedestrian_grid)
		batch_y_list.append(batch_y)
		batch_pos_final_list.append(batch_pos_final)
		traj_list.append(traj)

		for id in self.agent_container.agent_data.keys():
			if id != agent_id:
				t = self.agent_container.agent_data[id].getTrajectoryForTime(traj.time_vec[0])
				sequence_length = min(max_sequence_length,
									  t.pose_vec.shape[0] - self.output_sequence_length - self.prev_horizon)
				batch_x = np.zeros([1, sequence_length, (self.prev_horizon + 1) * self.input_dim])
				batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
				batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
				batch_grid = np.zeros([1, sequence_length,
									   int(self.submap_size),
									   int(self.submap_size)])
				batch_goal = np.zeros([1, sequence_length, 2])
				batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
				batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
				pedestrian_grid = np.zeros(
					[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
				other_agents_pos_list.append(self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, t,batch_pos_final, centered_grid=self.centered_grid))
				batch_x_list.append(batch_x)
				batch_vel_list.append(batch_vel)
				batch_pos_list.append(batch_pos)
				batch_goal_list.append(batch_goal)
				batch_grid_list.append(batch_grid)
				pedestrian_grid_list.append(pedestrian_grid)
				batch_y_list.append(batch_y)
				batch_pos_final_list.append(batch_pos_final)
				traj_list.append(t)

		return batch_x_list, batch_vel_list, batch_pos_list,batch_goal_list, batch_grid_list, pedestrian_grid_list, batch_y_list, batch_pos_final_list, other_agents_pos_list, traj_list

	def getTestTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000,freeze=False):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		print(trajectory_idx)
		id = self.test_trajectory_set[trajectory_idx][0]
		traj = self.test_trajectory_set[trajectory_idx][1]

		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])  # data fed for training
		batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
		if "future" in self.args.others_info:
			batch_vel = np.zeros(
				[1, sequence_length, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													 int(self.submap_size),
													 int(self.submap_size)])
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		if self.args.others_info == "relative":
			pedestrian_grid = np.zeros([1,sequence_length, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		else:
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
		other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj,batch_pos_final, centered_grid=self.centered_grid,testing=True)

		if freeze:
			pedestrian_grid = pedestrian_grid*0.0

		return batch_x, batch_vel, batch_pos,batch_goal, batch_grid, pedestrian_grid, batch_y, other_agents_pos, traj

	def getAgentDataAtTimestep(self, query_time):
		"""
		Get data from all agents which are active at a certain time.
		"""
		agent_dict = {}  # key is agent id
		for id in self.agent_container.agent_data.keys():
			traj_for_time = self.agent_container.agent_data[id].getTrajectoryForTime(query_time)
			if traj_for_time != None:
				pose = traj_for_time.getPoseAtTime(query_time)
				vel = traj_for_time.getVelAtTime(query_time)
				heading = math.atan2(vel[1], vel[0])
				grid = self.agent_container.occupancy_grid.getSubmapByCoords(pose[0], pose[1], self.submap_size, self.submap_size)
				rotated_grid = sup.rotate_grid_around_center(grid, -heading * 180 / math.pi)
				radial_pedestrian_grid = np.zeros([self.pedestrian_vector_dim])  # Initialize pedestriangrid with 0, will be evaluated later
				goal = traj_for_time.goal
				future_pos_gt = np.zeros([0, 2])
				future_vel_gt = np.zeros([0, 2])
				for future_step in range(1, self.output_sequence_length+1):
					t_future = query_time + future_step * self.dt
					if t_future <= traj_for_time.getMaxTime():
						pose_step = traj_for_time.getPoseAtTime(t_future)
						vel_step = traj_for_time.getVelAtTime(t_future)
						future_pos_gt = np.append(future_pos_gt, np.expand_dims(pose_step[:2], axis=0), axis=0)
						future_vel_gt = np.append(future_vel_gt, np.expand_dims(vel_step[:2], axis=0), axis=0)

				agent_dict[id] = [pose, vel, rotated_grid, radial_pedestrian_grid, goal, future_pos_gt, future_vel_gt]

		# For each pedestrian, find other agents position for this time and compute pedestrian grid
		for id in agent_dict.keys():
			other_agents_pos = np.zeros([0,2])
			for other_id in agent_dict.keys():
				if other_id != id:
					other_agents_pos = np.append(other_agents_pos, np.expand_dims(agent_dict[other_id][0][:2], axis=0), axis=0)
			pos = agent_dict[id][0][:2]
			vel = agent_dict[id][1][:2]
			heading = math.atan2(vel[1], vel[0])
			other_pos_local_frame = sup.positions_in_local_frame(pos, heading, other_agents_pos)
			agent_dict[id][3] = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																													max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																													normalize=True)

		return agent_dict

	def plot_global_scenario(self,batch_grid,batch_x,batch_agent_real_traj,batch_goal,other_agents_pos,batch_agent_predicted_traj,t,seq_index):
		"""
			inputs:
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_x: batch of initial states [x,y, vx, vy] on the global frame
				batch_agent_real_traj: agent ground truth velocity in global frame also know as batch_y
				batch_goal: batch with the trajectory goal positons
				other_agents_pos: other agent position in the global frame
				batch_agent_predicted_traj: agent predicted velocities in global frame
			"""
		global_grid = self.agent_container.occupancy_grid
		fig = pl.figure("Global Trajectory Predictions")
		ax_in = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)

		ax_in.clear()
		#plot scenario grid
		sup.plot_grid(ax_in, np.array([0.0, 0.0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
		ax_in.set_xlim([-global_grid.center[0], global_grid.center[0]])
		ax_in.set_ylim([-global_grid.center[1], global_grid.center[1]])
		#ax_in.set_aspect("equal")

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index, t, 1]*self.norm_const_y
		ax_in.plot(x_init_global_frame, y_init_global_frame, color='g', marker='o', label='Agent initial pos')
		ax_in.plot(batch_goal[seq_index,t, 0]*self.norm_const_x, batch_goal[seq_index,t, 1]*self.norm_const_y,
							 color='purple', marker='o', label='Agent goal')

		#plot predicted trajectory global frame
		vel_pred = np.zeros((self.args.prediction_horizon, 2))
		for mix_idx in range(self.args.n_mixtures):
			# plot predicted trajectory global frame
			for pred_step in range(self.args.prediction_horizon):
				idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
				idy = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures
				mu_x = batch_agent_predicted_traj[seq_index, t, idx]
				mu_y = batch_agent_predicted_traj[seq_index, t, idy]
				vel_pred[pred_step, :] = [mu_x, mu_y]

			traj_pred = sup.path_from_vel(initial_pos=np.array([batch_x[seq_index,t, 0]*self.norm_const_x, batch_x[seq_index,t, 1]*self.norm_const_y]),
																		pred_vel=vel_pred, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
		x_pred = traj_pred[:, 0] #* self.norm_const_x
		y_pred = traj_pred[:, 1] #* self.norm_const_y
		ax_in.plot(x_pred, y_pred, color='r',label='Predicted trajectory')

		#plot real trajectory global frame
		color_value = scalar_color_map.to_rgba(r)
		real_traj = np.zeros((self.args.prediction_horizon, 2))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = batch_agent_real_traj[seq_index, t, idx]
			mu_y = batch_agent_real_traj[seq_index, t, idy]
			real_traj[i, :] = [mu_x, mu_y]
		real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([batch_x[seq_index, t, 0]*self.norm_const_x, batch_x[seq_index, t, 1]*self.norm_const_y]),
												pred_vel=real_traj, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
		x = real_traj_global_frame[:, 0] #* self.norm_const_x
		y = real_traj_global_frame[:, 1] #* self.norm_const_y
		ax_in.plot(x, y, marker='x',color='b',label='Real trajectory')

		#Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			try:
				ax_in.plot(other_pos[0, 0], other_pos[0, 1], color='r', marker='o', label='Other agents')
				for jj in range(1,other_pos.shape[0]):
					ax_in.plot(other_pos[jj, 0], other_pos[jj, 1], color='r', marker='o') #other agent pos positon
			except IndexError:
				print("Oops!  That was no valid number.  Try again...")
		ax_in.legend()
		fig.canvas.draw()
		pl.show(block=False)

	def plot_local_scenario(self,batch_grid,batch_x,batch_agent_real_traj,batch_goal,other_agents_pos,batch_agent_predicted_traj,t,seq_index):
		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth trajectory in global frame
				batch_agent_predicted_traj: agent predicted trajectory in global frame
			"""
		#pl.close('all')
		global_grid = self.agent_container.occupancy_grid

		fig = pl.figure("Local Trajectory Predictions")
		ax_out = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)
		color_value = scalar_color_map.to_rgba(r)
		#Convert initial position to local frame
		#ini_pos_local_frame = sup.rotate_batch_x_to_local_frame(batch_x)
		ax_out.clear()

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index, t, 1]*self.norm_const_y

		# plot initial and goal pose
		ax_out.plot(0, 0, color='g', marker='o', label='Agent initial pos')

		x_lim = [-self.submap_size/2, self.submap_size/2]
		y_lim = [-self.submap_size/2, self.submap_size/2]
		ax_out.set_xlim(x_lim)
		ax_out.set_ylim(y_lim)

		# plot predicted trajectory local frame
		vel_pred = np.zeros((self.args.prediction_horizon, 2))
		for mix_idx in range(self.args.n_mixtures):
			# plot predicted trajectory global frame
			for pred_step in range(self.args.prediction_horizon):
				idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
				idy = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures
				mu_x = batch_agent_predicted_traj[seq_index, t, idx]
				mu_y = batch_agent_predicted_traj[seq_index, t, idy]
				vel_pred[pred_step, :] = [mu_x, mu_y]

			traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),
																		pred_vel=vel_pred, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
			x_pred = traj_pred[:, 0] #* self.norm_const_x
			y_pred = traj_pred[:, 1] #* self.norm_const_y
			ax_out.plot(x_pred, y_pred, color='r',label='Predicted trajectory')

		# plot real trajectory local frame does it make sense? isnt batch_y in the robot frame
		real_traj = np.zeros((self.args.prediction_horizon, 2))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = batch_agent_real_traj[seq_index, t, idx]
			mu_y = batch_agent_real_traj[seq_index, t, idy]
			real_traj[i, :] = [mu_x, mu_y]
		real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]),
																								 pred_vel=real_traj,
																							 dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)

		# Query Agent Local frame plot
		if real_traj[0, 0] > 0.1:
			sup.plot_grid(ax_out, np.array([3.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))
		elif real_traj[0, 0] < -0.1:
			sup.plot_grid(ax_out, np.array([-3.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))
		else:
			sup.plot_grid(ax_out, np.array([0.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))

		x = real_traj_global_frame[:, 0]
		y = real_traj_global_frame[:, 1]
		ax_out.plot(x, y, marker='x', color='b', label='Agent real trajectory')
		#ax_out.plot(batch_x[seq_index,:, 0] * self.norm_const_x-x_init_global_frame*np.ones_like(batch_x[seq_index,:, 0]),
		#            batch_x[seq_index, :, 1] * self.norm_const_y-y_init_global_frame*np.ones_like(batch_x[seq_index,:, 0]), color='k')

		# Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			heading = math.atan2(batch_x[seq_index, t, 3], batch_x[seq_index, t, 2])

			try:
				ax_out.plot(other_pos[0, 0]-x_init_global_frame, other_pos[0, 1]-y_init_global_frame, color='r', marker='o',label='Other agents')
				for jj in range(1,other_pos.shape[0]):
					ax_out.plot(other_pos[jj, 0]-x_init_global_frame, other_pos[jj, 1]-y_init_global_frame, color='r', marker='o')#,markersize=23)  # other agent pos positon
					#ax_out.annotate(str(jj), xy=(other_pos[jj, 0]-x_init_global_frame , other_pos[jj, 1]-y_init_global_frame))
			except IndexError:
				print("Oops!  That was no valid number.  Try again...")

		circle = pl.Circle((0, 0), self.max_range_ped_grid, color='blue',
											 fill=False)
		ax_out.add_artist(circle)
		ax_out.set_aspect("equal")
		ax_out.legend()
		fig.canvas.draw()
				#sleep(0.5)  # Time in seconds.
		pl. show(block=False)

	def normalize_angle_positive(self,angle):
		""" Normalizes the angle to be 0 to 2*pi
				It takes and returns radians. """
		return angle % (2.0 * np.pi)

	def normalize_angle(self,angle):
		""" Normalizes the angle to be -pi to +pi
				It takes and returns radians."""
		a = self.normalize_angle_positive(angle)
		if a > np.pi:
			a -= 2.0 * np.pi
		return a

	def shortest_angular_distance(self,from_angle, to_angle):
		""" Given 2 angles, this returns the shortest angular
				difference.  The inputs and ouputs are of course radians.

				The result would always be -pi <= result <= pi. Adding the result
				to "from" will always get you an equivelent angle to "to".
		"""
		return self.normalize_angle(to_angle - from_angle)

	def get_radius(self,cutoff, A,var):

		return np.sqrt(-2 * var * np.log(cutoff / A))

	def getGaussian(self,x,y,x0,y0,varx,vary,skew):

		dx = x - x0
		dy = y - y0

		h = np.linalg.norm(np.array([dx,dy]))
		angle = np.arctan2(dy, dx)
		mx = np.cos(angle - skew) * h
		my = np.sin(angle - skew) * h
		f1 = np.square(mx) / (2.0 * varx)
		f2 = np.square(my) / (2.0 * vary)

		return self.amplitude_ * np.exp(-(f1 + f2))

	def plot_angular_grid(self, batch_ped_grid,t,seq_index):
		if self.multi_pedestrian:
			fig = pl.figure("Angular grid")
			ax_ped_grid = pl.subplot()
			ax_ped_grid.clear()
			n_sectors = batch_ped_grid.shape[2]
			grid = batch_ped_grid[seq_index,t,:]
			grid = np.expand_dims(grid, axis=1)
			grid_flipped = np.zeros_like(grid)
			grid_flipped[0:int(self.pedestrian_vector_dim/2)] = grid[-int(self.pedestrian_vector_dim/2):]
			grid_flipped[int(self.pedestrian_vector_dim / 2):] = grid[0:int(self.pedestrian_vector_dim / 2)]
			sup.plot_radial_distance_vector(ax_ped_grid, grid_flipped, max_range=1.0, min_angle=0.0,
																			max_angle=2 * np.pi)
			ax_ped_grid.plot(30, 30, color='r', marker='o', markersize=4)
			ax_ped_grid.arrow(0, 0, 1, 0, head_width=0.1,head_length=self.max_range_ped_grid) #agent poiting direction
		# x- and y-range only need to be [-1, 1] since the pedestrian grid is normalized
			ax_ped_grid.set_xlim([-self.max_range_ped_grid-1, self.max_range_ped_grid+1])
			ax_ped_grid.set_ylim([-self.max_range_ped_grid-1, self.max_range_ped_grid+1])

			fig.canvas.draw()
			# sleep(0.5)  # Time in seconds.
			pl.show(block=False)
			sleep(0.5)  # Time in seconds.

	def query_agent_velocity_direction(self, batch_grid, batch_x, batch_agent_real_traj,other_agents_pos,t,seq_index):
		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth velocity in global frame
				batch_agent_predicted_traj: agent predicted trajectory in global frame
			"""
		#pl.close('all')
		global_grid = self.agent_container.occupancy_grid
		fig = pl.figure("Query agent direction")
		ax_in = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)

		number_samples =1
		ax_in.clear()
		#clear plots
		#plot scenario grid
		sup.plot_grid(ax_in, np.array([0.0, 0.0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
		ax_in.set_xlim([-global_grid.center[0], global_grid.center[0]])
		ax_in.set_ylim([-global_grid.center[1], global_grid.center[1]])

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index,t, 1]*self.norm_const_y
		ax_in.plot(x_init_global_frame, y_init_global_frame, color='g', marker='o', label='start')

		#plot real trajectory global frame
		color_value = scalar_color_map.to_rgba(r)

		#Plot velocity direction
		ax_in.arrow(x_init_global_frame, y_init_global_frame, batch_x[seq_index,t, 2]*self.norm_const_vx, batch_x[seq_index,t, 3]*self.norm_const_vy
				, head_width=0.1, head_length=self.max_range_ped_grid+0.5,alpha=0.4,label='Agent velocity direction') # angular grid detection distance
		circle = pl.Circle((x_init_global_frame,y_init_global_frame),self.max_range_ped_grid,color='blue',fill=False,label='Detection range')
		ax_in.add_artist(circle)
		#Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			for jj in range(other_pos.shape[0]):
				ax_in.plot(other_pos[jj, 0], other_pos[jj, 1], color='r', marker='o') #other agent pos positon

		ax_in.set_aspect("equal")
		fig.canvas.draw()
		#sleep(0.5)  # Time in seconds.
		pl. show(block=False)

build/lib/data_utils/DataHandlerRoboat.py
eliatrevisan/IA-MPPI-LBM
#!/usr/bin/env python
import sys
import numpy as np
from scipy.linalg import hankel

if sys.version_info[0] < 3:
    sys.path.append('../src/data_utils')
    import Support as sup
    from DataHandler import *
else:
    import src.data_utils.Support as sup
    from src.data_utils.DataHandlerLSTM import *

class DataHandler_Keras(DataHandlerLSTM):
    def __init__(self, args):
        super().__init__(args)
        self.tbpl = self.prev_horizon+1
        
        # Batches size redefinition
        self.batch_grid = np.zeros((self.batch_size, self.prev_horizon+1, int(self.args.submap_width / self.args.submap_resolution),
                                    int(self.args.submap_height / self.args.submap_resolution)))
        self.batch_x = np.zeros((self.batch_size, self.prev_horizon+1, self.args.input_dim))
        self.batch_vel = np.zeros((self.batch_size, self.prev_horizon+1, self.args.input_state_dim))
        self.batch_pos = np.zeros((self.batch_size, self.prev_horizon+1, self.args.input_state_dim))
        self.batch_goal = np.zeros((self.batch_size, self.prev_horizon+1, 2))
        self.other_agents_info = np.zeros((self.batch_size, self.prev_horizon+1, self.args.pedestrian_vector_dim))
        self.other_agents_pos = np.zeros((self.batch_size, self.prev_horizon+1, self.args.pedestrian_vector_state_dim))
        self.other_agents_vel = np.zeros((self.batch_size, self.prev_horizon+1, self.args.pedestrian_vector_state_dim))
        self.batch_y = np.zeros((self.batch_size, self.prediction_horizon, self.args.output_dim))
        self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])

    def getBatch(self):
        """
        		Get the next batch of training data.
        		"""
        # Update sequences
        # If batch sequences are empty and need to be filled
        trajectory = []
        if len(self.batch_sequences) == 0:
            for b in range(0, min(self.batch_size, len(self.trajectory_set))):
                id, trajectory = self.trajectory_set[self.data_idx]
                self.data_idx += 1
                self.batch_sequences.append(trajectory)
                self.batch_ids.append(id)
        # If batch sequences are filled and can be used or need to be updated.
        other_agents_pos = []
        new_epoch = False
        for ii, traj in enumerate(self.batch_sequences):
            if self.sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
                id, trajectory = self.trajectory_set[self.data_idx]
                self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set) * self.train_set)
                if self.data_idx == 0:
                    new_epoch = True
                self.batch_sequences[ii] = trajectory
                self.batch_ids[ii] = id
                self.sequence_idx[ii] = self.args.prev_horizon
                self.sequence_reset[ii] = 1
            else:
                self.sequence_reset[ii] = 0

        # Fill the batch
        other_agents_pos = []
        for ii in range(0, min(self.batch_size, len(self.trajectory_set) - len(self.trajectory_set) % self.batch_size)):
            traj = self.batch_sequences[ii]
            agent_id = self.batch_ids[ii]
            other_agents_pos.append(
                self.fillBatch(agent_id, ii, int(self.sequence_idx[ii]), self.tbpl, self.batch_x, self.batch_vel,
                               self.batch_pos, self.batch_grid, self.pedestrian_grid, self.batch_goal, self.batch_y,
                               traj, centered_grid=self.centered_grid))
            self.sequence_idx[ii] += self.tbpl
        
        batch = {
            'x': self.batch_x,
            'vel': self.batch_vel,
            'pos': self.batch_pos,
            'goal': self.batch_goal,
            'grid': self.batch_grid,
            'other_agents_info': self.other_agents_info,
            'y': self.batch_y,
            'other_agents_pos': self.other_agents_pos,
            'other_agents_vel': self.other_agents_vel,
            'new_epoch': new_epoch
        }
        return batch

    def fillBatch(self, agent_id, batch_idx, start_idx, truncated_backprop_length, batch_x, batch_vel, batch_pos,
                  batch_grid, pedestrian_grid, batch_goal, batch_y, trajectory, centered_grid=False):
        """
				Fill the data batches of batch_idx with data for all truncated backpropagation steps.
				"""

        for prev_step in range(self.prev_horizon,-1,-1):

            # Input values
            current_pos = np.array([trajectory.pose_vec[start_idx - prev_step, 0], trajectory.pose_vec[start_idx - prev_step, 1]])
            current_vel = np.array([trajectory.vel_vec[start_idx - prev_step, 0] , trajectory.vel_vec[start_idx - prev_step, 1]])

            if self.args.normalize_data:
                self.normalize_pos(current_pos)
                self.normalize_vel(current_vel)

            batch_x[batch_idx, prev_step, :] = np.array([current_pos[0],current_pos[1],current_vel[0],current_vel[1]])
            batch_vel[batch_idx, prev_step, :] = np.array([current_vel[0],current_vel[1]])

            heading = math.atan2(current_vel[1], current_vel[0])

            # Find positions of other pedestrians at the current timestep
            other_poses = trajectory.other_agents_positions[start_idx - prev_step]
            other_agents_pos = other_poses
            n_other_agents = other_poses.shape[0]
            if n_other_agents>0:
                other_velocities = trajectory.other_agents_velocities[start_idx - prev_step]
                other_pos_local_frame = sup.positions_in_local_frame(current_pos, heading, other_poses)

                # TODO: it only works for one agent now
                rel_pos = other_poses - current_pos
                rel_vel = other_velocities - current_vel
                distance = np.linalg.norm(rel_pos)
                pedestrian_grid[batch_idx, prev_step, :] = np.concatenate(
                  (np.array([np.linalg.norm(rel_pos)]), rel_vel[0]))

        # Output values
        for pred_step in range(self.output_sequence_length):
            vx = trajectory.vel_vec[start_idx +1+ pred_step, 0]
            vy = trajectory.vel_vec[start_idx +1+ pred_step, 1]
            px = trajectory.pose_vec[start_idx +1+ pred_step, 0]
            py = trajectory.pose_vec[start_idx +1+ pred_step, 1]
            batch_y[batch_idx, pred_step, 0] = vx
            batch_y[batch_idx, pred_step, 1] = vy
            if self.args.normalize_data:
                self.normalize_vel(batch_y[batch_idx, pred_step,:])
            batch_pos[batch_idx, pred_step, 0] = px
            batch_pos[batch_idx, pred_step, 1] = py

        return other_agents_pos
    
    def getTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000,unit_testing=False):
        """
        Get a trajectory out of the trajectory set in the same format as for the standard training data
        (e.g. for validation purposes).
        """
        if unit_testing:
            traj = self.test_trajectory_set[trajectory_idx][1]
        else:
            traj = self.trajectory_set[trajectory_idx][1]

        sequence_length = min(max_sequence_length, len(traj) - self.prediction_horizon) - self.prev_horizon
        
        # Old data structures, used for plotting
        batch_x = np.zeros([sequence_length,self.prev_horizon+1, self.args.input_dim])  # data fed for training
        batch_pos = np.zeros([sequence_length,self.prev_horizon+1, self.args.input_state_dim])  # data fed for training
        batch_vel = np.zeros([sequence_length,self.prev_horizon+1, self.args.input_state_dim])  # data fed for training
        batch_goal = np.zeros([sequence_length, 2])
        batch_target = np.zeros([sequence_length, self.prediction_horizon,self.args.output_dim ])

        other_agents_info = np.zeros([sequence_length, self.prev_horizon+1, self.args.pedestrian_vector_dim])
        batch_grid = np.zeros((sequence_length, int(self.args.submap_width / self.args.submap_resolution),
                            int(self.args.submap_height / self.args.submap_resolution)))
        n_other_agents =  traj.other_agents_positions[0].shape[0]
        other_agents_pos = np.zeros((sequence_length, n_other_agents,self.args.pedestrian_vector_state_dim))
        other_agents_vel = np.zeros((sequence_length, self.args.pedestrian_vector_state_dim))
        
        # New data structures, fed to the network
        batch_vel_seq = np.zeros([sequence_length, self.prev_horizon+1, self.args.input_state_dim])
        other_agents_info_seq = np.zeros([sequence_length, self.prev_horizon+1, self.args.pedestrian_vector_dim])

        for batch_idx in range(sequence_length):
            other_agents_pos[batch_idx,:,:] = self.fillBatch(id, batch_idx, self.prev_horizon+batch_idx, sequence_length, batch_x, batch_vel, batch_pos,
                                              batch_grid, other_agents_info, batch_goal, batch_target, traj,
                                              centered_grid=self.centered_grid)

        batch = {
            'x': batch_x,
            'vel': batch_vel,
            'pos': batch_pos,
            'goal': batch_goal,
            'grid': batch_grid,
            'other_agents_info': other_agents_info,
            'y': batch_target,
            'other_agents_pos': other_agents_pos,
            'other_agents_vel': other_agents_vel,
            'traj': traj,
            'vel_seq': batch_vel_seq,
            'other_agents_info_seq': other_agents_info_seq
        }
        return batch
        
        
def expand_sequence(sequence_array, horizon): 
    # Sequence array has shape [features, time steps]
    # Expanded sequence will have shape [time steps, horizon, features]
    expanded_sequence = np.zeros((sequence_array.shape[1]-horizon+1,\
                                  horizon,\
                                  sequence_array.shape[0]))
    
    for i in range(sequence_array.shape[0]): # For each feature
        sequence = sequence_array[i, :]
        expanded_sequence[:, :, i] = hankel(sequence[0:horizon],\
                                            sequence[horizon-1:]).transpose()
    
    return expanded_sequence


def reduce_sequence(hankel_matrix):
    aux = []
    for feature_idx in range(hankel_matrix.shape[2]):
        aux.append( np.concatenate([hankel_matrix[0, :, feature_idx], hankel_matrix[1:, -1, feature_idx]], axis = 0) )
    return np.stack(aux, axis = 1)

build/lib/data_utils/DataHandler_Keras.py
eliatrevisan/IA-MPPI-LBM
"""
New independent Keras DataHandler
Requires Python 3 and TensorFlow 2

List of admitted input 1 types:
> 'vel'
> 'polarvel'

List of admitted input 2 types (both strings for position and velocity components need to be contained in input2_type):
> For the position component:
>> 'posabs'
>> 'posrel'
>> 'posinvdist'
>> 'polarposabs'
>> 'polarposrel'
>> 'polarposinvdist'
> For the velocity component:
>> 'velabs'
>> 'velrel'
>> 'polarvelabs'
>> 'polarvelreldiff'
>> 'polarvelrelcos'
"""
#!/usr/bin/env python
import sys
import numpy as np
import pickle as pkl
from numpy.linalg import norm
from scipy.linalg import hankel
from scipy.spatial.distance import pdist
from math import atan2, cos, sin

class DataHandler():
    def __init__(self, args):
        self.args = args
        self.tbpl = self.args.truncated_backprop_length # unused
        self.prev_horizon = self.args.prev_horizon
        self.prediction_horizon = self.args.prediction_horizon
        
        self.input1_type = self.args.input1_type
        self.input2_type = self.args.input2_type

        self.shuffle = True
        self.batch_traj_idx = 0
        self.batch_size = self.args.batch_size
        self.input1_batch = np.zeros( [self.batch_size, self.prev_horizon, self.args.input_dim] )
        self.input2_batch = np.zeros( [self.batch_size, self.prev_horizon, self.args.pedestrian_vector_dim] )
        self.target_batch = np.zeros( [self.batch_size, self.prediction_horizon, self.args.output_pred_state_dim] )
        
        self.trajectory_set = []
        """ 
        trajectory_set is a list of trajectories
        > Each trajectory is a list of dictionary variables representing each data point
        >> Each datapoint contains the follwing keys:
        >>> 'control_command': Command of the robot
        >>> 'future_positions': 
        >>> 'goal_position': Goal position of the robot
        >>> 'pedestrian_goal_position': Goal position of the pedestrian
        >>> 'pedestrian_state': Dictionary which contains the current 'position' and 'velocity' of the pedestrian
        >>> 'predicted_cmd': Not used
        >>> 'robot_state': Current position and orientation of the robot
        >>> 'robot_velocity': Current speed of the robot
        """
        
        self.file = None

    def load_data(self):
        min_length_trajectory = self.prev_horizon + self.prediction_horizon
        self.file = open(self.args.data_path + self.args.dataset, 'rb')
        tmp_trajectory_set = pkl.load(self.file, encoding='latin1')
        trajectory_set = []
        
        for traj in tmp_trajectory_set:
            sub_traj = []
            # Sub-sample trajectories to make dt = 0.4
            for t_id in range(0, len(traj), 4):
            # for t_id in range(0, len(traj), 2):
                sub_traj.append(traj[t_id])
            # Append zero velocities to goal position
            if traj:
                goal_data = traj[-1].copy()
                goal_data["pedestrian_state"]["velocity"] = np.array([0, 0])
                for t_id in range(10):
                    sub_traj.append(goal_data)
                if len(sub_traj) > min_length_trajectory:
                    trajectory_set.append(sub_traj)
        self.trajectory_set = trajectory_set

        input1_arrays = []
        input2_arrays = []
        Y_arrays = []
        
        for trajectory_idx in range(len(self.trajectory_set)):
            feedArrays = self.getTrajectoryAsBatch(trajectory_idx)
            input1_arrays.append( feedArrays['input'][0] )
            input2_arrays.append( feedArrays['input'][1] )
            Y_arrays.append( feedArrays['target'] )
            
        self.input1 = np.concatenate(input1_arrays, axis = 0)
        self.input2 = np.concatenate(input2_arrays, axis = 0)
        self.target = np.concatenate(Y_arrays, axis = 0)
        self.n_data_samples = self.target.shape[0]
        
        if self.shuffle:
            idxs = np.random.permutation(self.n_data_samples)
            self.input1 = self.input1[idxs]
            self.input2 = self.input2[idxs]
            self.target = self.target[idxs]    

    def getBatch(self):
        if self.batch_traj_idx + self.batch_size <= self.n_data_samples: # Regular batches
            self.input1_batch = self.input1[self.batch_traj_idx:self.batch_traj_idx+self.batch_size,:,:]
            self.input2_batch = self.input2[self.batch_traj_idx:self.batch_traj_idx+self.batch_size,:,:]
            self.target_batch = self.target[self.batch_traj_idx:self.batch_traj_idx+self.batch_size,:,:]
            self.batch_traj_idx += self.batch_size
            new_epoch = False
            
        else: # Last batch, fill with random data points
            remaining_samples = self.n_data_samples - self.batch_traj_idx
            self.input1_batch[0:remaining_samples, :, :] = self.input1[self.batch_traj_idx:, :, :]
            self.input2_batch[0:remaining_samples, :, :] = self.input2[self.batch_traj_idx:, :, :]
            self.target_batch[0:remaining_samples, :, :] = self.target[self.batch_traj_idx:, :, :]
            
            # Take random data points outside of the remaining ones in order to completely fill up the batches
            idxs = np.random.permutation(self.batch_traj_idx)[0:self.batch_size-remaining_samples] 
            self.input1_batch[remaining_samples:, :, :] = self.input1[idxs, :, :]
            self.input2_batch[remaining_samples:, :, :] = self.input2[idxs, :, :]
            self.target_batch[remaining_samples:, :, :] = self.target[idxs, :, :]
            
            self.batch_traj_idx = 0
            new_epoch = True

        return {"input": [self.input1_batch, self.input2_batch],
                "target": self.target_batch,
                "new_epoch": new_epoch}

    def getTrajectoryAsBatch(self, trajectory_idx):
        trajectory = self.trajectory_set[trajectory_idx]
        input1_seq_list = []
        input2_seq_list = []

        for sample in trajectory:
            input1_data = self.getInput1Data(sample)
            input2_data = self.getInput2Data(sample)

            input1_seq_list.append(input1_data)
            input2_seq_list.append(input2_data)
            
        vel_seq =np.stack(input1_seq_list, axis = 1) # (features, timesteps)
        other_seq = np.stack(input2_seq_list, axis = 1) # (features, timesteps)
        
        input1_feedArray = expand_sequence(vel_seq[:, 0:-self.prediction_horizon], self.prev_horizon)
        input2_feedArray = expand_sequence(other_seq[:, 0:-self.prediction_horizon], self.prev_horizon)
        target_feedArray = expand_sequence(vel_seq[:, self.prev_horizon:], self.prediction_horizon)
        
        return {
            'input': [input1_feedArray, input2_feedArray],
            'target': target_feedArray,
            'trajectory': trajectory
        }
        
    def getPaddedTrajectoryAsBatch(self, trajectory_idx):
        trajectory = self.trajectory_set[trajectory_idx]
        input1_seq_list = []
        input2_seq_list = []

        for sample in trajectory:            
            input1_data = self.getInput1Data(sample)
            input2_data = self.getInput2Data(sample)

            input1_seq_list.append(input1_data)
            input2_seq_list.append( input2_data )
        vel_seq =np.stack(input1_seq_list, axis = 1) # (features, timesteps)
        other_seq = np.stack(input2_seq_list, axis = 1) # (features, timesteps)

        padded_vel_seq = np.concatenate( [np.zeros((2, self.prev_horizon-1)), vel_seq], axis = 1 )
        padded_other_seq = np.concatenate( [np.zeros((4, self.prev_horizon-1)), other_seq], axis = 1 )

        input1_feedArray = expand_sequence(padded_vel_seq, self.prev_horizon)
        input2_feedArray = expand_sequence(padded_other_seq, self.prev_horizon)

        return {
            'input': [input1_feedArray, input2_feedArray],
            'trajectory': trajectory
        }
        
    def getInput1Data(self, sample):
        ped_vel = sample['pedestrian_state']['velocity']
        
        if self.input1_type == "vel": # Velocity in cartesian coordinates
            # Velocity as data
            data = ped_vel
        elif self.input1_type == "polarvel": # Velocity in polar coordinates
            # Velocity modulus and direction as data
            modulus = norm(ped_vel)
            direction = atan2(ped_vel[1], ped_vel[0])
            data = np.array([modulus, direction])
        else:
            raise Exception("No valid type for Input 1")
        
        return data
    
    def getInput2Data(self, sample):
        ped_pos = sample['pedestrian_state']['position']
        ped_vel = sample['pedestrian_state']['velocity']
        robot_pos = sample['robot_state'][0:2]
        robot_vel = sample['robot_velocity']
        
        # Position component
        if "polarpos" in self.input2_type: # Position in polar coordinates
            if "posabs" in self.input2_type:
                modulus = norm(robot_pos)
                direction = atan2(robot_pos[1], robot_pos[0])
                pos = np.array([modulus, direction])
            elif "posrel" in self.input2_type:
                distance = robot_pos - ped_pos
                modulus = norm(distance)
                direction = atan2(distance[1], distance[0])
                pos = np.array([modulus, direction])
            elif "posinvdist" in self.input2_type:
                distance = robot_pos - ped_pos
                invdist = np.zeros(2)
                invdist[0] = 1/(distance[0] + 1)
                invdist[1] = 1/(distance[1] + 1)
                modulus = norm(invdist)
                direction = atan2(invdist[1], invdist[0])
                pos = np.array([modulus, direction])
            else:
                raise Exception("No valid position-component type for Input 2")
            
        else: # Position in cartesian coordinates
            if "posabs" in self.input2_type:
                pos = robot_pos
            elif "posrel" in self.input2_type:
                pos = robot_pos - ped_pos
            elif "posinvdist" in self.input2_type:
                distance = robot_pos - ped_pos
                invdist = np.zeros(2)
                invdist[0] = 1/(distance[0] + 1)
                invdist[1] = 1/(distance[1] + 1)
                pos = invdist
            else:
                raise Exception("No valid position-component type for Input 2")
        
        # Velocity component
        if "polarvel" in self.input2_type: # Velocity in polar coordinates
            if "velabs" in self.input2_type: # Polar absolute velocity
                modulus = norm(robot_vel)
                direction = atan2(robot_vel[1], robot_vel[0])
                vel = np.array([modulus, direction])
            elif "velreldiff" in self.input2_type: # Polar relative velocity  
                modulus = norm(robot_vel - ped_vel)
                direction_rob = atan2(robot_vel[1], robot_vel[0])
                direction_ped = atan2(robot_vel[1], robot_vel[0])
                angular_distance = direction_rob - direction_ped
                if angular_distance > 180:
                    angular_distance -= 360
                elif angular_distance < -180:
                    angular_distance += 360 
                vel = np.array([modulus, angular_distance])
            elif "velrelcos" in self.input2_type: # Polar relative velocity using cosine distance
                modulus = norm(robot_vel - ped_vel)
                distance = pdist([robot_vel, ped_vel], 'cosine')[0]
                if np.isnan(distance):
                    distance = 0
                vel = np.array([modulus, distance])
            else:
                raise Exception("No valid velocity-component type for Input 2")
                
        else: # Velocity in cartesian coordinates
            if "velabs" in self.input2_type:
                vel = robot_vel
            elif "velrel" in self.input2_type:
                vel = robot_vel - ped_vel
            else:
                raise Exception("No valid velocity-component type for Input 2")
        
        data =  np.concatenate( [pos, vel] )
        
        return data


def prediction_to_trajectory(initial_position, prediction, args):
    dt = args.dt
    output_type = args.input1_type
    
    trajectory_list = [ initial_position ]

    for pred_step in prediction:
        if output_type == "vel": # Velocity in cartesian coordinates
            new_point = trajectory_list[-1] + pred_step * dt
            
        elif output_type == "polarvel": # Velocity in polar coordinates
            modulus = pred_step[0]
            direction = pred_step[1]
            new_point = trajectory_list[-1] + modulus * np.array((cos(direction), sin(direction))) * dt
        else:
            raise Exception("Unimplemented output type")
        
        trajectory_list.append(new_point)
    trajectory = np.stack(trajectory_list, axis = 0)
    return trajectory

def expand_sequence(sequence_array, horizon): 
    # Sequence array has shape [features, time steps]
    # Expanded sequence will have shape [time steps, horizon, features]
    expanded_sequence = np.zeros((sequence_array.shape[1]-horizon+1,\
                                  horizon,\
                                  sequence_array.shape[0]))
    
    for i in range(sequence_array.shape[0]): # For each feature
        sequence = sequence_array[i, :]
        expanded_sequence[:, :, i] = hankel(sequence[0:horizon],\
                                            sequence[horizon-1:]).transpose()
    
    return expanded_sequence

def reduce_sequence(hankel_matrix):
    aux = []
    for feature_idx in range(hankel_matrix.shape[2]):
        aux.append( np.concatenate([hankel_matrix[0, :, feature_idx], hankel_matrix[1:, -1, feature_idx]], axis = 0) )
    return np.stack(aux, axis = 1)

build/lib/data_utils/DataHandler_Keras_v2.py
eliatrevisan/IA-MPPI-LBM
import random

"""Code inspired in: https://github.com/fhennecker/deepdoom/blob/master/src/memory.py"""


class MemoryBuffer:
    def __init__(self, min_size=20, max_size=100):
        self.buffer = []
        self.min_size, self.max_size = min_size, max_size

    def full(self):
        return len(self.buffer) >= self.max_size

    def initialized(self):
        return len(self.buffer) >= self.min_size

    def add(self, step):
        if self.full():
            self.buffer.pop(0)
        self.buffer.append(step)

    def sample(self, batch_size):
        return [random.choice(self.buffer) for _ in range(batch_size)]
build/lib/data_utils/MemoryBuffer.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import cv2 as cv
import src.data_utils.Support as sup
import matplotlib.pyplot as plt
import matplotlib.patches as patches
class OccupancyGrid():
	"""
	Occupancy grid class for capturing static object information.
	This occupancy grid is aligned with the Cartesian coordinate frame: 
		index 0: x-axis
		index 1: y-axis
	"""

	def __init__(self):
		self.gridmap = None
		self.resolution = None
		self.map_size = None
		#self.center = np.array([0.0, 0.0])
		self.center = np.array([None, None])
		
	def getIdx(self, pos_x, pos_y):
		"""
		Get indices of position. 
		pos_x and pos_y are the positions w.r.t. the center of the map.
		"""

		# Custom function to extract correct indices from Herengracht map
		idx_y = int((pos_x / self.resolution) + (self.center[0]*-1 / self.resolution)) 
		idx_x = self.map_size[1] - int((pos_y / self.resolution) + (self.center[1]*-1 / self.resolution))

		
		# Projecting index on map if out of bounds
		idx_x = max(0, min(idx_x, self.map_size[0] / self.resolution))
		idx_y = max(0, min(idx_y, self.map_size[1] / self.resolution))

		return idx_x, idx_y
		
	
	def getSubmapByIndices(self, center_idx_x, center_idx_y, span_x, span_y):
		"""
		Extract a submap of span (span_x, span_y) around 
		center index (center_idx_x, center_idx_y)
		"""
		debug_info = {}
		start_idx_x = max(0, int(center_idx_x - np.floor(span_x / 2)))
		start_idx_y = max(0, int(center_idx_y - np.floor(span_y / 2)))
		
		# Compute end indices (assure size of submap is correct, if out pf bounds)
		max_idx_x = self.gridmap.shape[0] - 1
		max_idx_y = self.gridmap.shape[1] - 1
		
		end_idx_x = start_idx_x + span_x
		if end_idx_x > max_idx_x:
			end_idx_x = max_idx_x
			start_idx_x = end_idx_x - span_x
		end_idx_y = start_idx_y + span_y
		if end_idx_y > max_idx_y:
			end_idx_y = max_idx_y
			start_idx_y = end_idx_y - span_y
		
		# Collect debug information
		debug_info["start_x"] = start_idx_x
		debug_info["start_y"] = start_idx_y
		debug_info["end_x"] = end_idx_x
		debug_info["end_y"] = end_idx_y

		#print("Shape in func: ", self.gridmap[start_idx_x:end_idx_x, start_idx_y:end_idx_y].shape)

		submap = self.gridmap[start_idx_x:end_idx_x, start_idx_y:end_idx_y]
		
		"""
		print(self.resolution)
		print(center_idx_x, center_idx_y)
		print(span_x, span_y)

		plt.figure()
		plt.imshow(self.gridmap, cmap="gray_r")
		plt.show()

		plt.figure()
		plt.imshow(submap, cmap="gray_r")
		plt.show()
		"""
		
		return submap, debug_info
		
	def getSubmapByCoords(self, center_pos_x, center_pos_y, size_x, size_y):
		"""
		Get submap around specified coordinates. 
		The sizes in x and y direction are within the same coordinate frame as the center coordinates.
		"""
		center_idx_x, center_idx_y = self.getIdx(center_pos_x, center_pos_y)
		span_x = int(size_x / self.resolution)
		span_y = int(size_y / self.resolution)

		#print("coordinates: ", center_pos_x, center_pos_y)
		
		return self.getSubmapByIndices(center_idx_x, center_idx_y, span_x, span_y)[0]

build/lib/data_utils/OccupancyGrid.py
eliatrevisan/IA-MPPI-LBM
from copyreg import add_extension
import os
import numpy as np
import math
import pylab as pl
import sys
if sys.version_info[0] < 3:
	import Support as sup
else:
	import src.data_utils.Support as sup


def compute_trajectory_prediction_mse(args,ground_truth, predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	avg_mse = 0
	cnt = 0
	mse_list = []
	for pred, gt in zip(predictions, ground_truth):
		avg_mse = 0
		cnt = 0
		# compute average per trajectory
		for t in range(len(pred)):
			# real trajectory global frame
			real_vel_global_frame = gt.vel_vec[t+args.prev_horizon+1:t+args.prediction_horizon+args.prev_horizon+1,:2]
			real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=real_vel_global_frame,
			                                           dt=args.dt)
			vel_pred = np.zeros((args.prediction_horizon, args.output_dim))
			error = 0
			pred_t = pred[t][0]
			if args.n_mixtures<=1:
				# plot predicted trajectory global frame
				for sample_id in range(1):
					for i in range(args.prediction_horizon):
						idx = i * args.output_pred_state_dim
						idy = i * args.output_pred_state_dim + 1
						mu_x = pred_t[sample_id, idx]
						mu_y = pred_t[sample_id, idy]
						vel_pred[i, :] = [mu_x, mu_y]

				pred_vel_global_frame = vel_pred
				traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),pred_vel=np.squeeze(pred_vel_global_frame), dt=args.dt)
				for pred_step in range(args.prediction_horizon):
					error += np.linalg.norm(real_traj_global_frame[pred_step, :] - traj_pred[pred_step, :])/args.prediction_horizon
				avg_mse = (avg_mse*cnt+error)/(cnt+1)
				cnt += 1
			else:
				min_error = np.zeros((pred_t.shape[0]))
				for sample_id in range(pred_t.shape[0]):
					error = np.zeros((args.n_mixtures))
					for mix_idx in range(args.n_mixtures):
						# plot predicted trajectory global frame
						for i in range(args.prediction_horizon):
							idx = i * args.output_pred_state_dim * args.n_mixtures + mix_idx
							idy = i * args.output_pred_state_dim * args.n_mixtures + mix_idx + args.n_mixtures
							if args.normalize_data:
								mu_x = pred_t[sample_id,idx] / args.sx_vel + args.min_vel_x
								mu_y = pred_t[sample_id,idy] / args.sy_vel + args.min_vel_y
							else:
								mu_x = pred_t[sample_id,idx]
								mu_y = pred_t[sample_id,idy]
							vel_pred[i,:] = [mu_x, mu_y]

						traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),pred_vel=vel_pred, dt=args.dt)
						for pred_step in range(args.prediction_horizon):
							error[mix_idx] += np.linalg.norm(real_traj_global_frame[pred_step, :] - traj_pred[pred_step, :])/args.prediction_horizon
					min_error[sample_id] = min(error)
				avg_mse = (avg_mse * cnt + min(min_error)) / (cnt + 1)
				cnt += 1
		mse_list.append(avg_mse)
	return avg_mse, mse_list

def compute_trajectory_fde(args,ground_truth, predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	avg_fde = 0
	cnt = 0
	avg_fde_list = []
	real_traj_global_frame = np.zeros((args.prediction_horizon, args.output_dim))
	for pred, gt in zip(predictions, ground_truth):
		avg_fde = 0
		cnt = 0
		# compute average per trajectory
		for t in range(len(pred)):
			# real trajectory global frame
			real_vel_global_frame = gt.vel_vec[t+args.prev_horizon+1:t+args.prediction_horizon+args.prev_horizon+1,:2]
			real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=real_vel_global_frame,
			                                           dt=args.dt)
			vel_pred = np.zeros((args.prediction_horizon, args.output_dim))
			error = 0
			pred_t = pred[t][0]
			if args.n_mixtures<=1:
				# plot predicted trajectory global frame
				for sample_id in range(1):
					for i in range(args.prediction_horizon):
						idx = i * args.output_pred_state_dim
						idy = i * args.output_pred_state_dim + 1
						mu_x = pred_t[sample_id, idx]
						mu_y = pred_t[sample_id, idy]
						vel_pred[i, :] = [mu_x, mu_y]

				pred_vel_global_frame = vel_pred
				traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),pred_vel=np.squeeze(pred_vel_global_frame), dt=args.dt)
				error = np.linalg.norm(real_traj_global_frame[-1, :] - traj_pred[-1, :])
				avg_fde = (avg_fde*cnt+error)/(cnt+1)
				cnt += 1
			else:
				min_error = np.zeros((pred_t.shape[0]))
				for sample_id in range(pred_t.shape[0]):
					error = np.zeros((args.n_mixtures))
					for mix_idx in range(args.n_mixtures):
						# plot predicted trajectory global frame
						for i in range(args.prediction_horizon):
							idx = i * args.output_pred_state_dim * args.n_mixtures + mix_idx
							idy = i * args.output_pred_state_dim * args.n_mixtures + mix_idx + args.n_mixtures
							if args.normalize_data:
								mu_x = pred_t[sample_id,idx] / args.sx_vel + args.min_vel_x
								mu_y = pred_t[sample_id,idy] / args.sy_vel + args.min_vel_y
							else:
								mu_x = pred_t[sample_id,idx]
								mu_y = pred_t[sample_id,idy]
							vel_pred[i,:] = [mu_x, mu_y]

						traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),pred_vel=vel_pred, dt=args.dt)
						error[mix_idx] = np.linalg.norm(real_traj_global_frame[-1, :] - traj_pred[-1, :])
					min_error[sample_id] = min(error)
				avg_fde = (min(min_error)+ avg_fde*cnt)/(cnt+1)
				cnt +=1
			avg_fde_list.append(avg_fde)
	return avg_fde, avg_fde_list

def compute_rolling_trajectory_fde(args, horizon, ground_truth, predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	avg_fde = 0
	cnt = 0
	avg_fde_list = []
	real_traj_global_frame = np.zeros((args.prediction_horizon, args.output_dim))
	for pred, gt in zip(predictions, ground_truth):
		avg_fde = 0
		cnt = 0
		# compute average per trajectory
		for t in range(len(pred)):
			# real trajectory global frame
			real_vel_global_frame = gt.vel_vec[t+args.prev_horizon+1:t+args.prediction_horizon+args.prev_horizon+1,:2]
			#real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=real_vel_global_frame,
			#                                           dt=args.dt)
			real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=real_vel_global_frame,
			                                           dt=args.dt)
			vel_pred = np.zeros((args.prediction_horizon, args.output_dim))
			error = 0
			pred_t = pred[t][0]
			if args.n_mixtures<=1:
				# plot predicted trajectory global frame
				for sample_id in range(1):
					for i in range(args.prediction_horizon):
						idx = i * args.output_pred_state_dim
						idy = i * args.output_pred_state_dim + 1
						mu_x = pred_t[sample_id, idx]
						mu_y = pred_t[sample_id, idy]
						vel_pred[i, :] = [mu_x, mu_y]

				pred_vel_global_frame = vel_pred
				traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),pred_vel=np.squeeze(pred_vel_global_frame), dt=args.dt)
				error = np.linalg.norm(real_traj_global_frame[horizon, :] - traj_pred[horizon, :])
				avg_fde = (avg_fde*cnt+error)/(cnt+1)
				cnt += 1
			else:
				min_error = np.zeros((pred_t.shape[0]))
				for sample_id in range(pred_t.shape[0]):
					error = np.zeros((args.n_mixtures))
					for mix_idx in range(args.n_mixtures):
						# plot predicted trajectory global frame
						for i in range(args.prediction_horizon):
							idx = i * args.output_pred_state_dim * args.n_mixtures + mix_idx
							idy = i * args.output_pred_state_dim * args.n_mixtures + mix_idx + args.n_mixtures
							if args.normalize_data:
								mu_x = pred_t[sample_id,idx] / args.sx_vel + args.min_vel_x
								mu_y = pred_t[sample_id,idy] / args.sy_vel + args.min_vel_y
							else:
								mu_x = pred_t[sample_id,idx]
								mu_y = pred_t[sample_id,idy]
							vel_pred[i,:] = [mu_x, mu_y]

						traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),pred_vel=vel_pred, dt=args.dt)
						error[mix_idx] = np.linalg.norm(real_traj_global_frame[horizon, :] - traj_pred[horizon, :])
					min_error[sample_id] = min(error)
				avg_fde = (min(min_error)+ avg_fde*cnt)/(cnt+1)
				cnt +=1
			avg_fde_list.append(avg_fde)
			std = np.std(avg_fde_list)
	return avg_fde, avg_fde_list, std

def compute_2_wasserstein(args, predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	avg_wasser = 0
	cnt = 0
	wasser_list = []
	for prediction in predictions: # number of trajectories
		avg_wasser = 0
		cnt = 0
		# compute average per trajectory
		for pred_step in prediction: # trajectory length
			error = 0
			for i in range(args.prediction_horizon):  # prediction length
				if args.n_mixtures == 2:
					idx = i * args.output_pred_state_dim * args.n_mixtures
					idy = i * args.output_pred_state_dim * args.n_mixtures + args.n_mixtures
					mu_x = pred_step[0][ 0, idx]
					mu_y = pred_step[0][ 0, idy]
					mu_x1 = pred_step[0][ 0, idx+1]
					mu_y1 = pred_step[0][ 0, idy+1]
					sigma_x = pred_step[0][0, idy + args.n_mixtures]
					sigma_y = pred_step[0][0, idy + 2*args.n_mixtures]
					sigma_x1 = pred_step[0][0, idy + args.n_mixtures+1]
					sigma_y1 = pred_step[0][0, idy + 2*args.n_mixtures+1]
					wasser1 = np.linalg.norm(np.array([mu_x,mu_y])-np.array([mu_x1,mu_y1])) + sigma_x+sigma_x1+sigma_y+sigma_y1-2*np.sqrt(sigma_x*sigma_x1)-2*np.sqrt(sigma_y*sigma_y1)
					error += wasser1/args.n_mixtures
					avg_wasser = (avg_wasser * cnt + error) / (cnt + 1)
					cnt += 1
				if args.n_mixtures == 3:
					wasser1 = 0
					for mix_idx in range(args.n_mixtures):
						idx = i * args.output_pred_state_dim * args.n_mixtures + mix_idx
						idy = i * args.output_pred_state_dim * args.n_mixtures + args.n_mixtures + mix_idx
						mu_x = pred_step[0][ 0, idx]
						mu_y = pred_step[0][ 0, idy]
						sigma_x = pred_step[0][ 0, idy + args.n_mixtures]
						sigma_y = pred_step[0][ 0, idy + 2*args.n_mixtures]
						if mix_idx+1 >= args.n_mixtures:
							mu_x1 = pred_step[0][ 0, idx - mix_idx]
							mu_y1 = pred_step[0][ 0, idy - mix_idx]
							sigma_x1 = pred_step[0][ 0, idy + args.n_mixtures - mix_idx]
							sigma_y1 = pred_step[0][ 0, idy + 2 * args.n_mixtures - mix_idx]
						else:
							mu_x1 = pred_step[0][ 0, idx+1]
							mu_y1 = pred_step[0][ 0, idy+1]
							sigma_x1 = pred_step[0][ 0, idy + args.n_mixtures+1]
							sigma_y1 = pred_step[0][ 0, idy + 2*args.n_mixtures+1]
						wasser1 += np.linalg.norm(np.array([mu_x,mu_y])-np.array([mu_x1,mu_y1]))+sigma_x+sigma_x1+sigma_y+sigma_y1-2*np.sqrt(sigma_x*sigma_x1)-2*np.sqrt(sigma_y*sigma_y1)
					error += wasser1/args.n_mixtures
					avg_wasser = (avg_wasser * cnt + error) / (cnt + 1)
					cnt += 1
		wasser_list.append(avg_wasser)
	return avg_wasser, wasser_list

def compute_nll(args, ground_truth,predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	cnt = 0
	avg_list = []
	for pred, gt in zip(predictions, ground_truth): # number of trajectories
		avg_nll = 0
		cnt = 0
		# compute average per trajectory
		for t in range(pred.shape[0]): # trajectory length
			nll = 0
			for i in range(args.prediction_seq_length):  # prediction length
				term1 = 0
				for mix_id in range(args.n_mixtures):
					id = i * args.output_state_dim
					idx = i * args.output_pred_state_dim * args.n_mixtures
					idy = i * args.output_pred_state_dim * args.n_mixtures + args.n_mixtures
					mu_x = pred[t, idx+ mix_id]
					mu_y = pred[t, idy+ mix_id]
					sigma_x = max(pred[t, idy + args.n_mixtures+ mix_id],0)
					sigma_y = max(pred[t, idy + 2 * args.n_mixtures+ mix_id],0)
					pi = pred[t, idy + 3 * args.n_mixtures+ mix_id]
					normx = gt[t, id] - mu_x
					normy = gt[t, id + 1] - mu_y
					# Calculate sx*sy
					sxsy = sigma_x*sigma_y
					# Calculate the exponential factor
					z = np.square(normx/sigma_x) + np.square(normy/sigma_y)
					# Numerator
					result = np.exp(-z/2)
					# Normalization constant
					denom = 2 * np.pi * sxsy
					# Final PDF calculation
					result = result / denom
					term1 += pi*result
				nll += -np.log(np.maximum(term1, 1e-5))
			nll = nll / args.prediction_seq_length
			avg_nll = (avg_nll * cnt + nll) / (cnt + 1)
			cnt += 1
		avg_list.append(avg_nll)
	return avg_nll, avg_list

def compute_nll2(args, ground_truth,predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	cnt = 0
	avg_list = []
	for pred, gt in zip(predictions, ground_truth): # number of trajectories
		avg_nll = 0
		cnt = 0
		# compute average per trajectory
		for t in range(pred.shape[0]): # trajectory length
			nll = 0
			pi = pred[t, args.prediction_seq_length*args.n_mixtures*args.output_pred_state_dim:]
			for i in range(args.prediction_seq_length):  # prediction length
				term1 = 0
				for mix_id in range(args.n_mixtures):
					id = i * args.output_state_dim
					idx = i * args.output_pred_state_dim * args.n_mixtures
					idy = i * args.output_pred_state_dim * args.n_mixtures + args.n_mixtures
					mu_x = pred[t, idx+ mix_id]
					mu_y = pred[t, idy+ mix_id]
					sigma_x = pred[t, idy + args.n_mixtures+ mix_id]
					sigma_y = pred[t, idy + 2 * args.n_mixtures+ mix_id]
					w = pred[t, idy + 3 * args.n_mixtures+ mix_id]
					normx = gt[t, id] - mu_x
					normy = gt[t, id + 1] - mu_y
					# Calculate sx*sy
					sxsy = sigma_x*sigma_y
					# Calculate the exponential factor
					z = (np.square(normx/sigma_x) + np.square(normy/sigma_y))/2.0
					# Numerator
					result = math.exp(-z)
					# Normalization constant
					denom = 2 * np.pi * sxsy
					# Final PDF calculation
					result2 = result / denom #pi[mix_id]*
					term1 += result2 / args.n_mixtures
				nll += -np.log(np.maximum(term1, 1e-5))
			nll = nll / args.prediction_seq_length
			avg_nll = (avg_nll * cnt + nll) / (cnt + 1)
			cnt += 1
		avg_list.append(avg_nll)
	return avg_nll, avg_list


def compute_ade(args, trajectories, predictions):
	"""
	Compute Average Displacement Error for the constant velocity predictions
	"""
	ade = []
	
	for idx in range(0,len(trajectories)):
		traj = trajectories[idx].pose_vec
		pred = predictions[idx]

		for step in range(0, len(pred)):
			tr = traj[step+args.prev_horizon+1:step+args.prev_horizon+args.prediction_horizon+1]
			pr = pred[step] 

			sum = 0
			for i in range(0, args.prediction_horizon):
				sum += np.linalg.norm(tr[i, :2] - pr[i, :])/args.prediction_horizon
			mse = sum #/ (i + 1)
			ade.append(mse)
	return np.mean(ade)

def compute_fde(args, trajectories, predictions):
	# Calculate Final Displacement Error for the CV predictions
	fde = []
	for idx in range(0, len(trajectories)):
		traj = trajectories[idx].pose_vec
		pred = predictions[idx]
		for step in range(0, len(pred)):
			tr = traj[step + args.prev_horizon + args.prediction_horizon] # Trajectory sample at horizon
			pr = pred[step][args.prediction_horizon-1] # Prediction at horizon
			error = np.linalg.norm(tr[:2]-pr[:])
			fde.append(error)
	return np.mean(fde)

def compute_rolling_fde(args, trajectories, predictions):
	res = []
	
	for idx in range(0,len(trajectories)):
		traj = trajectories[idx].pose_vec
		pred = predictions[idx]

		for step in range(0, len(pred)):

			fde = []
			for p in range(0, args.prediction_horizon):
				tr = traj[step + args.prev_horizon + p + 1] # Trajectory sample at horizon
				pr = pred[step][p] # Prediction at horizon
				error = np.linalg.norm(tr[:2]-pr[:])
				fde.append(error)
			res.append(fde)
	return np.mean(res, axis=0), np.std(res, axis=0)



build/lib/data_utils/Performance.py
eliatrevisan/IA-MPPI-LBM
import os
import numpy as np
import cv2
import pylab as pl
from matplotlib.animation import FFMpegWriter
import sys
if sys.version_info[0] < 3:
	import Support as sup
else:
	import src.data_utils.Support as sup
import math
from matplotlib.patches import Ellipse
from scipy import interpolate
import matplotlib.pyplot as plt

def rgba2rgb(rgba):
    # rgba is a list of 4 color elements btwn [0.0, 1.0]
    # or a 2d np array (num_colors, 4)
    # returns a list of rgb values between [0.0, 1.0] accounting for alpha and background color [1, 1, 1] == WHITE
    if isinstance(rgba, list):
        alpha = rgba[3]
        r = max(min((1 - alpha) * 1.0 + alpha * rgba[0],1.0),0.0)
        g = max(min((1 - alpha) * 1.0 + alpha * rgba[1],1.0),0.0)
        b = max(min((1 - alpha) * 1.0 + alpha * rgba[2],1.0),0.0)
        return [r,g,b]
    elif rgba.ndim == 2:
        alphas = rgba[:,3]
        r = np.clip((1 - alphas) * 1.0 + alphas * rgba[:,0], 0, 1)
        g = np.clip((1 - alphas) * 1.0 + alphas * rgba[:,1], 0, 1)
        b = np.clip((1 - alphas) * 1.0 + alphas * rgba[:,2], 0, 1)
        return np.vstack([r,g,b]).T

class Recorder():

	def __init__(self,args,grid):

		# Parameters
		self.args = args
		self.gridmap = grid
		# Animation
		self.fig_animate = pl.figure('Animation')
		self.fig_width = 12  # width in inches
		self.fig_height = 12  # height in inches
		self.fig_size = [self.fig_width, self.fig_height]
		self.fontsize = 9
		self.colors = ['r', 'g', 'b','y','m','c','k']

		self.params = {'backend': 'ps',
		          'axes.labelsize': self.fontsize,
		          'font.size': self.fontsize,
		          'xtick.labelsize': self.fontsize,
		          'ytick.labelsize': self.fontsize,
		          'figure.figsize': self.fig_size}
		pl.rcParams.update(self.params)
		self.ax_pos = pl.subplot()
		#self.ax_pos_local = pl.subplot2grid((6, 12), (0, 6), colspan=6, rowspan=6)
		#self.ax_grid = pl.subplot2grid((9, 12), (7, 0), colspan=2, rowspan=2)
		#self.ax_ped_grid = pl.subplot2grid((9, 12), (7, 2), colspan=2, rowspan=2)
		#self.ax_vel = pl.subplot2grid((9, 12), (7, 4), colspan=2, rowspan=2)
		pl.rcParams.update(self.params)
		#pl.rcParams['animation.ffmpeg_path'] = "/home/bbrito/ffmpeg/ffmpeg"

		metadata = dict(title='Movie Test', artist='Matplotlib')
		self.writer = FFMpegWriter(fps=10, metadata=metadata,codec='mpeg4')

	def animate_local(self,input_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,trajectories,test_args):

		fig_animate = pl.figure('Local Scenario')
		fig_width = 12  # width in inches
		fig_height = 10  # height in inches
		fig_size = [fig_width, fig_height]
		fontsize = 26
		colors = ['r', 'g', 'y']
		params = {'backend': 'ps',
		               'axes.labelsize': fontsize,
		               'font.size': fontsize,
		               'xtick.labelsize': fontsize,
		               'ytick.labelsize': fontsize,
		               'figure.figsize': fig_size}

		pl.rcParams.update(params)
		ax_pos = pl.subplot()

		for animation_idx in range(len(input_list)):
			input = input_list[animation_idx]
			grid = grid_list[animation_idx]

			traj = trajectories[animation_idx]
			gt_vel = y_ground_truth_list[animation_idx]

			if not (y_pred_list_global is None):
				model_vel_pred = y_pred_list_global[animation_idx]
			else:
				model_vel_pred = None
			other_agents_pos = other_agents_list[animation_idx]

			for step in range(input.shape[0]): # trajectory length considering data from getTrajectoryAsBatch

				"""
				fig, ax = plt.subplots(1,2)
				sup.plot_grid_roboat(ax[0], self.gridmap.center, self.gridmap.gridmap, self.gridmap.resolution, self.gridmap.map_size)
				print("xy: ",input[step][0],input[step][1])
				ax[0].plot(input[step][0],input[step][1], markersize=50)
				ax[1].imshow(grid[step], cmap="gray_r")
				plt.show()

				"""
				self.plot_local_scenario(ax_pos,input, grid, model_vel_pred[step], gt_vel, other_agents_pos, step,
					                    rotate=False,n_samples=test_args.n_samples)

				ax_pos.axis("on")
				ax_pos.set_xlabel('x [m]',fontsize=26)
				ax_pos.set_ylabel('y [m]',fontsize=26)
				ax_pos.set_xlim([-10.0, 10.0])
				ax_pos.set_ylim([-10.0, 10.0])
				ax_pos.set_aspect('equal')
				pl.rcParams.update(params)
				ax_pos.set_aspect('equal')

				if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario+"/figs/"):
					os.makedirs(self.args.model_path + '/results/' + self.args.scenario+"/figs/")
				if test_args.save_figs:
					fig_animate.savefig(self.args.model_path + '/results/' + self.args.scenario+"/figs/local_"+str(animation_idx)+"_"+str(step)+".png")


	def animate(self,input_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,trajectories,test_args):

		#pl.show(block=False)
		if not os.path.exists(self.args.model_path + '/../videos/'):
			os.makedirs(self.args.model_path + '/../videos/')
		if test_args.unit_testing:
			video_file = self.args.model_path + '/../videos/' + str(self.args.exp_num) + "_unit_tests.mp4"
		elif test_args.freeze_other_agents:
			video_file = self.args.model_path + '/../videos/' + str(self.args.exp_num) + "_frozen.mp4"
		else:
			video_file = self.args.model_path + '/../videos/' + str(self.args.exp_num) + "_final.mp4"
		print("Recoding to: " + video_file)
		with self.writer.saving(self.fig_animate, video_file, 100):
			for animation_idx in range(len(input_list)):
				input = input_list[animation_idx]
				grid = grid_list[animation_idx]
				ped_grid = ped_grid_list[animation_idx]
				traj = trajectories[animation_idx]
				gt_vel = y_ground_truth_list[animation_idx]

				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				other_agents_pos = other_agents_list[animation_idx]

				for step in range(input.shape[0]): # trajectory length considering data from getTrajectoryAsBatch
					# Eucldean / position space
					self.ax_pos.clear()
					#self.ax_pos.plot(input[:, 0], input[:, 1], color='c', alpha=0.4, lw=1)
					self.ax_pos.plot(traj.pose_vec[-1, 0], traj.pose_vec[-1, 1], color='g', marker='o')
					#Real Trajectory
					self.ax_pos.plot(traj.pose_vec[:, 0], traj.pose_vec[:, 1], color='b', lw=2)

					self.ax_pos.plot(input[step, 0], input[step, 1], color='c', marker='o')
					bottom_left_submap = input[step, 0:2] - np.array([self.args.submap_size * self.gridmap.resolution / 2.0, self.args.submap_size * self.gridmap.resolution / 2.0])
					self.ax_pos.add_patch(pl.Rectangle(bottom_left_submap, self.args.submap_width * self.gridmap.resolution, self.args.submap_size * self.gridmap.resolution, fill=None))
					self.plot_local_scenario(input, grid, model_vel_pred[step], gt_vel, other_agents_pos, step,
					                    rotate=False,n_samples=test_args.n_samples)

					# Other agents
					for id in range(other_agents_pos[step].shape[0]): # number of agents
						self.ax_pos.plot(other_agents_pos[step][id, 0], other_agents_pos[step][id, 1], marker='o', color='b')
					vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
					sigmax = np.zeros((self.args.prediction_horizon, 1))
					sigmay = np.zeros((self.args.prediction_horizon, 1))
					#Predicted trajectory
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							# plot predicted trajectory global frame
							prediction_sample = model_vel_pred[step][0]
							for i in range(self.args.prediction_horizon):
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								mu_x = prediction_sample[0, idx]
								mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]

							pred_vel_global_frame = vel_pred
							traj_pred = sup.path_from_vel(initial_pos=np.array(
								[input[step, 0],
								 input[step, 1]]),
								pred_vel=pred_vel_global_frame, dt=self.args.dt)
							self.ax_pos.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[0], label='Predicted trajectory')
						else:
							for sample_id in range(test_args.n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								for mix_idx in range(self.args.n_mixtures):
									# plot predicted trajectory global frame
									for pred_step in range(self.args.prediction_horizon):
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										mu_x = prediction_sample[0, idx]
										mu_y = prediction_sample[0, idx + self.args.n_mixtures]
										if self.args.output_pred_state_dim>2:
											sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
											sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										vel_pred[pred_step, :] = [mu_x, mu_y]

									pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array(
										[input[step, 0],
										 input[step, 1]]),
										pred_vel=pred_vel_global_frame, dt=self.args.dt)

									self.ax_pos.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')

									if self.args.output_pred_state_dim > 2:
										# prior of 0 on the uncertainty of the pedestrian velocity
										sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
										sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

										e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
										             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
										e1.set_alpha(0.5)
										self.ax_pos.add_patch(e1)
										#self.ax_pos.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')
										for pred_step in range(1, self.args.prediction_horizon):
											sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
											sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt
											e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
											             height=np.sqrt(sigma_y) / 2,
											             angle=0 / np.pi * 180)
											self.ax_pos.add_patch(e1)

					# Real trajectory
					model_vel_real = np.zeros((self.args.prediction_horizon, self.args.output_dim))
					for i in range(self.args.prediction_horizon):
						idx = i * self.args.output_dim
						idy = i * self.args.output_dim + 1
						mu_x = gt_vel[step, idx]
						mu_y = gt_vel[step, idy]
						model_vel_real[i, :] = [mu_x, mu_y]
					if False:
						real_vel = sup.rotate_predicted_vel_to_global_frame(model_vel_real, batch_x[step, 2:])
					else:
						real_vel = model_vel_real
					traj_real = sup.path_from_vel(initial_pos=np.array([0, 0]),
					                              pred_vel=real_vel, dt=self.args.dt)
					self.ax_pos_local.plot(traj_real[:, 0], traj_real[:, 1], color='m', label='Agent real trajectory')

					sup.plot_grid(self.ax_pos, np.array([0.0, 0.0]), self.gridmap.gridmap, self.gridmap.resolution,
					              self.gridmap.map_size)

					ax_pos.axis('equal')
					self.ax_pos.axis("on")
					self.ax_pos.set_xlabel('x [m]')
					self.ax_pos.set_ylabel('y [m]')
					self.ax_pos.set_aspect('equal')

					#self.fig_animate.canvas.draw()
					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario+"/figs/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario+"/figs/")
					if test_args.save_figs:
						self.fig_animate.savefig(self.args.model_path + '/results/' + self.args.scenario+"/figs/result_"+str(animation_idx)+"_"+str(step)+".jpg")
					self.writer.grab_frame()
				

	def animate_group_of_agents(self,trajectories,y_pred_list_global,test_args=None):

		if not os.path.exists(self.args.model_path + '/videos/'):
			os.makedirs(self.args.model_path + '/videos/')

		fig_animate = pl.figure('Global Scenario')
		fig_width = 12  # width in inches
		fig_height = 10  # height in inches
		fig_size = [fig_width, fig_height]
		fontsize = 26

		params = {'backend': 'ps',
		          'axes.labelsize': fontsize,
		          'font.size': fontsize,
		          'xtick.labelsize': fontsize,
		          'ytick.labelsize': fontsize,
		          'figure.figsize': fig_size}

		plt_colors = []
		plt_colors.append([0.8500, 0.3250, 0.0980])  # orange
		plt_colors.append([0.0, 0.4470, 0.7410])  # blue
		plt_colors.append([0.4660, 0.6740, 0.1880])  # green
		plt_colors.append([0.4940, 0.1840, 0.5560])  # purple
		plt_colors.append([0.9290, 0.6940, 0.1250])  # yellow
		plt_colors.append([0.3010, 0.7450, 0.9330])  # cyan
		plt_colors.append([0.6350, 0.0780, 0.1840])  # chocolate
		plt_colors.append([0.6350, 0.0780, 0.1840])  # chocolate
		plt_colors.append([0.6350, 0.0780, 0.1840])  # chocolate
		scenario = self.args.scenario.split('/')[0]
		ax_pos = pl.subplot()

		ax_pos.set_xlabel('x [m]')
		ax_pos.set_ylabel('y [m]')
		ax_pos.axis('on')
		ax_pos.set_aspect('equal')

		n_agents =  len(trajectories[0])
		axbackground = fig_animate.canvas.copy_from_bbox(ax_pos.bbox)

		# Create plotting variables
		current_pos_list_ = []
		pred_line_list_ = []
		real_line_list_ = []
		ellipses_list_ = []
		for ag_id in range(n_agents):
			current_pos_, = ax_pos.plot([], color=plt_colors[ag_id], marker='o')
			current_pos_list_.append(current_pos_)
			pred_line, = ax_pos.plot([], color=plt_colors[ag_id], label='Predicted trajectory Agent ' + str(ag_id))
			pred_line_list_.append(pred_line)
			real_line, = ax_pos.plot([], color=plt_colors[ag_id], alpha=0.5, label='Real trajectory Agent ' + str(ag_id))
			real_line_list_.append(real_line)
			c = rgba2rgb(plt_colors[ag_id] + [float(1)])
			ellipses = [Ellipse(xy=(0, 0), width=0.6,
			                    height=0.6,
			                    fc=c, ec=plt_colors[ag_id],
			                    fill=True) for pred_step in range(self.args.prediction_horizon)]
			for ellipse in ellipses:
				ax_pos.add_artist(ellipse)
			ellipses_list_.append(ellipses)

		ax_pos.legend()

		with self.writer.saving(fig_animate, self.args.model_path + '/videos/' +str(self.args.exp_num)+scenario+".mp4", 100):
			for animation_idx in range(len(trajectories)):
				traj = trajectories[animation_idx]

				model_vel_pred = y_pred_list_global[animation_idx]

				vel_pred = np.zeros((self.args.prediction_horizon, self.args.output_pred_state_dim))
				sigmax = np.zeros((self.args.prediction_horizon, 1))
				sigmay = np.zeros((self.args.prediction_horizon, 1))

				ax_pos.set_xlim([0.0, 20.0])
				ax_pos.set_ylim([-5.0, 5.0])

				for step in range(traj[0].time_vec.shape[0]):
					for ag_id in range(n_agents):
					  # trajectory length considering data from getTrajectoryAsBatch
						fig_animate.canvas.restore_region(axbackground)

						try:
							current_pos_list_[ag_id].set_data(traj[ag_id].pose_vec[step,0]/self.args.sx_pos,traj[ag_id].pose_vec[step, 1]/self.args.sy_pos)
						except:
							current_pos_list_[ag_id].set_data(traj[ag_id].pose_vec[-1, 0] / self.args.sx_pos,
							                                  traj[ag_id].pose_vec[-1, 1] / self.args.sy_pos)
						#Predicted trajectory
						for sample_id in range(1):
							try:
								prediction_sample = model_vel_pred[ag_id][step][sample_id]
								init_pose = traj[ag_id].pose_vec[step,:2]
							except:
								prediction_sample =  np.zeros([self.args.output_dim*self.args.prediction_horizon])
								init_pose = traj[ag_id].pose_vec[-1, :2]
							for mix_idx in range(self.args.n_mixtures):
								# plot predicted trajectory global frame
								for pred_step in range(self.args.prediction_horizon):
									idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
									mu_x = prediction_sample[idx]/self.args.sx_vel
									mu_y = prediction_sample[idx + self.args.n_mixtures]/self.args.sy_vel
									if self.args.output_pred_state_dim > 2:
										sigmax[pred_step, :] = prediction_sample[idx + 2 * self.args.n_mixtures]
										sigmay[pred_step, :] = prediction_sample[idx + 3 * self.args.n_mixtures]
									vel_pred[pred_step, :] = [mu_x, mu_y]

								pred_vel_global_frame = vel_pred
								traj_pred = sup.path_from_vel(initial_pos=init_pose/self.args.sx_pos,
								                              pred_vel=pred_vel_global_frame, dt=self.args.dt)

								pred_line_list_[ag_id].set_data(traj_pred[:, 0], traj_pred[:, 1])

								if self.args.output_pred_state_dim > 2:
									# TODO: ADAPT TO NEW CODE
									# prior of 0 on the uncertainty of the pedestrian velocity
									sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
									sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

									e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
									             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
									e1.set_alpha(0.5)

									for pred_step in range(1, self.args.prediction_horizon):
										sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
										sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt

										e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
										             height=np.sqrt(sigma_y) / 2,
										             angle=0 / np.pi * 180)
										ax_pos.add_patch(e1)
								else:
									for pred_step in range(0, self.args.prediction_horizon):
										alpha = 1/ (1+pred_step)
										c = rgba2rgb(plt_colors[ag_id] + [float(alpha)])
										ellipses_list_[ag_id][pred_step].set_facecolor(c)
										ellipses_list_[ag_id][pred_step].set_center((traj_pred[pred_step, 0], traj_pred[pred_step, 1]))

						# Real trajectory
						real_line_list_[ag_id].set_data(traj[ag_id].pose_vec[:,0]/self.args.sx_pos, traj[ag_id].pose_vec[:,1]/self.args.sy_pos)

					fig_animate.canvas.blit(ax_pos.bbox)
					self.writer.grab_frame()

					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario + "/figs/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario + "/figs/")
					if test_args.save_figs:
						fig_animate.savefig(self.args.model_path + '/results/' + self.args.scenario + "/figs/result_" + str(
							animation_idx) + "_" + str(step) + ".jpg")

					fig_animate.canvas.flush_events()

	def animate_global(self,input_list,grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,y_ground_truth_list_exp,all_traj_likelihood, cv_pred_list, test_args=None ):

		if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario):
			os.makedirs(self.args.model_path + '/results/' + self.args.scenario)

		fig_animate = pl.figure('Global Scenario')
		fig_animate.tight_layout()
		fig_width = 12  # width in inches
		fig_height = 10  # height in inches
		fig_size = [fig_width, fig_height]
		fontsize = 26

		params = {'backend': 'ps',
		          'axes.labelsize': fontsize,
		          'font.size': fontsize,
		          'xtick.labelsize': fontsize,
		          'ytick.labelsize': fontsize,
		          'figure.figsize': fig_size,
		          'legend.loc': 'upper right',
		          'legend.fontsize': 14}

		#pl.rcParams.update(params)

		plt_colors = []
		plt_colors.append([0.8500, 0.3250, 0.0980])  # orange
		plt_colors.append([0.0, 0.4470, 0.7410])  # blue
		scenario = self.args.scenario.split('/')[0]
		ax_pos = pl.subplot()

		ax_pos.set_xlabel('x [m]')
		ax_pos.set_ylabel('y [m]')
		ax_pos.axis('on')
		#ax_pos.set_aspect('equal')

		# Load Map
		"""
		map_file = os.path.join(self.args.data_path + self.args.scenario, 'map.png')
		if os.path.exists(map_file):
			#img = np.uint8(cv2.imread(map_file) * -1 + 255)
			img = cv2.imread(map_file)
			#img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
			resolution = 0.1
			width = int(img.shape[1] * resolution)
			height = int(img.shape[0] * resolution)
			dim = (width, height)
			resized = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
		else:
			print('[INF] No map file')
		ax_pos.imshow(resized,extent=[0, 30, -3, 2])
		"""
		#sup.plot_grid(ax_pos, np.array([0.0, 0.0]), self.gridmap.gridmap, self.gridmap.resolution,  # np.array([0.0, 0.0])
		#              self.gridmap.map_size)

		sup.plot_grid_roboat(ax_pos, self.gridmap.center, self.gridmap.gridmap, self.gridmap.resolution,  # np.array([0.0, 0.0])
		              self.gridmap.map_size)

		axbackground = fig_animate.canvas.copy_from_bbox(ax_pos.bbox)
		current_pos_, = ax_pos.plot([], color='b', marker='o')
		other_agent_traj, = ax_pos.plot([], color='r', marker='.', alpha=0.5, label='Real Trajectory Agent 2')
		other_agent_traj_, = ax_pos.plot([], color='r', marker='.', alpha=0.5, label='Real Trajectory Agent 3')
		other_agent_traj_2, = ax_pos.plot([], color='r', marker='.', alpha=0.5, label='Real Trajectory Agent 4')
		cv_predictions, = ax_pos.plot([], color='g', marker='.', alpha=0.5, label='Constant Velocity')
		other_agents_line_ = []
		pred_line_list = []
		for i in range(self.args.n_mixtures):
			pred_line, = ax_pos.plot([], color=self.colors[2], label='Predicted trajectory Agent ' + str(i))
			pred_line_list.append(pred_line)

		real_line, = ax_pos.plot([], color='c',marker='.',alpha=0.5, label='Real trajectory Agent 1')
		c = rgba2rgb(plt_colors[1] + [float(1)])

		ellipses = []
		for i in range(self.args.n_mixtures):
			ellipse = [Ellipse(xy=(0,0), width=0.6,
				             height=0.6,
				             fc=c, ec=plt_colors[1],
				             fill=True) for pred_step in range(self.args.prediction_horizon)]
			ellipses.append(ellipse)

		for ellipse in ellipses:
			for el in ellipse:
				ax_pos.add_artist(el)

		#ax_pos.legend()


		with self.writer.saving(fig_animate, self.args.model_path + '/results/' + self.args.scenario+"/video.mp4", 100):
			for animation_idx in range(0,len(input_list)):
				input = input_list[animation_idx] #* (1 / self.gridmap.resolution)
				grid = grid_list[animation_idx]
				if test_args.constant_velocity:
					cv_pred = cv_pred_list[animation_idx]

				#ax_pos.set_xlim([int(np.min(input/self.args.sx_pos))-2.0, int(np.max(input/self.args.sx_pos))+2.0])
				#ax_pos.set_ylim([int(np.min(input/self.args.sy_pos))-2.0, int(np.max(input/self.args.sy_pos))+2.0])
				#ax_pos.set_xlim([-self.gridmap.map_size[0]/2, self.gridmap.map_size[0]/2])
				#ax_pos.set_ylim([-self.gridmap.map_size[1]/2, self.gridmap.map_size[1]/2])
				#ax_pos.set_xlim([self.gridmap.map_size[0], -self.gridmap.map_size[0]])
				#ax_pos.set_ylim([self.gridmap.map_size[1], -self.gridmap.map_size[1]])

				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				other_agents_pos = other_agents_list[animation_idx]
				traj_likelihood = all_traj_likelihood[animation_idx]

				number_of_agents = len(other_agents_pos[0])
				for ag_id in range(number_of_agents):
					other_agent, = ax_pos.plot([], color='r', marker='o')
					other_agents_line_.append(other_agent)
				vel_pred = np.zeros((self.args.prediction_horizon, 2))
				sigmax = np.zeros((self.args.prediction_horizon))
				sigmay = np.zeros((self.args.prediction_horizon))

				### 
				# Plotting agent real trajectories

				other_agent_traj.set_data(None,None)
				other_agent_traj_.set_data(None,None)
				other_agent_traj_2.set_data(None,None)
			
				x_real = list()
				y_real = list()
				for test in range(input.shape[0]):
					x_real.append(other_agents_pos[test][:,0])
					y_real.append(other_agents_pos[test][:,1])

				if len(x_real[0]) == 1:
					other_agent_traj.set_data(x_real[:], y_real[:])
				elif len(x_real[0]) == 2:
					# Two (other) agents
					n_agents = len(x_real[0])
 
					x_real_one = list()
					x_real_two = list()
					y_real_one = list()
					y_real_two = list()

					for test in range(input.shape[0]):
						x_real_one.append(x_real[test][0])
						x_real_two.append(x_real[test][1])
						y_real_one.append(y_real[test][0])	
						y_real_two.append(y_real[test][1])

					other_agent_traj.set_data(x_real_one[:], y_real_one[:])
					other_agent_traj_.set_data(x_real_two[:], y_real_two[:])
				elif len(x_real[0]) == 3:
					# Three (other) agents
					n_agents = len(x_real[0])
 
					x_real_one = list()
					x_real_two = list()
					x_real_three = list()
					y_real_one = list()
					y_real_two = list()
					y_real_three = list()

					for test in range(input.shape[0]):
						x_real_one.append(x_real[test][0])
						x_real_two.append(x_real[test][1])
						x_real_three.append(x_real[test][2])
						y_real_one.append(y_real[test][0])	
						y_real_two.append(y_real[test][1])
						y_real_three.append(y_real[test][2])

					other_agent_traj.set_data(x_real_one[:], y_real_one[:])
					other_agent_traj_.set_data(x_real_two[:], y_real_two[:])
					other_agent_traj_2.set_data(x_real_three[:], y_real_three[:])
	

				for step in range(input.shape[0]): # trajectory length considering data from getTrajectoryAsBatch
					fig_animate.canvas.restore_region(axbackground)
					# Eucldean / position space
					current_pos_.set_data(input[step, 0]/self.args.sx_pos, input[step, 1]/self.args.sy_pos)

					if test_args.constant_velocity:
						cv_predictions.set_data(cv_pred[step][:,0], cv_pred[step][:,1])
					
					bottom_left_submap = input[step, 0:2] - np.array([self.args.submap_size / 2.0, self.args.submap_size / 2.0])

					# Other agents
					for id in range(other_agents_pos[step].shape[0]): # number of agents
						other_agents_line_[id].set_data(other_agents_pos[step][id, 0], other_agents_pos[step][id, 1])
					#Predicted trajectory
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							# plot predicted trajectory global frame
							for i in range(self.args.prediction_horizon):
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								mu_x = model_vel_pred[step, idx]/self.args.sx_vel
								mu_y = model_vel_pred[step, idy]/self.args.sy_vel
								vel_pred[0, 2 * i:2 * i + 2] = [mu_x, mu_y]

							pred_vel_global_frame = vel_pred
							traj_pred = sup.path_from_vel(initial_pos=np.array(
								[input[step, 0],
								 input[step, 1]]),
								pred_vel=np.squeeze(pred_vel_global_frame), dt=self.args.dt)
							pred_line_list[mix_idx].set_data(traj_pred[:, 0], traj_pred[:, 1])
						else:
							for sample_id in range(1):
								prediction_sample = model_vel_pred[step][sample_id]
								most_likely_traj = np.argmax(traj_likelihood[step][0,0])
								for mix_idx in range(self.args.n_mixtures):
									# plot predicted trajectory global frame
									for pred_step in range(self.args.prediction_horizon):
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										mu_x = prediction_sample[0, idx]/self.args.sx_vel
										mu_y = prediction_sample[0, idx + self.args.n_mixtures]/self.args.sy_vel
										if self.args.output_pred_state_dim > 2:
											sigmax[pred_step] = np.maximum(np.minimum(prediction_sample[0, idx + 2 * self.args.n_mixtures],0.5),0.05)
											sigmay[pred_step] = np.maximum(np.minimum(prediction_sample[0, idx + 3 * self.args.n_mixtures],0.5),0.05)
										vel_pred[pred_step, :] = [mu_x, mu_y]

									pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=input[step,:2]/self.args.sx_pos,
									                              pred_vel=pred_vel_global_frame, dt=self.args.dt)
									# hack for prediciting positions
									#traj_pred = vel_pred+input[step,:2]
									if mix_idx != most_likely_traj:
										pred_line_list[mix_idx].set_data(traj_pred[:, 0], traj_pred[:, 1])
										pred_line_list[mix_idx].set_color(self.colors[2])
									else:
										pred_line_list[mix_idx].set_data(traj_pred[:, 0], traj_pred[:, 1])
										pred_line_list[mix_idx].set_color(self.colors[1])
									pred_line_list[mix_idx].set_label('Predicted trajectory ' + str(traj_likelihood[step][0,0,mix_idx]))
									
									if self.args.output_pred_state_dim > 2:
										# prior of 0 on the uncertainty of the pedestrian velocity
										sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
										sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

										alpha = 1
										c = rgba2rgb(plt_colors[0] + [float(alpha)])
										ellipses[mix_idx][0].set_facecolor(c)
										ellipses[mix_idx][0].set_center((traj_pred[0, 0], traj_pred[0, 1]))
										ellipses[mix_idx][0].width = np.sqrt(sigma_x)
										ellipses[mix_idx][0].height = np.sqrt(sigma_y)
										for pred_step in range(1, self.args.prediction_horizon):
											sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
											sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt
											#sigma_x += sigmax[pred_step]
											#sigma_y += sigmay[pred_step]

											alpha = 1 / (1 + pred_step)
											c = rgba2rgb(plt_colors[0] + [float(alpha)])
											ellipses[mix_idx][pred_step].set_facecolor(c)
											ellipses[mix_idx][pred_step].set_center((traj_pred[pred_step, 0], traj_pred[pred_step, 1]))
											ellipses[mix_idx][pred_step].width = np.sqrt(sigma_x)
											ellipses[mix_idx][pred_step].height = np.sqrt(sigma_y)
									else:
										for pred_step in range(0, self.args.prediction_horizon):
											alpha = 1/ (1+pred_step)
											c = rgba2rgb(plt_colors[1] + [float(alpha)])
											ellipses[mix_idx][pred_step].set_facecolor(c)
											ellipses[mix_idx][pred_step].set_center((traj_pred[pred_step, 0], traj_pred[pred_step, 1]))
									
									"""
									# From animate local script
									if self.args.output_pred_state_dim > 2:
										# prior of 0 on the uncertainty of the pedestrian velocity
										sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
										sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

										e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
																height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
										e1.set_alpha(0.5)

										for pred_step in range(1, self.args.prediction_horizon):
											sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
											sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt

											e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
																	height=np.sqrt(sigma_y) / 2,
																	angle=0 / np.pi * 180)
											ax_pos.add_patch(e1)
										"""

					# Real trajectory
					vel_real = np.zeros((self.args.prediction_horizon, self.args.output_dim))

					# plot trajectory global frame
					real_line.set_data(input[:, 0]/self.args.sx_pos, input[:, 1]/self.args.sy_pos)

					#sup.plot_grid(ax_pos, np.array([0.0, 0.0]), self.gridmap.gridmap, self.gridmap.resolution,
					#              self.gridmap.map_size)
					fig_animate.set_tight_layout(True)
					fig_animate.canvas.blit(ax_pos.bbox)

					self.writer.grab_frame()

					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario + "/figs/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario + "/figs/")
					if test_args.save_figs:
						fig_animate.savefig(self.args.model_path + '/results/' + self.args.scenario + "/figs/result_" + str(
							animation_idx) + "_" + str(step) + ".png")

					fig_animate.canvas.flush_events()


	def plot_local_scenario(self,ax_pos_local, batch_x, batch_grid, model_vel_pred, model_vel_real, other_agents_pos,step,
	            rotate=False,n_samples=1):

		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth trajectory in global frame
				model_vel_pred: agent predicted trajectory in global frame
			"""
		ax_pos_local.clear()


		#plot initial and goal pose
		x_init_global_frame = batch_x[step, 0]
		y_init_global_frame = batch_x[step, 1]

		# plot initial and goal pose
		ax_pos_local.plot(0, 0, color='c', marker='o', label='Agent initial pos')
		

		# Local Frame Goal Position
		x_goal = batch_x[-1, 0] - x_init_global_frame
		y_goal = batch_x[-1, 1] - y_init_global_frame
		#self.ax_pos_local.plot(x_goal, y_goal, color='g', marker='o', label='Agent goal pos')

		# plot predicted trajectory local frame
		if not (model_vel_pred is None):
			vel_pred = np.zeros((self.args.prediction_horizon, self.args.output_dim))
			sigmax = np.zeros((self.args.prediction_horizon, 1))
			sigmay = np.zeros((self.args.prediction_horizon, 1))
			if self.args.n_mixtures == 0:
				# plot predicted trajectory global frame
				for i in range(self.args.prediction_horizon):
					idx = i * self.args.output_pred_state_dim
					idy = i * self.args.output_pred_state_dim + 1
					mu_x = model_vel_pred[0][0, idx]
					mu_y = model_vel_pred[0][0, idy]
					vel_pred[i, :] = [mu_x, mu_y]
				if rotate:
					pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, batch_x[step, 2:])
				else:
					pred_vel_global_frame = vel_pred
				traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),
					pred_vel=pred_vel_global_frame, dt=self.args.dt)
				ax_pos_local.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[0], label='Predicted trajectory')
			else:
				for sample_id in range(n_samples):
					prediction_sample = model_vel_pred[sample_id]
					for mix_idx in range(self.args.n_mixtures):
						# plot predicted trajectory global frame
						for pred_step in range(self.args.prediction_horizon):
							idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
							mu_x = prediction_sample[0, idx]
							mu_y = prediction_sample[0, idx + self.args.n_mixtures]
							if self.args.output_pred_state_dim > 2:
								sigmax[pred_step, :] = prediction_sample[0, idx + 2 * self.args.n_mixtures]
								sigmay[pred_step, :] = prediction_sample[0, idx + 3 * self.args.n_mixtures]
							vel_pred[pred_step, :] = [mu_x, mu_y]

						if rotate:
							pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, batch_x[step, 2:])
						else:
							pred_vel_global_frame = vel_pred
						traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),
							pred_vel=pred_vel_global_frame, dt=self.args.dt)

						ax_pos_local.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx],
							                       label='Predicted trajectory')

						if self.args.output_pred_state_dim > 2:
							# prior of 0 on the uncertainty of the pedestrian velocity
							sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
							sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

							e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
							             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
							e1.set_alpha(0.5)

							for pred_step in range(1, self.args.prediction_horizon):
								sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
								sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt

								e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
								             height=np.sqrt(sigma_y) / 2,
								             angle=0 / np.pi * 180)
								ax_pos_local.add_patch(e1)

		# Plot real trajectory
		gt_vel = np.zeros((self.args.prediction_horizon, self.args.output_dim))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = model_vel_real[step, idx]
			mu_y = model_vel_real[step, idy]
			gt_vel[i, :] = [mu_x, mu_y]
		if False:
			real_vel = sup.rotate_predicted_vel_to_global_frame(gt_vel, batch_x[step, 2:])
		else:
			real_vel = gt_vel
		traj_real = sup.path_from_vel(initial_pos=np.array([0, 0]),
		                              pred_vel=real_vel, dt=self.args.dt)
		ax_pos_local.plot(traj_real[:, 0], traj_real[:, 1], color='m', label='Agent real trajectory')

		# Other agents
		for id in range(other_agents_pos[step].shape[0]):  # number of agents
			ax_pos_local.plot(other_agents_pos[step][id, 0]-x_init_global_frame
			                 , other_agents_pos[step][id, 1]-y_init_global_frame, marker='o', color='b')
		
		# Local grid
		
		if rotate:
			heading = math.atan2(batch_x[step, 3], batch_x[step, 2])
			grid = sup.rotate_grid_around_center(batch_grid[step, :, :], -heading * 180 / math.pi)  # rotation in degrees
			#sup.plot_grid(ax_pos_local, np.array([0.0, 0.0]), grid, self.gridmap.resolution,
			#              np.array([self.args.submap_width, self.args.submap_height]))
		else:
			"""
			if real_vel[0,0] > 0.1:
				sup.plot_grid(ax_pos_local, np.array([3.0, 0.0]), batch_grid[step, :, :], self.gridmap.resolution,
				              np.array([self.args.submap_width, self.args.submap_height]))
			elif real_vel[0,0] < -0.1:
				sup.plot_grid(ax_pos_local, np.array([-3.0, 0.0]), batch_grid[step, :, :], self.gridmap.resolution,
				              np.array([self.args.submap_width, self.args.submap_height]))
			else:
			"""
			#batch_grid[step, :, :] = list(zip(*batch_grid[step, :, :][::-1]))

			#fig, ax = plt.subplots()
			#ax.imshow(batch_grid[step, :, :])
			#plt.show()
			#ax_pos_local.imshow(batch_grid[step,:,:], cmap="gray_r")
			
			#sup.plot_submap(ax_pos_local, self.gridmap.center, batch_grid[step, :, :], self.gridmap.resolution)

			sup.plot_grid(ax_pos_local, self.gridmap.center, batch_grid[step, :, :], self.gridmap.resolution,
		            np.array([self.args.submap_size, self.args.submap_size]))



	def save_global_scenario(self,input_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,rotate=False,n_samples=1,format=".pdf"):
		if not os.path.exists(self.args.model_path + '/results/'+ self.args.scenario+"/figs"):
			os.makedirs(self.args.model_path + '/results/'+ self.args.scenario+"/figs")
		for animation_idx in range(len(input_list)):
			input = input_list[animation_idx]
			grid = grid_list[animation_idx]
			ped_grid = ped_grid_list[animation_idx]

			gt_vel = y_ground_truth_list[animation_idx]
			if not (y_pred_list_global is None):
				model_vel_pred = y_pred_list_global[animation_idx]
			other_agents_pos = other_agents_list[animation_idx]

			for step in range(input.shape[0]):  # trajectory length considering data from getTrajectoryAsBatch
				# Eucldean / position space
				self.ax_global_scenario.clear()
				#self.ax_global_scenario.plot(input[:, 0], input[:, 1], color='b', alpha=0.4, lw=1)
				#self.ax_global_scenario.plot(input[:step, 0], input[:step, 1], color='b', lw=2)
				self.ax_global_scenario.plot(input[step, 0], input[step, 1], color='c', marker='o')
				bottom_left_submap = input[step, 0:2] - np.array([self.args.submap_width / 2.0, self.args.submap_height / 2.0])
				#self.ax_global_scenario.add_patch(
				#	pl.Rectangle(bottom_left_submap, self.args.submap_width, self.args.submap_height, fill=None))

				# Plot real trajectory
				real_vel = np.zeros((1, self.args.prediction_horizon * self.args.output_dim))
				for i in range(self.args.prediction_horizon):
					idx = i * self.args.output_dim
					idy = i * self.args.output_dim + 1
					mu_x = gt_vel[step, idx]
					mu_y = gt_vel[step, idy]
					real_vel[0, 2 * i:2 * i + 2] = [mu_x, mu_y]

				traj_real = sup.path_from_vel(initial_pos=np.array([input[step, 0], input[step, 1]]),
				                              pred_vel=np.squeeze(real_vel), dt=self.args.dt)
				self.ax_global_scenario.plot(traj_real[:, 0], traj_real[:, 1], color='m', label='Agent real trajectory')

				# Other agents
				for id in range(other_agents_pos[step].shape[0]):  # number of agents
					self.ax_global_scenario.plot(other_agents_pos[step][id, 0], other_agents_pos[step][id, 1], marker='o', color='b',markersize=10)
				vel_pred = np.zeros((1, self.args.prediction_horizon * self.args.output_dim))
				# Predicted trajectory
				if not (y_pred_list_global is None):
					if self.args.n_mixtures == 0:
						# plot predicted trajectory global frame
						for i in range(self.args.prediction_horizon):
							idx = i * self.args.output_pred_state_dim
							idy = i * self.args.output_pred_state_dim + 1
							mu_x = model_vel_pred[step, idx]
							mu_y = model_vel_pred[step, idy]
							vel_pred[0, 2 * i:2 * i + 2] = [mu_x, mu_y]
						if rotate:
							pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
						else:
							pred_vel_global_frame = vel_pred
						traj_pred = sup.path_from_vel(initial_pos=np.array(
							[input[step, 0] * self.args.normalization_constant_x,
							 input[step, 1] * self.args.normalization_constant_y]),
							pred_vel=np.squeeze(pred_vel_global_frame), dt=self.args.dt)
						self.ax_global_scenario.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[0], label='Predicted trajectory')
					else:
						for sample_id in range(n_samples):
							for mix_idx in range(self.args.n_mixtures):
								# plot predicted trajectory global frame
								time = np.zeros([self.args.prediction_horizon])
								for pred_step in range(self.args.prediction_horizon):
									time[pred_step] = pred_step * self.args.dt
									idx = i * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id
									idy = i * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id
									mu_x = model_vel_pred[step, idx]
									mu_y = model_vel_pred[step, idy]
									vel_pred[0, 2 * i:2 * i + 2] = [mu_x, mu_y]

								if rotate:
									pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
								else:
									pred_vel_global_frame = vel_pred
								traj_pred = sup.path_from_vel(initial_pos=np.array(
									[input[step, 0] * self.args.normalization_constant_x,
									 input[step, 1] * self.args.normalization_constant_y]),
									pred_vel=np.squeeze(pred_vel_global_frame), dt=self.args.dt)
								# sub-sample to make smoother trajectories
								_, traj_pred_smooth, _ = self.smoothenTrajectory(time, traj_pred, vel_pred)


								self.ax_global_scenario.plot(traj_pred_smooth[:, 0], traj_pred_smooth[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')

								# prior of 0 on the uncertainty of the pedestrian velocity
								sigma_x = np.square(
									np.maximum(model_vel_pred[0, self.args.n_mixtures * 2 + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id],0.75)) * self.args.dt * self.args.dt + 0.5
								sigma_y = np.square(
									np.maximum(model_vel_pred[0, self.args.n_mixtures * 3 + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id],0.75)) * self.args.dt * self.args.dt +0.5

								e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
								             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
								e1.set_alpha(0.5)
								self.ax_global_scenario.add_patch(e1)

								for i in range(1, self.args.prediction_horizon):
									idx = i * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures * 2 + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id
									idy = i * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures * 3 + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id
									sigma_x += np.square(
										np.maximum(model_vel_pred[i - 1, idx],0.75)) * self.args.dt * self.args.dt
									sigma_y += np.square(
										np.maximum(model_vel_pred[i - 1, idy],0.75)) * self.args.dt * self.args.dt
									e1 = Ellipse(xy=(traj_pred[i, 0], traj_pred[i, 1]), width=np.sqrt(sigma_x) / 2,
									             height=np.sqrt(sigma_y) / 2,
									             angle=0 / np.pi * 180)
									e1.set_alpha(0.5)
									self.ax_global_scenario.add_patch(e1)

				sup.plot_grid(self.ax_global_scenario, np.array([0.0, 0.0]), self.gridmap.gridmap, self.gridmap.resolution,
				              self.gridmap.map_size)
				#self.ax_global_scenario.set_xlim([-self.gridmap.center[0], self.gridmap.center[0]])
				#self.ax_global_scenario.set_ylim([-self.gridmap.center[1], self.gridmap.center[1]])
				self.ax_global_scenario.set_xlim([input[step, 0]-7, input[step, 0]+7])
				self.ax_global_scenario.set_ylim([input[step, 1]-7, input[step, 1]+7])

				#self.ax_global_scenario.set_xlabel('x [m]')
				#self.ax_global_scenario.set_ylabel('y [m]')
				self.ax_global_scenario.axis('off')
				self.ax_global_scenario.set_aspect('equal')
				self.fig_global_scenario.savefig(self.args.model_path + '/results/'+ self.args.scenario+"/figs/"+"scene"+str(animation_idx)+"_step_"+str(step)+format, bbox_inches='tight')

	def smoothenTrajectory(self, time_vec,pose_vec,vel_vec,dt=0.05):
		"""
		Cubic interpolation with provided values in order to smoothen the trajectory and obtain
		a sequence with the specified dt.
		"""
		x_interpolator = interpolate.interp1d(time_vec, pose_vec[:, 0], kind='cubic', axis=0,fill_value='extrapolate')
		y_interpolator = interpolate.interp1d(time_vec, pose_vec[:, 1], kind='cubic', axis=0,fill_value='extrapolate')
		vx_interpolator = interpolate.interp1d(time_vec, vel_vec[:, 0], kind='cubic', axis=0,fill_value='extrapolate')
		vy_interpolator = interpolate.interp1d(time_vec, vel_vec[:, 1], kind='cubic', axis=0,fill_value='extrapolate')
		n_elem = self.args.prediction_horizon*int(self.args.dt/dt)
		new_time_vec = np.linspace(0, 0 + (n_elem - 1) * dt, n_elem)

		new_time_vec = np.zeros([n_elem])
		new_pose_vec = np.zeros([n_elem, 2])
		new_vel_vec = np.zeros([n_elem, 2])

		for ii in range(n_elem):
			t = 0 + ii * dt
			new_time_vec[ii] = t
			new_pose_vec[ii, 0] = x_interpolator(t)
			new_pose_vec[ii, 1] = y_interpolator(t)
			new_vel_vec[ii, 0] = vx_interpolator(t)
			new_vel_vec[ii, 1] = vy_interpolator(t)

		return new_time_vec, new_pose_vec, new_vel_vec

	def plot_on_video(self,input_list,grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,traj_list,all_traj_likelihood,test_args,social_trajectories=None):
		homography_file = os.path.join( self.args.data_path+ self.args.scenario, 'H.txt')
		if os.path.exists(homography_file):
			Hinv = np.linalg.inv(np.loadtxt(homography_file))
		else:
			print('[INF] No homography file')
		resolution = np.abs(1/Hinv[1,0])
		print("Resolution: " + str(resolution))
		scenario = self.args.scenario.split('/')[-1]
		video_file = self.args.data_path + self.args.scenario+'/' + scenario + '.avi'
		if os.path.exists(video_file):
			print('[INF] Using video file ' + video_file)
			cap = cv2.VideoCapture(video_file)
			if "zara_02" in self.args.scenario:
				time_offset = -22
			else:
				time_offset = 0

			if cap:
				dt = 0.4  # seconds (equivalent to 2.5 fps)
				#if 'seq_eth' in self.args.scenario:
				#	frames_between_annotation = 6.0
				#else:
				frames_between_annotation = 10.0

				frame_width = int(cap.get(3))
				frame_height = int(cap.get(4))

				# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
				if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario):
					os.makedirs(self.args.model_path + '/results/' + self.args.scenario)
				if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario+"/figs"):
					os.makedirs(self.args.model_path + '/results/' + self.args.scenario+"/figs")
				if social_trajectories:
					out = cv2.VideoWriter(self.args.model_path + '/results/' + self.args.scenario+'/outpy'+str(test_args.n_samples)+'_social.avi', cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (frame_width, frame_height))
				else:
					out = cv2.VideoWriter(self.args.model_path + '/results/' + self.args.scenario+'/outpy'+str(test_args.n_samples)+'.avi', cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (frame_width, frame_height))

			vel_real = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			sigmax = np.zeros((self.args.prediction_horizon, 1))
			sigmay = np.zeros((self.args.prediction_horizon, 1))
			pis = np.zeros((self.args.prediction_horizon, 3))
			for animation_idx in range(len(input_list)):
				traj = traj_list[animation_idx]

				if social_trajectories:
					social_traj = social_trajectories[animation_idx]
				gt_vel = y_ground_truth_list[animation_idx]
				input = input_list[animation_idx]
				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				traj_likelihood = all_traj_likelihood[animation_idx]
				# plot real trajectory global frame
				for step in range(input.shape[0]):

					time_stamp = traj.time_vec[step] / dt * frames_between_annotation + time_offset
					cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, time_stamp +self.args.prev_horizon* frames_between_annotation))
					ret, im = cap.read()
					overlay = im.copy()

					overlay_ellipses = im.copy()
					if not ret:
						print("Could not open the video")

					# Initial positions
					if self.args.normalize_data:
						x0 = input[step, 0] / self.args.sx_pos + self.args.min_pos_x
						y0 = input[step, 1] / self.args.sy_pos + self.args.min_pos_y
					else:
						x0 = input[step, 0]
						y0 = input[step, 1]

					# Real Trajectory
					for i in range(self.args.prediction_horizon):
						idx = i * self.args.output_dim
						idy = i * self.args.output_dim + 1
						mu_x = gt_vel[step, idx]
						mu_y = gt_vel[step, idy]
						vel_real[i, :] = [mu_x, mu_y]
					if self.args.rotated_grid:
						real_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_real, input[step, 2:])
					else:
						real_vel_global_frame = vel_real
					traj_real = sup.path_from_vel(initial_pos=np.array([x0,
					                                                    y0]),
					                              pred_vel=real_vel_global_frame, dt=self.args.dt)

					obsv_XY = sup.to_image_frame(Hinv, traj_real)
					sup.line_cv(overlay, obsv_XY, (255, 0, 0), 3) # bgr convention

					# Plot social trajectory
					if social_trajectories:
						for sample_id in range(3):
							social_traj_pred = sup.path_from_vel(initial_pos=np.array([input[step, 0],
						                                                    input[step, 1]]),
						                              pred_vel=social_traj[step][sample_id], dt=1)
							obsv_XY = sup.to_image_frame(Hinv, social_traj_pred)
							sup.line_cv(overlay, obsv_XY, (255, 0, 0), 3)  # bgr convention

					# Plot real predicted traj from positions
					traj_real = traj.pose_vec[step:step+self.args.prev_horizon+self.args.prediction_horizon,:2]
					obsv_XY = sup.to_image_frame(Hinv, traj_real)
					sup.line_cv(overlay, obsv_XY, (0, 0, 0), 3)  # bgr convention

					# Predicted trajectory
					colors = [(0,0,255),(0,255,0),(0,255,255)]
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							prediction_sample = model_vel_pred[step][0]
							# plot predicted trajectory global frame
							time = np.zeros([self.args.prediction_horizon])
							for i in range(self.args.prediction_horizon):
								time[i] = i*self.args.dt
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								if self.args.normalize_data:
									mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
									mu_y = prediction_sample[0, idy] / self.args.sy_vel + self.args.min_vel_y
								else:
									mu_x = prediction_sample[0, idx]
									mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]
							if self.args.rotated_grid:
								pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
							else:
								pred_vel_global_frame = vel_pred

							if self.args.predict_positions:
								traj_pred = vel_pred
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array(
									[x0,
									 y0]),
									pred_vel=pred_vel_global_frame, dt=self.args.dt)
							# sub-sample to make smoother trajectories
							_, traj_pred_smooth, _ = self.smoothenTrajectory(time,traj_pred,vel_pred)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred_smooth)
							sup.line_cv(overlay, obsv_XY, (0, 0, 255),2)
						else:
							for sample_id in range(test_args.n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								for mix_idx in range(self.args.n_mixtures): #self.args.n_mixtures
									# plot predicted trajectory global frame
									time = np.zeros([self.args.prediction_horizon])
									for pred_step in range(self.args.prediction_horizon):
										time[pred_step] = pred_step * self.args.dt
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										if self.args.normalize_data:
											mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
											mu_y = prediction_sample[0, idx + self.args.n_mixtures] / self.args.sy_vel + self.args.min_vel_y
										else:
											mu_x = prediction_sample[0, idx]
											mu_y = prediction_sample[0, idx + self.args.n_mixtures]
										if self.args.output_pred_state_dim > 2:
											sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
											sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										if self.args.output_pred_state_dim ==5:
											pis[pred_step, mix_idx] = prediction_sample[0][idx + 4 * self.args.n_mixtures]
										if math.isnan(mu_x) | math.isnan(mu_y):
											continue
										vel_pred[pred_step, :] = [mu_x, mu_y]

									if self.args.rotated_grid:
										pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
									else:
										pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array(
										[x0,
										 y0]),
										pred_vel=pred_vel_global_frame, dt=self.args.dt)
									# sub-sample to make smoother trajectories
									self.smoothenTrajectory(time, traj_pred, vel_pred)
									obsv_XY = sup.to_image_frame(Hinv, traj_pred)

									sup.line_cv(overlay, obsv_XY, colors[mix_idx], 2)

									if self.args.output_pred_state_dim > 2:
										# prior of 0 on the uncertainty of the pedestrian velocity
										sigma_x = np.square(min(sigmax[0],0.5)) * self.args.dt * self.args.dt +0.2
										sigma_y = np.square(min(sigmay[0],0.5)) * self.args.dt * self.args.dt +0.2
										if math.isnan(sigma_x) | math.isnan(sigma_y):
											continue
										axis = (int(np.sqrt(sigma_x) /resolution),int(np.sqrt(sigma_y) /resolution))
										center = (obsv_XY[0, 1], obsv_XY[0, 0])
										cv2.ellipse(overlay_ellipses, center,       axis,    0, 0, 360, (255,153,51),-1)

										for pred_step in range(1, self.args.prediction_horizon):
											sigma_x += np.square(min(sigmax[pred_step],0.5)) * self.args.dt * self.args.dt
											sigma_y += np.square(min(sigmay[pred_step],0.5)) * self.args.dt * self.args.dt
											if math.isnan(sigma_x) | math.isnan(sigma_y):
												continue
											axis = (int(np.sqrt(sigma_x) / resolution),
											        int(np.sqrt(sigma_y) / resolution))
											if (sigma_x<0) or (sigma_y<0):
												print("Negative ellipse")
											center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
											cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (255,153,51),-1)

					# Plot other agents
					plot_other_agents = True
					if plot_other_agents:
						other_agents = other_agents_list[animation_idx][step]
						axis = (int(np.sqrt(0.3) / resolution),
						        int(np.sqrt(0.3) / resolution))

						for agent_id in range(other_agents.shape[0]):
							obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(other_agents[agent_id,:2],axis=0))
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (153, 153, 51), -1)

					image_new = cv2.addWeighted(overlay, 0.6, im, 0.4, 0)
					image_new = cv2.addWeighted(image_new, 0.6, overlay_ellipses, 0.4, 0)

					# Adding legend
					font = cv2.FONT_HERSHEY_SIMPLEX
					cv2.line(image_new, (frame_width - 270, 30), (frame_width - 250, 30), (0, 0, 0), 4)
					cv2.putText(image_new, "Real Trajectory", (frame_width - 240, 30), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)
					cv2.line(image_new, (frame_width - 270, 50), (frame_width - 250, 50), (0, 0, 255), 4)
					if self.args.output_pred_state_dim == 5:
						cv2.putText(image_new, "Predicted Trajectory Prob " + str(np.sum(pis[:,0])/self.args.prediction_horizon*100), (frame_width - 240, 50), font, 0.5, (0, 0, 255), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 70), (frame_width - 250, 70), (0, 255, 0), 4)
						cv2.putText(image_new, "Predicted Trajectory Prob " + str(np.sum(pis[:,1])/self.args.prediction_horizon*100), (frame_width - 240, 70), font, 0.5, (0, 255, 0), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 90), (frame_width - 250, 90), (0, 255, 255), 4)
						cv2.putText(image_new, "Predicted Trajectory Prob " + str(np.sum(pis[:,2])/self.args.prediction_horizon*100), (frame_width - 240, 90), font, 0.5, (0, 255, 255), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 110), (frame_width - 250, 110), (255, 153, 51), 4)
					else:
						cv2.putText(image_new, "Predicted Trajectory: " + str(np.around(traj_likelihood[step][0,0,0],decimals=2)), (frame_width - 240, 50), font, 0.5, (0, 0, 255), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 70), (frame_width - 250, 70), (0, 255, 0), 4)
						cv2.putText(image_new, "Predicted Trajectory: " + str(np.around(traj_likelihood[step][0,0,1],decimals=2)), (frame_width - 240, 70), font, 0.5, (0, 255, 0), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 90), (frame_width - 250, 90), (0, 255, 255), 4)
						cv2.putText(image_new, "Predicted Trajectory: " + str(np.around(traj_likelihood[step][0,0,2],decimals=2)), (frame_width - 240, 90), font, 0.5, (0, 255, 255), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 110), (frame_width - 250, 110), (255, 153, 51), 4)
					cv2.putText(image_new, "1-sigma Uncertainty", (frame_width - 240, 110), font, 0.5, (255,153,51), 2,
					            cv2.LINE_AA)

					out.write(image_new)

					#cv2.imwrite(self.args.model_path + '/results/' + self.args.scenario+"/figs/result_"+str(animation_idx)+"_"+str(step)+".jpg", image_new);

			# When everything done, release the video capture and video write objects
			cap.release()
			out.release()

			# Closes all the frames
			cv2.destroyAllWindows()

		else:
			"Video not found"

	def plot_on_image(self,input_list, grid_list, y_pred_list_global, y_ground_truth_list, other_agents_list, traj_list, test_args, social_trajectories=None):
		"""
		Not working yet for Herengracht map. Seems to be something wrong with the H.txt file and size of plotting
		"""

		scenario = self.args.scenario.split('/')[-1]

		if not os.path.exists(self.args.model_path + '/../videos/'):
			os.makedirs(self.args.model_path + '/../videos/')
		if test_args.freeze_other_agents:
			video_file = self.args.model_path + '/../videos/'+scenario + str(self.args.exp_num) + "_frozen.avi"
			map_file = self.args.data_path + self.args.scenario + '/map.png'
			homography_file = os.path.join(self.args.data_path + self.args.scenario, 'H.txt')
		elif test_args.unit_testing:
			video_file = self.args.model_path + '/../videos/'+scenario + str(self.args.exp_num) + "_unit_tests.avi"
			map_file = self.args.data_path + self.args.scenario + '/unit_map.png'
			homography_file = os.path.join(self.args.data_path + self.args.scenario, 'unit_H.txt')
		else:
			video_file = self.args.model_path + '/../videos/'+scenario + str(self.args.exp_num) + "_final.avi"
			map_file = self.args.data_path + self.args.scenario + '/map.png'
			homography_file = os.path.join(self.args.data_path + self.args.scenario, 'H.txt')

		if os.path.exists(homography_file):
			Hinv = np.linalg.inv(np.loadtxt(homography_file))
		else:
			print('[INF] No homography file')
		resolution = self.args.submap_resolution
		print("Resolution: " + str(resolution))

		im = np.uint8(cv2.imread(map_file) * -1 + 255)
		#im = cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE)
		frame_height, frame_width, layers = im.shape
		scale_factor = 2
		resized_img = cv2.resize(im, (frame_width * scale_factor, frame_height * scale_factor),
		                         interpolation=cv2.INTER_AREA)
		print(resized_img.shape)

		h, w, c = resized_img.shape
		blank_image = 255 * np.ones(shape=(h, w, c), dtype=np.uint8)


		if os.path.exists(map_file):
			print('[INF] Using video file ' + video_file)
			cap = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 20,
			                      (frame_width * scale_factor, frame_height * scale_factor))

			vel_real = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			sigmax = np.zeros((self.args.prediction_horizon, 1))
			sigmay = np.zeros((self.args.prediction_horizon, 1))
			pis = np.zeros((self.args.prediction_horizon, 3))
			for animation_idx in range(len(input_list)):
				traj = traj_list[animation_idx]

				if social_trajectories:
					social_traj = social_trajectories[animation_idx]
				gt_vel = y_ground_truth_list[animation_idx]
				input = input_list[animation_idx]
				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				# plot real trajectory global frame
				for step in range(input.shape[0]):

					overlay = resized_img.copy()
					#overlay = blank_image.copy()

					# Initial positions
					if self.args.normalize_data:
						x0 = input[step, 0] / self.args.sx_pos + self.args.min_pos_x
						y0 = input[step, 1] / self.args.sy_pos + self.args.min_pos_y
					else:
						x0 = input[step, 0]
						y0 = input[step, 1]

					# Real Trajectory
					for i in range(self.args.prediction_horizon):
						idx = i * self.args.output_dim
						idy = i * self.args.output_dim + 1
						mu_x = gt_vel[step, idx] + input[step, 2]
						mu_y = gt_vel[step, idy] + input[step, 3]
						vel_real[i, :] = [mu_x, mu_y]

					real_vel_global_frame = vel_real
					traj_real = sup.path_from_vel(initial_pos=np.array([x0,
					                                                    y0]),
					                              pred_vel=real_vel_global_frame, dt=self.args.dt)

					obsv_XY = sup.to_image_frame(Hinv, traj_real)*scale_factor
					sup.line_cv(overlay, obsv_XY, (255, 0, 0), 3) # bgr convention

					# Plot social trajectory
					if social_trajectories:
						for sample_id in range(3):
							social_traj_pred = sup.path_from_vel(initial_pos=np.array([input[step, 0],
						                                                    input[step, 1]]),
						                              pred_vel=social_traj[step][sample_id], dt=1)
							obsv_XY = sup.to_image_frame(Hinv, social_traj_pred)*scale_factor
							sup.line_cv(overlay, obsv_XY, (255, 0, 0), 3)  # bgr convention

					# Plot real predicted traj from positions
					#traj_real = traj.pose_vec[step:step+self.args.prev_horizon+self.args.prediction_horizon,:2]
					#obsv_XY = sup.to_image_frame(Hinv, traj_real)*scale_factor
					#sup.line_cv(overlay, obsv_XY, (0, 0, 0), 3)  # bgr convention

					# Predicted trajectory
					colors = [(0,0,255),(0,255,0),(0,255,255)]
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							prediction_sample = model_vel_pred[step][0]
							# plot predicted trajectory global frame
							time = np.zeros([self.args.prediction_horizon])
							for i in range(self.args.prediction_horizon):
								time[i] = i*self.args.dt
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								if self.args.normalize_data:
									mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
									mu_y = prediction_sample[0, idy] / self.args.sy_vel + self.args.min_vel_y
								else:
									mu_x = prediction_sample[0, idx]
									mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]

							pred_vel_global_frame = vel_pred

							if self.args.predict_positions:
								traj_pred = vel_pred
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array(
									[x0,
									 y0]),
									pred_vel=pred_vel_global_frame, dt=self.args.dt)
							# sub-sample to make smoother trajectories
							#_, traj_pred_smooth, _ = self.smoothenTrajectory(time,traj_pred,vel_pred)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)*scale_factor
							sup.line_cv(overlay, obsv_XY, (0, 0, 255),2)
							#print(obsv_XY)
						else:
							for sample_id in range(test_args.n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								for mix_idx in range(self.args.n_mixtures): #self.args.n_mixtures
									# plot predicted trajectory global frame
									time = np.zeros([self.args.prediction_horizon])
									for pred_step in range(self.args.prediction_horizon):
										time[pred_step] = pred_step * self.args.dt
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										if self.args.normalize_data:
											mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
											mu_y = prediction_sample[0, idx + self.args.n_mixtures] / self.args.sy_vel + self.args.min_vel_y
										else:
											mu_x = prediction_sample[0, idx] + input[step, 2]
											mu_y = prediction_sample[0, idx + self.args.n_mixtures] + input[step, 3]
										if self.args.output_pred_state_dim > 2:
											sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
											sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										if self.args.output_pred_state_dim ==5:
											pis[pred_step, mix_idx] = prediction_sample[0][idx + 4 * self.args.n_mixtures]
										if math.isnan(mu_x) | math.isnan(mu_y):
											continue
										vel_pred[pred_step, :] = [mu_x, mu_y]

									pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array(
										[x0,
										 y0]),
										pred_vel=pred_vel_global_frame, dt=self.args.dt)
									# sub-sample to make smoother trajectories
									#self.smoothenTrajectory(time, traj_pred, vel_pred)
									obsv_XY = sup.to_image_frame(Hinv, traj_pred)*scale_factor

									sup.line_cv(overlay, obsv_XY, colors[mix_idx], 2)
									print(obsv_XY)

									try:
										if self.args.output_pred_state_dim > 2:
											# prior of 0 on the uncertainty of the pedestrian velocity
											sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt +0.2
											sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt +0.2
											if math.isnan(sigma_x) | math.isnan(sigma_y):
												continue
											axis = (int(np.sqrt(sigma_x) /resolution*scale_factor),int(np.sqrt(sigma_y) /resolution*scale_factor))
											y = obsv_XY[0, 1]*scale_factor
											x = obsv_XY[0, 0]*scale_factor
											center = (y , x)
											cv2.ellipse(overlay, center,       axis,    0, 0, 360, (255,153,51),-1)

											for pred_step in range(1, self.args.prediction_horizon):
												sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
												sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt
												if math.isnan(sigma_x) | math.isnan(sigma_y):
													continue
												axis = (int(np.sqrt(sigma_x) / resolution*scale_factor),
												        int(np.sqrt(sigma_y) / resolution*scale_factor))
												if (sigma_x<0) or (sigma_y<0):
													print("Negative ellipse")
												y = obsv_XY[pred_step, 1] * scale_factor
												x = obsv_XY[pred_step, 0] * scale_factor
												center = (y, x)
												cv2.ellipse(overlay, center, axis, 0, 0, 360, (255,153,51),-1)
									except:
										print("Failed to add ellipse")

					# Plot other agents
					plot_other_agents = True
					if plot_other_agents:
						other_agents = other_agents_list[animation_idx][step]
						axis = (int(np.sqrt(0.3) / resolution),
						        int(np.sqrt(0.3) / resolution))

						for agent_id in range(other_agents.shape[0]):
							obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(other_agents[agent_id,:2],axis=0))
							center = (int(obsv_XY[0, 1]*scale_factor), int(obsv_XY[0, 0]*scale_factor))
							cv2.ellipse(overlay, center, axis, 0, 0, 360, (153, 153, 51), -1)

					# Adding legend
					font = cv2.FONT_HERSHEY_SIMPLEX
					cv2.line(overlay, (int(frame_height*scale_factor*0.9-50), 30), (int(frame_height*scale_factor*0.9-40), 30), (255, 0, 0), 4)
					cv2.putText(overlay, "Real Trajectory", (int(frame_height*scale_factor*0.9-30), 30), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
					cv2.line(overlay, (int(frame_height*scale_factor*0.9-50), 50), (int(frame_height*scale_factor*0.9-40), 50), (0, 0, 255), 4)
					cv2.putText(overlay, "Predicted Trajectory", (int(frame_height*scale_factor * 0.9 - 40), 50), font, 0.5, (0, 0, 255), 2,
					            cv2.LINE_AA)

					cap.write(overlay)

					#cv2.imwrite(self.args.model_path + '/results/' + self.args.scenario+"/figs/result_"+str(animation_idx)+"_"+str(step)+".jpg", overlay);

			# When everything done, release the video capture and video write objects
			cap.release()

			# Closes all the frames
			cv2.destroyAllWindows()

		else:
			print("map_file not found")

	def plot_GA3C(self,input_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,traj_list,rotate=False,n_samples=1,social_trajectories=None):
		homography_file = os.path.join( self.args.data_path+ self.args.scenario, 'H.txt')
		if os.path.exists(homography_file):
			Hinv = np.linalg.inv(np.loadtxt(homography_file))
		else:
			print('[INF] No homography file')
		resolution = np.abs(1/Hinv[0,0])/2
		print("Resolution: " + str(resolution))
		scenario = self.args.scenario.split('/')[-1]
		map_file = self.args.data_path + self.args.scenario+ '/map.png'
		im = cv2.imread(map_file) +255
		height, width, layers = im.shape
		scale_factor = 2
		resized_img = cv2.resize(im, (height*scale_factor, width*scale_factor), interpolation=cv2.INTER_AREA)

		video_file = self.args.model_path + '/results/' + self.args.scenario+'/outpy'+'.avi'
		if os.path.exists(map_file):
			print('[INF] Using map file ' + map_file)
			cap = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width*scale_factor, height*scale_factor))


			vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			sigmax = np.zeros((self.args.prediction_horizon, 1))
			sigmay = np.zeros((self.args.prediction_horizon, 1))
			pis = np.zeros((self.args.prediction_horizon, 3))

			for animation_idx in range(len(input_list)):
				traj = traj_list[animation_idx]

				gt_vel = y_ground_truth_list[animation_idx]
				input = input_list[animation_idx]
				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None

				# Real Trajectory
				traj_real = np.zeros((len(traj), self.args.output_dim))
				for i in range(len(traj)):
					x = traj[i]["pedestrian_state"]["position"][0]
					y = traj[i]["pedestrian_state"]["position"][1]
					traj_real[i, :] = [x, y]

				# plot real trajectory global frame
				for step in range(len(model_vel_pred)):

					overlay = resized_img.copy()

					# Initial positions
					x0 = traj[0]["pedestrian_state"]["position"][0]
					y0 = traj[0]["pedestrian_state"]["position"][1]
					obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(np.array([x0,y0]),axis=0))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					axis = (int(np.sqrt(0.3) / resolution),
					        int(np.sqrt(0.3) / resolution))
					cv2.ellipse(overlay, center, axis, 0, 0, 360, (255, 0, 0), -1)

					# Final positions
					x0 = traj[-1]["pedestrian_state"]["position"][0]
					y0 = traj[-1]["pedestrian_state"]["position"][1]
					obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(np.array([x0, y0]), axis=0))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					axis = (int(np.sqrt(0.3) / resolution),
					        int(np.sqrt(0.3) / resolution))
					cv2.ellipse(overlay, center, axis, 0, 0, 360, (0, 255, 0), -1)

					obsv_XY = sup.to_image_frame(Hinv, traj_real)
					sup.line_cv(overlay, obsv_XY, (0, 0, 0), 3) # bgr convention

					# Predicted trajectory
					colors = [(0,0,255),(0,255,0),(0,255,255)]
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							prediction_sample = model_vel_pred[step][0]
							# plot predicted trajectory global frame
							time = np.zeros([self.args.prediction_horizon])
							for i in range(self.args.prediction_horizon):
								time[i] = i*self.args.dt
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								if self.args.normalize_data:
									mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
									mu_y = prediction_sample[0, idy] / self.args.sy_vel + self.args.min_vel_y
								else:
									mu_x = prediction_sample[0, idx]
									mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]

							if True: #self.args.predict_positions
								traj_pred = vel_pred
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array(
									[x0,
									 y0]),
									pred_vel=pred_vel_global_frame, dt=self.args.dt)
							# sub-sample to make smoother trajectories
							#_, traj_pred_smooth, _ = self.smoothenTrajectory(time,traj_pred,vel_pred)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)
							sup.line_cv(overlay, obsv_XY, (0, 0, 255),2)
						else:
							for sample_id in range(n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								for mix_idx in range(self.args.n_mixtures): #self.args.n_mixtures
									# plot predicted trajectory global frame
									time = np.zeros([self.args.prediction_horizon])
									for pred_step in range(self.args.prediction_horizon):
										time[pred_step] = pred_step * self.args.dt
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										if self.args.normalize_data:
											mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
											mu_y = prediction_sample[0, idx + self.args.n_mixtures] / self.args.sy_vel + self.args.min_vel_y
										else:
											mu_x = prediction_sample[0, idx]
											mu_y = prediction_sample[0, idx + self.args.n_mixtures]
										sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
										sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										if self.args.output_pred_state_dim ==5:
											pis[pred_step, mix_idx] = prediction_sample[0][idx + 4 * self.args.n_mixtures]
										if math.isnan(mu_x) | math.isnan(mu_y):
											continue
										vel_pred[pred_step, :] = [mu_x, mu_y]

									if rotate:
										pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
									else:
										pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array(
										[x0,
										 y0]),
										pred_vel=pred_vel_global_frame, dt=self.args.dt)
									# sub-sample to make smoother trajectories
									#self.smoothenTrajectory(time, traj_pred, vel_pred)
									obsv_XY = sup.to_image_frame(Hinv, traj_pred)

									sup.line_cv(overlay, obsv_XY, colors[mix_idx], 2)

									# prior of 0 on the uncertainty of the pedestrian velocity
									sigma_x = np.square(min(sigmax[0],0.5)) * self.args.dt * self.args.dt +0.2
									sigma_y = np.square(min(sigmay[0],0.5)) * self.args.dt * self.args.dt +0.2
									if math.isnan(sigma_x) | math.isnan(sigma_y):
										continue
									axis = (int(np.sqrt(sigma_x) /resolution),int(np.sqrt(sigma_y) /resolution))
									center = (obsv_XY[0, 1], obsv_XY[0, 0])
									cv2.ellipse(overlay, center,       axis,    0, 0, 360, (255,153,51),-1)

									for pred_step in range(1, self.args.prediction_horizon):
										sigma_x += np.square(min(sigmax[pred_step],0.5)) * self.args.dt * self.args.dt
										sigma_y += np.square(min(sigmay[pred_step],0.5)) * self.args.dt * self.args.dt
										if math.isnan(sigma_x) | math.isnan(sigma_y):
											continue
										axis = (int(np.sqrt(sigma_x) / resolution),
										        int(np.sqrt(sigma_y) / resolution))
										if (sigma_x<0) or (sigma_y<0):
											print("Negative ellipse")
										center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
										cv2.ellipse(overlay, center, axis, 0, 0, 360, (255,153,51),-1)

					# Plot other agents
					plot_other_agents = True
					if plot_other_agents:
						other_agents = np.expand_dims(other_agents_list[animation_idx][0,step],axis=0)
						axis = (int(np.sqrt(0.3) / resolution),
						        int(np.sqrt(0.3) / resolution))

						for agent_id in range(other_agents.shape[0]):
							obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(other_agents[agent_id],axis=0))
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							cv2.ellipse(overlay, center, axis, 0, 0, 360, (153, 153, 51), -1)

					# Adding legend
					font = cv2.FONT_HERSHEY_SIMPLEX
					cv2.line(overlay, (width - 70, 30), (width - 50, 30), (0, 0, 0), 4)
					cv2.putText(overlay, "Real Trajectory", (width - 40, 30), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)
					cv2.line(overlay, (width - 70, 50), (width - 50, 50), (0, 0, 255), 4)
					if self.args.output_pred_state_dim == 5:
						cv2.putText(overlay, "Predicted Trajectory Prob " + str(np.sum(pis[:,0])/self.args.prediction_horizon*100), (width - 40, 50), font, 0.5, (0, 0, 255), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 70), (width - 50, 70), (0, 255, 0), 4)
						cv2.putText(overlay, "Predicted Trajectory Prob " + str(np.sum(pis[:,1])/self.args.prediction_horizon*100), (width - 40, 70), font, 0.5, (0, 255, 0), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 90), (width - 50, 90), (0, 255, 255), 4)
						cv2.putText(overlay, "Predicted Trajectory Prob " + str(np.sum(pis[:,2])/self.args.prediction_horizon*100), (width - 40, 90), font, 0.5, (0, 255, 255), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 110), (width - 50, 110), (255, 153, 51), 4)
					else:
						cv2.putText(overlay, "Predicted Trajectory Mode 0", (width - 40, 50), font, 0.5, (0, 0, 255), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 70), (width - 50, 70), (0, 255, 0), 4)
						cv2.putText(overlay, "Goal Position", (width - 40, 70), font, 0.5, (0, 255, 0), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 90), (width - 50, 90), (255, 0, 0), 4)
						cv2.putText(overlay, "Initial", (width - 40, 90), font, 0.5, (255, 0, 0), 2,
						            cv2.LINE_AA)

					cap.write(overlay)

					cv2.imwrite(self.args.model_path + '/results/' + self.args.scenario+"/figs/result_"+str(animation_idx)+"_"+str(step)+".jpg", overlay);

			# When everything done, release the video capture and video write objects
			cap.release()

			# Closes all the frames
			cv2.destroyAllWindows()

		else:
			print("Video not found")

	def animate_GA3C(self,input_list,goal_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,traj_list,test_args):
		#self.args.dt = 0.1
		pl.show(block=False)
		if not os.path.exists(self.args.model_path + '/results/'+ self.args.scenario):
			os.makedirs(self.args.model_path + '/results/'+ self.args.scenario)
		if test_args.freeze_other_agents:
			video_file = self.args.model_path + '/results/' + self.args.scenario+"/frozen.mp4"
		else:
			video_file = self.args.model_path + '/results/' + self.args.scenario + "/final.mp4"
		with self.writer.saving(self.fig_ga3c, video_file, 100):
			for animation_idx in range(len(input_list)):
				input = input_list[animation_idx]
				goal = goal_list[animation_idx]
				traj = traj_list[animation_idx]
				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				other_agents_pos = other_agents_list[animation_idx]

				# Real Trajectory
				traj_real = np.zeros((len(traj), self.args.output_dim))
				for i in range(len(traj)):
					x = traj[i]["pedestrian_state"]["position"][0]
					y = traj[i]["pedestrian_state"]["position"][1]
					traj_real[i, :] = [x, y]

				for step in range(input.shape[0]): # trajectory length considering data from getTrajectoryAsBatch
					# Eucldean / position space
					self.ga3c_plot.clear()
					#self.ga3c_plot.plot(input[:, 0], input[:, 1], color='c', alpha=0.4, lw=1)
					#self.ga3c_plot.plot(input[:step, 0], input[:step, 1], color='b', lw=2)
					# Initial positions
					x0 = traj[0]["pedestrian_state"]["position"][0]
					y0 = traj[0]["pedestrian_state"]["position"][1]
					self.ga3c_plot.plot(x0, y0, color='r', marker='o')

					xf = traj[-1]["pedestrian_state"]["position"][0]
					yf = traj[-1]["pedestrian_state"]["position"][1]
					self.ga3c_plot.plot(xf, yf, color='g', marker='x')

					xf = goal[0,0]
					yf = goal[0,1]
					self.ga3c_plot.plot(xf, yf, color='g', marker='o')

					# Other agents
					for id in range(other_agents_pos.shape[0]): # number of agents
						self.ga3c_plot.plot(other_agents_pos[id][step, 0], other_agents_pos[id][step, 1], marker='o', color='b')
					vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
					sigmax = np.zeros((self.args.prediction_horizon, 1))
					sigmay = np.zeros((self.args.prediction_horizon, 1))
					#Predicted trajectory
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							prediction_sample = model_vel_pred[step][0]
							# plot predicted trajectory global frame
							time = np.zeros([self.args.prediction_horizon])
							for i in range(self.args.prediction_horizon):
								time[i] = i*self.args.dt
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								if self.args.normalize_data:
									mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
									mu_y = prediction_sample[0, idy] / self.args.sy_vel + self.args.min_vel_y
								else:
									mu_x = prediction_sample[0, idx]
									mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]
							x0 = traj[step]["pedestrian_state"]["position"][0]
							y0 = traj[step]["pedestrian_state"]["position"][1]
							if self.args.predict_positions:
								traj_pred = vel_pred
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array([x0,y0]),
									pred_vel=vel_pred, dt=self.args.dt)
							self.ga3c_plot.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[0], label='Predicted trajectory')
						else:
							for sample_id in range(test_args.n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								x0 = traj[step]["pedestrian_state"]["position"][0]
								y0 = traj[step]["pedestrian_state"]["position"][1]
								for mix_idx in range(self.args.n_mixtures):
									# plot predicted trajectory global frame
									for pred_step in range(self.args.prediction_horizon):
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										mu_x = prediction_sample[0, idx]
										mu_y = prediction_sample[0, idx + self.args.n_mixtures]
										sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
										sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										vel_pred[pred_step, :] = [mu_x, mu_y]

									pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array([x0,y0]),pred_vel=pred_vel_global_frame, dt=self.args.dt)

									self.ga3c_plot.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')

									# prior of 0 on the uncertainty of the pedestrian velocity
									sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
									sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

									e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
									             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
									e1.set_alpha(0.5)
									self.ga3c_plot.add_patch(e1)
									#self.ga3c_plot.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')
									for pred_step in range(1, self.args.prediction_horizon):
										sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
										sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt
										e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
										             height=np.sqrt(sigma_y) / 2,
										             angle=0 / np.pi * 180)
										self.ga3c_plot.add_patch(e1)

						# Plot other agents
						plot_other_agents = True
						if plot_other_agents:
							other_agents = np.expand_dims(other_agents_list[animation_idx][0, step], axis=0)

							for agent_id in range(other_agents.shape[0]):
								self.ga3c_plot.plot(other_agents[agent_id,0], other_agents[agent_id,1], color='b', marker='o')

					# Real trajectory
					self.ga3c_plot.plot(traj_real[:, 0], traj_real[:, 1], color="k",marker='x', label='Real trajectory')

					model_vel_real = np.zeros((self.args.prediction_horizon, self.args.output_dim))
					for i in range(self.args.prediction_horizon):
						mu_x = traj[step+i]["pedestrian_state"]["velocity"][0]
						mu_y = traj[step+i]["pedestrian_state"]["velocity"][1]
						model_vel_real[i, :] = [mu_x, mu_y]
					if False:
						real_vel = sup.rotate_predicted_vel_to_global_frame(model_vel_real, batch_x[step, 2:])
					else:
						real_vel = model_vel_real
					traj_real_from_v = sup.path_from_vel(initial_pos=np.array([x0, y0]),
					                              pred_vel=real_vel, dt=self.args.dt)
					self.ga3c_plot.plot(traj_real_from_v[:, 0], traj_real_from_v[:, 1], color='m', label='Agent real trajectory')


					self.ga3c_plot.set_xlim([-15, 15])
					self.ga3c_plot.set_ylim([-15, 15])
					self.ga3c_plot.axis("on")
					self.ga3c_plot.set_xlabel('x [m]')
					self.ga3c_plot.set_ylabel('y [m]')
					self.ga3c_plot.set_aspect('equal')

					self.fig_ga3c.canvas.draw()
					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario+"/figs/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario+"/figs/")
					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario + "/figs_frozen/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario + "/figs_frozen/")
					if test_args.save_figs:
						if test_args.freeze_other_agents:
							self.fig_ga3c.savefig(self.args.model_path + '/results/' + self.args.scenario+"/figs_frozen/result_"+str(animation_idx)+"_"+str(step)+".jpg")
						else:
							self.fig_ga3c.savefig(self.args.model_path + '/results/' + self.args.scenario + "/figs/result_" + str(
								animation_idx) + "_" + str(step) + ".jpg")
					self.writer.grab_frame()

build/lib/data_utils/Recorder.py
eliatrevisan/IA-MPPI-LBM
import os
import numpy as np
from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter
from matplotlib.patches import Ellipse
# from mpl_toolkits.mplot3d import Axes3D
# import mpl_toolkits.mplot3d.axes3d as p3

from src.data_utils.DataHandler_Keras_v2 import prediction_to_trajectory

class Recorder(object): 
    def __init__(self, all_trajectories, all_predictions, args, save = False, display = True, figsize = (8,6), nframes = 10000): 
        self.dt = args.dt
        self.all_predictions = all_predictions
        self.all_trajectories = all_trajectories
        self.args = args
        # self.fps = int(1/self.dt)
        self.fps = int(10/self.dt)
        
        self.fig = plt.figure(figsize=figsize)
        self.ax = plt.axes(xlim=(-15, 15), ylim=(-15, 15))
        
        self.stream = self.data_stream()

        total_frames = 0
        for prediction in all_predictions:
            total_frames += prediction.shape[0]

        self.nframes = min(nframes, total_frames) 

        self.ani = animation.FuncAnimation(self.fig, self.animate, frames = self.nframes, init_func=self.setup_plot, interval=1/self.fps, blit=False, repeat=False)

        if save:
            self.save_animation(self.ani)

        if display:
            plt.show()


    def setup_plot(self):
        # Axis setup
        self.ax.set_xlabel('X (m)')
        self.ax.set_ylabel('Y (m)')

        pedestrian, = self.ax.plot([], [], marker="o", color='b')
        pedestrian_goal, = self.ax.plot([], [],  marker="D", color='b')
        robot, = self.ax.plot([], [], marker="o", color='r')
        robot_goal, = self.ax.plot([], [], marker="D", color='r')
        complete_path, = self.ax.plot([], [], marker="o", fillstyle='none', markersize=5, lw=1.5)
        predicted_path, = self.ax.plot([], [], marker="x", markersize=5, lw=1.5, color='g')
        self.objects = {
            "pedestrian": pedestrian,
            "pedestrian_goal": pedestrian_goal,
            "robot": robot,
            "robot_goal": robot_goal,
            "complete_path": complete_path,
            "predicted_path": predicted_path
        }
        return self.objects,

    def animate(self, iteration):
        data = next(self.stream)
        
        prediction = prediction_to_trajectory(data["pedestrian_position"], data["prediction"], self.args)
        
        self.objects["pedestrian"].set_data(data["pedestrian_position"][0], data["pedestrian_position"][1])
        self.objects["robot"].set_data(data["robot_position"][0], data["robot_position"][1])
        self.objects["pedestrian_goal"].set_data(data["pedestrian_goal"][0], data["pedestrian_goal"][1])
        self.objects["robot_goal"].set_data(data["robot_goal"][0], data["robot_goal"][1])
        self.objects["complete_path"].set_data(data["complete_path"][:, 0], data["complete_path"][:, 1])
        self.objects["predicted_path"].set_data( prediction[:, 0], prediction[:, 1] )
        
        return self.objects,

    def data_stream(self):
        """ Generator for the data to be plotted """
        while True:
            for prediction_array, trajectory in zip(self.all_predictions, self.all_trajectories):
                complete_path = np.zeros((len(trajectory), 2))
                for idx, datapoint in enumerate(trajectory):
                    complete_path[idx, :] = datapoint['pedestrian_state']['position']
                    
                for prediction, datapoint in zip(prediction_array, trajectory):
                    yield {
                        "prediction": prediction,
                        "pedestrian_position": datapoint['pedestrian_state']['position'],
                        "robot_position": datapoint['robot_state'],
                        "pedestrian_goal": datapoint["pedestrian_goal_position"],
                        # "robot_goal": trajectory[-1]["robot_state"][0:2],
                        "robot_goal": datapoint["goal_position"],
                        "complete_path": complete_path
                    }

    def save_animation(self, ani):
        self.writer =  animation.writers['ffmpeg']
        self.writer = self.writer(fps=self.fps, metadata=dict(artist='Me'), bitrate=1800, extra_args=['-vcodec', 'libx264'])
        video_dir = os.path.join(self.args.model_path, self.args.scenario)
        Path(video_dir).mkdir(parents=True, exist_ok=True)
        video_name = os.path.join(video_dir, 'video.mp4')
        self.ani.save(video_name, writer=self.writer)

build/lib/data_utils/Recorder_Keras.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import pylab as pl
import math
import sys
import os

ros_path = '/opt/ros/kinetic/lib/python2.7/dist-packages'
if ros_path in sys.path:
    sys.path.remove(ros_path)
import cv2
import cv2 as cv
sys.path.append('../../src')

def compute_radial_distance_vector(number_elements, relative_positions, max_range=10.0, min_angle=0, max_angle=2*np.pi, normalize=False):
  """
  Compute the distance to surrounding objects in a radially discretized way.
  For each element there will be a floating point distance to the closest object in this sector.
  This allows to preserve the continuous aspect of the distance vs. a standard grid.

  !!! Attention: 0 angle is at the negative x-axis.

  number_elements: radial discretization
  relative_positions: relative positions of the surrounding objects in the local frame (numpy array)
  max_range: maximum range of the distance measurement

  returns:
  radial_dist_vector: contains the distance to the closest object in each sector
  """
  radial_dist_vector = max_range * np.ones([number_elements])
  radial_resolution = (max_angle-min_angle) / float(number_elements)
  if np.any(np.isnan(relative_positions)):
    print("inf")
  for ii in range(relative_positions.shape[0]):
    phi = math.atan2(relative_positions[ii, 1], relative_positions[ii, 0]) + np.pi
    rad_idx = int((phi - min_angle) / radial_resolution)
    if rad_idx >= 0 and rad_idx < number_elements:
      radial_dist_vector[rad_idx] = min(radial_dist_vector[rad_idx], np.linalg.norm(relative_positions[ii,:]))
  if normalize:
    radial_dist_vector /= max_range

  return radial_dist_vector

def sigmoid(x):
  return 1.0/(1.0+np.exp(-x))

def str2bool(v):
  if isinstance(v, bool):
    return v
  if v.lower() in ('yes', 'true', 't', 'y', '1'):
    return True
  elif v.lower() in ('no', 'false', 'f', 'n', '0'):
    return False
  else:
    raise argparse.ArgumentTypeError('Boolean value expected.')

def positions_in_local_frame(ego_pos, heading, other_pos):
  proj_matrix = np.array([[np.cos(heading), np.sin(heading)], [-np.sin(heading), np.cos(heading)]])
  local_frame_pos = np.zeros([0, 2])

  for ii in range(other_pos.shape[0]):
    rel_pos = other_pos[ii, :] - ego_pos
    local_frame_pos = np.append(local_frame_pos, np.expand_dims(np.matmul(proj_matrix, rel_pos), axis=0), axis=0)

  return local_frame_pos

def rotate_grid_around_center(grid, angle):
    """
    inputs:
      grid: numpy array (gridmap) that needs to be rotated
      angle: rotation angle in degrees
    """
    # Rotate grid into direction of initial heading
    grid = grid.copy()
    rows, cols = grid.shape
    M = cv2.getRotationMatrix2D(center=(rows / 2, cols / 2), angle=angle, scale=1)
    grid = cv2.warpAffine(grid, M, (rows, cols))

    return grid

def rotate_batch(batch_y, batch_x):
  """
    inputs:
      batch_y: trainning velocities batch containing vx, vy on the global frame
      batch_x: trainning or prediction state batch containing x,y, vx, vy on the global frame
    """
  # rotate initial vx and vy in global frame to quary-agent local frame
  by = batch_y.copy()  # rotated output batch values
  heading = math.atan2(batch_x[1], batch_x[0])
  rot_mat = np.array([[np.cos(-heading), -np.sin(-heading)], [np.sin(-heading), np.cos(-heading)]])
  for pred_step in range(by.shape[0]):
    by[pred_step,:] = np.dot(rot_mat, by[pred_step,:])
  return by

def rotate_batch_to_local_frame(batch_y, batch_x):
  """
    inputs:
      batch_y: trainning velocities batch containing vx, vy on the global frame
      batch_x: trainning or prediction state batch containing x,y, vx, vy on the global frame
    """
  # rotate initial vx and vy in global frame to quary-agent local frame
  bx = batch_x.copy()  # rotated input batch values
  by = batch_y.copy()  # rotated output batch values
  for batch_idx in range(batch_x.shape[0]):
    for tbp_step in range(batch_x.shape[1]):
      heading = math.atan2(bx[batch_idx, tbp_step, 3], bx[batch_idx, tbp_step, 2])
      rot_mat = np.array([[np.cos(-heading), -np.sin(-heading)], [np.sin(-heading), np.cos(-heading)]])
      bx[batch_idx, tbp_step, 2:] = np.dot(rot_mat, bx[batch_idx, tbp_step, 2:])
      for pred_step in range(int(by.shape[2] / 2)):
        by[batch_idx, tbp_step, 2 * pred_step:2 * pred_step + 2] = np.dot(rot_mat, by[batch_idx, tbp_step,
                                                                                   2 * pred_step:2 * pred_step + 2])
  return bx , by

def plot_submap(ax, origin, grid, resolution, submap_size):
    ax.imshow(grid, cmap='gray_r')

                    
    pl.show()
    

def plot_grid(ax, center, grid, grid_resolution, submap_size):
    """
    Plot a binary occupancy grid in the axis ax.
    Specify center, resolution and size in order to align it with the coordinate frame.
    """
    plot_values = np.zeros([0, 2])
    for idx_x in range(grid.shape[0]):
        for idx_y in range(grid.shape[1]):
            grid_coordinate_local = np.array([idx_x, idx_y]) * grid_resolution - np.array(submap_size) / 2.0
            grid_coordinate_global = grid_coordinate_local + center
            if grid[idx_x, idx_y] == 1:
                plot_values = np.append(plot_values, np.array([[grid_coordinate_global[0], grid_coordinate_global[1]]]),
                                        axis=0)
    return ax.plot(plot_values[:, 0], plot_values[:, 1], marker='s', alpha=0.8, color='k', lw=0)

def plot_grid_roboat(ax, origin, grid, resolution, size):
    ax.imshow(grid,
               extent = (origin[0], origin[0] + size[0] * resolution,
                         origin[1], origin[1] + size[1] * resolution),
               cmap='gray_r')
    ax.set_facecolor('#262626')
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    ax.yaxis.set_label_coords(-0.08, .5)
    ax.xaxis.set_label_coords(0.5, -0.09)

    if resolution == 0.4378: # Amstel
      ax.set_xlim((-280,-230))
      ax.set_ylim((-25,25))
    
    if resolution == 0.22: # Bloemgracht
      ax.set_xlim((-20,50))
      ax.set_ylim((-40,30))

def create_map_from_png(file_name,resolution,map_size,map_center,data_path='../data',):
  # Create grid for SF model
  print("create_map_from_png")
  grid_map = {}
  grid_map['Resolution'] = resolution
  grid_map['Size'] = map_size  # map size in [m]
  grid_map['Map'] = np.zeros((int(grid_map['Size'][0] / grid_map['Resolution']),
                              int(grid_map['Size'][1] / grid_map['Resolution'])))

  #map_center = grid_map['Size'] / 2  # hack for my dataset
  H = np.genfromtxt(os.path.join(data_path, 'H.txt'),
                    delimiter='  ',
                    unpack=True).transpose()

  # Extract static obstacles
  obst_threshold = 200
  static_obst_img = cv.imread(os.path.join(data_path, file_name), 0)
  #static_obst_img = cv.imread(os.path.join(load_data_path, 'map.png'), 0)
  obstacles = np.zeros([0, 3])
  #grid = (static_obst_img/255*-1)+1
  #static_obst_img = np.transpose(static_obst_img)
  for xx in range(static_obst_img.shape[0]):
    for yy in range(static_obst_img.shape[1]):
      if static_obst_img[xx, yy] > obst_threshold:
        obstacles = np.append(obstacles,
                              np.dot(H, np.array([[xx], [yy], [1]])).transpose(),
                              axis=0)
  Hinv = np.linalg.inv(H)

  # Compute obstacles in 2D
  obstacles_2d = np.zeros([obstacles.shape[0], 2])
  obstacles_2d[:, 0] = obstacles[:, 0] / obstacles[:, 2]
  obstacles_2d[:, 1] = obstacles[:, 1] / obstacles[:, 2]

  # Get obstacle idx on map
  obst_idx = []
  for obst_ii in range(obstacles_2d.shape[0]):
    idx = np.dot(Hinv, np.array([[obstacles_2d[obst_ii, 0]], [obstacles_2d[obst_ii, 1]], [1]])).transpose()
    obst_idx.append((int(idx[0,0]),int(idx[0,1])))
    grid_map['Map'][obst_idx[-1]] = 1


  grid_map['Closest X'] = np.zeros_like(grid_map['Map']) + 10000
  grid_map['Closest Y'] = np.zeros_like(grid_map['Map']) + 10000
  """ """
  # Calculate closest obstacle
  obst_idx = np.array(obst_idx)
  for xx in range(int(grid_map['Size'][0] / grid_map['Resolution'])):
    for yy in range(int(grid_map['Size'][1] / grid_map['Resolution'])):
      delta_idx = obst_idx - np.array([xx, yy])
      distances = np.sqrt(np.sum(np.square(delta_idx), axis=1))
      closest_obj = np.argmin(distances)
      grid_map['Closest X'][xx, yy] = obst_idx[closest_obj][0]
      grid_map['Closest Y'][xx, yy] = obst_idx[closest_obj][1]

  grid_map['Closest X'] = grid_map['Closest X'] * grid_map['Resolution'] - map_center[0]
  grid_map['Closest Y'] = grid_map['Closest Y'] * grid_map['Resolution'] - map_center[1]

  np.save(os.path.join(data_path, 'map'), grid_map)

def idx_from_pos(x, y, center, res=0.1):
  idx_x = round((x + float(center[0])) / res)
  idx_y = round((y + float(center[1])) / res)

  # Projecting index on map if out of bounds
  idx_x = max(0, min(idx_x, -1 + center[0] * 2. / res))
  idx_y = max(0, min(idx_y, -1 + center[1] * 2. / res))
  return int(idx_x), int(idx_y)

def path_from_vel(initial_pos, pred_vel,v0=np.array([0,0]), dt=0.1, n_vx=1, n_vy=1):
  """
  Extract a path from vector of predicted velocities by applying Euler forward integration.
  """
  vel_x = pred_vel[:,0]*n_vx
  vel_y = pred_vel[:,1]*n_vy
  n_steps = len(vel_x)
  pos_vec = np.zeros([n_steps, 2])
  pos_vec[0, :] = initial_pos + v0*dt
  for ii in range(1, n_steps):
    pos_vec[ii, :] = pos_vec[ii-1,:] + dt * np.array([vel_x[ii], vel_y[ii]])

  return pos_vec

def to_image_frame(Hinv, loc):
  """
	Given H^-1 and world coordinates, returns (u, v) in image coordinates.
	"""
  locHomogenous = np.hstack((loc, np.ones((loc.shape[0], 1))))
  if locHomogenous.ndim > 1:
    loc_tr = np.transpose(locHomogenous)
    loc_tr = np.matmul(Hinv, loc_tr)  # to camera frame
    locXYZ = np.transpose(loc_tr / loc_tr[2])  # to pixels (from millimeters)
    return locXYZ[:, :2].astype(int)
  else:
    locHomogenous = np.dot(Hinv, locHomogenous)  # to camera frame
    locXYZ = locHomogenous / locHomogenous[2]  # to pixels (from millimeters)
    return locXYZ[:2].astype(int)

def to_pos_frame(Hinv, loc):
  """
	Given H^-1 and world coordinates, returns (u, v) in image coordinates.
	"""
  locHomogenous = np.hstack((loc, np.ones((loc.shape[0], 1))))
  if locHomogenous.ndim > 1:
    loc_tr = np.transpose(locHomogenous)
    loc_tr = np.matmul(Hinv, loc_tr)  # to camera frame
    locXYZ = np.transpose(loc_tr / loc_tr[2])  # to pixels (from millimeters)
    return locXYZ[:, :2]
  else:
    locHomogenous = np.dot(Hinv, locHomogenous)  # to camera frame
    locXYZ = locHomogenous / locHomogenous[2]  # to pixels (from millimeters)
    return locXYZ[:2]

def line_cv(im, ll, value, width):
  for tt in range(ll.shape[0] - 1):
    cv2.line(im, (int(ll[tt][1]), int(ll[tt][0])), (int(ll[tt + 1][1]), int(ll[tt + 1][0])), value, width)

def plotGrid(grid, ax, color='k', alpha=1.0):
  plot_values_ped = np.zeros([0, 2])
  plot_values_grid = np.zeros([0, 2])
  for idx_x in range(grid.shape[0]):
    for idx_y in range(grid.shape[1]):
      if grid[idx_x, idx_y] > 0.8 and grid[idx_x, idx_y] < 1.5:
        plot_values_grid = np.append(plot_values_grid, np.array([[idx_x, idx_y]]), axis=0)
      elif grid[idx_x, idx_y] > 1.5:
        plot_values_ped = np.append(plot_values_ped, np.array([[idx_x, idx_y]]), axis=0)
  ax.plot(plot_values_grid[:, 0], plot_values_grid[:, 1], marker='o', color=color, lw=0)
  ax.plot(plot_values_ped[:, 0], plot_values_ped[:, 1], marker='o', color='r', lw=0)
  ax.set_xlim([0, grid.shape[0]])
  ax.set_ylim([0, grid.shape[1]])


build/lib/data_utils/Support.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
from scipy import interpolate

class Trajectory():
  
  def __init__(self, time_vec = np.zeros([0]), pose_vec = np.zeros([0,3]), vel_vec = np.zeros([0,3]), goal = np.zeros([1,2])):
    self.time_vec = time_vec    # timesteps in [ns]
    self.pose_vec = pose_vec    # [x, y, heading]
    self.vel_vec = vel_vec      # [vx, vy, omega]
    self.goal = goal            # [x,y]
    self.pose_interp = None
    self.vel_interp = None
    self.other_agents_positions = []  # store indices of other agents trajectories for each time step of this trajectory (with which other positions does it need to be compared at a certain time)
    self.other_agents_velocities = []
    
  def __len__(self):
    if len(self.time_vec) == len(self.pose_vec) == len(self.vel_vec):
      return len(self.time_vec)
    else: 
      raise Exception('Member vectors are of different length. Cannot get unique length.')
    
  def addData(self, timestamp, pose, vel, goal):
    """
    Inserts the provided data at the right index in the trajectory.
    """
    assert np.all(goal == self.goal), "Goals have to be matching."
    # Add data at the right position. In the log the data might not be ordered properly.
    if self.time_vec.shape[0] > 0 and timestamp < self.time_vec[-1]:
      try: 
        idx_insert = np.where(timestamp > self.time_vec)[-1][-1] + 1
      except:
        idx_insert = 0  # In case the timestamp is smaller than the first element of time_vec
      self.time_vec = np.insert(self.time_vec, idx_insert, timestamp)
      self.pose_vec = np.insert(self.pose_vec, idx_insert, np.atleast_2d(pose), axis=0)
      self.vel_vec = np.insert(self.vel_vec, idx_insert, np.atleast_2d(vel), axis=0)
    else:
      self.time_vec = np.append(self.time_vec, timestamp)
      self.pose_vec = np.append(self.pose_vec, np.atleast_2d(pose), axis=0)
      self.vel_vec = np.append(self.vel_vec, np.atleast_2d(vel), axis=0)
    
    if self.time_vec.shape[0] > 1:
      self.updateInterpolators()
      
  def subsample(self, subsampling_factor):
    """
    Subsample the trajectory in order to reduce the overall number of samples and 
    increase the time difference between two samples.
    """
    if int(self.time_vec[0]*10) % 2 == 1:
      self.time_vec = self.time_vec[1:]
      self.pose_vec = self.pose_vec[1:, :]
      self.vel_vec = self.vel_vec[1:, :]
    self.time_vec = self.time_vec[0::subsampling_factor]
    self.pose_vec = self.pose_vec[0::subsampling_factor, :]
    self.vel_vec = self.vel_vec[0::subsampling_factor, :]

    
    #self.updateInterpolators()
        
  def updateInterpolators(self):
    self.pose_interp = interpolate.interp1d(self.time_vec, self.pose_vec, kind='linear', axis=0)  # TODO: correct heading interpolation
    self.vel_interp = interpolate.interp1d(self.time_vec, self.vel_vec, kind='linear', axis=0)
      
  def contains(self, query_time):
    return query_time >= self.time_vec[0] and query_time <= self.time_vec[-1]
   
  def getPoseAtTime(self, query_time):
    pose = self.pose_interp(query_time)
    return pose
   
  def getVelAtTime(self, query_time):
    return self.vel_interp(query_time)
  
  def getDataAtTime(self, query_time):
    pose_interpol = self.getPoseAtTime(query_time)
    vel_interpol = self.getVelAtTime(query_time)
    return pose_interpol, vel_interpol
  
  def getMinTime(self):
    """
    Return the starting time of the trajectory.
    """
    return self.time_vec[0]
  
  def getMaxTime(self):
    """
    Return the end time of the trajectory.
    """
    return self.time_vec[-1]  
  
  def getDuration(self):
    return self.getMaxTime() - self.getMinTime()
  
  def smoothenTrajectory(self, dt=0.3):
    """
    Cubic interpolation with provided values in order to smoothen the trajectory and obtain
    a sequence with the specified dt.
    """
    x_interpolator = interpolate.interp1d(self.time_vec, self.pose_vec[:,0], kind='cubic', axis=0)
    y_interpolator = interpolate.interp1d(self.time_vec, self.pose_vec[:,1], kind='cubic', axis=0)
    vx_interpolator = interpolate.interp1d(self.time_vec, self.vel_vec[:,0], kind='cubic', axis=0)
    vy_interpolator = interpolate.interp1d(self.time_vec, self.vel_vec[:,1], kind='cubic', axis=0)
    n_elem = int((self.getMaxTime() - self.getMinTime()) / dt)
    new_time_vec = np.linspace(self.getMinTime(), self.getMinTime() + (n_elem - 1) * dt, n_elem)
    
    new_time_vec = np.zeros([n_elem])
    new_pose_vec = np.zeros([n_elem, 3])
    new_vel_vec = np.zeros([n_elem, 3])
    
    for ii in range(n_elem):
      t = self.getMinTime() + ii*dt
      new_time_vec[ii] = t
      new_pose_vec[ii, 0] = x_interpolator(t)
      new_pose_vec[ii, 1] = y_interpolator(t)
      new_vel_vec[ii, 0] = vx_interpolator(t)
      new_vel_vec[ii, 1] = vy_interpolator(t)
    
    self.goal = new_pose_vec[-1,:2]
    self.time_vec = new_time_vec
    self.pose_vec = new_pose_vec
    self.vel_vec = new_vel_vec
    

  def __getstate__(self):
    d = dict(self.__dict__)
    del d['pose_interp']
    del d['vel_interp']
    return d
    
build/lib/data_utils/Trajectory.py
eliatrevisan/IA-MPPI-LBM
import numpy as np

class Config:
build/lib/data_utils/config.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import os
import random
import sys
import math
import cv2
import pickle as pkl
import time
if sys.version_info[0] < 3:
  sys.path.append('../src/data_utils')
  import Support as sup
else:
	import src.data_utils.Support as sup
import matplotlib.pyplot as pl
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter
from matplotlib.patches import Ellipse

def plot_scenario(enc_seq, dec_seq,traj, predictions, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	metadata = dict(title='Movie Test', artist='Matplotlib')
	writer = FFMpegWriter(fps=10, metadata=metadata)

	fig = pl.figure("Trajectory Predictions")
	ax_in = pl.subplot()

	ax_in.clear()

	#x_lim = [-global_grid.map_size[0]/2 ,global_grid.map_size[0]/2]
	#y_lim = [-global_grid.map_size[1]/2,global_grid.map_size[1]/2]

	x_lim = [-5 ,5]
	y_lim = [-5,5]

	ax_in.clear()
	pedestrian_trajectory = np.zeros((len(traj),2))
	for t in range(len(traj)):
		pedestrian_trajectory[t,:] = traj[t]["pedestrian_state"]["position"]
	with writer.saving(fig, args.data_path + "/movie.mp4",100):
		for t in range(len(enc_seq)):
				ax_in.clear()
				# plot scenario grid
				#sup.plot_grid(ax_in, np.array([0, 0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
				ax_in.plot(pedestrian_trajectory[t:, 0], pedestrian_trajectory[t:, 1], color='red', label='Agent Real Traj')
				ax_in.plot(enc_seq[t][:,0], enc_seq[t][:,1], color='g', marker='o', label='Robot Prev Real Traj')
				ax_in.plot(enc_seq[t][:,4], enc_seq[t][:,5], color='red', marker='o', label='Agent Prev Traj')
				ax_in.plot(dec_seq[t][:, 0], dec_seq[t][:, 1], color='blue',marker='o', label='Agent Real')
				ax_in.plot(predictions[t][:, 0], predictions[t][:, 1], color='blue', marker='x', label='Agent Pred Traj')


				ax_in.set_xlim(x_lim)
				ax_in.set_ylim(y_lim)

				ax_in.legend()
				fig.canvas.draw()
				pl.show(block=False)
				#time.sleep(t)
				writer.grab_frame()

def plot_scenario_vel_tbp(enc_seq, dec_seq,traj, predictions, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	metadata = dict(title='Movie Test', artist='Matplotlib')
	writer = FFMpegWriter(fps=10, metadata=metadata)

	fig = pl.figure("Trajectory Predictions")
	ax_in = pl.subplot()

	ax_in.clear()

	#x_lim = [-global_grid.map_size[0]/2 ,global_grid.map_size[0]/2]
	#y_lim = [-global_grid.map_size[1]/2,global_grid.map_size[1]/2]

	x_lim = [-5 ,5]
	y_lim = [-5,5]

	ax_in.clear()
	pedestrian_trajectory = np.zeros((len(traj),2))
	for t in range(len(traj)):
		pedestrian_trajectory[t,:] = traj[t]["pedestrian_state"]["position"]
	with writer.saving(fig, args.data_path + "/movie.mp4",100):
		for t in range(len(enc_seq)):
				ax_in.clear()
				# plot scenario grid
				#sup.plot_grid(ax_in, np.array([0, 0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
				ax_in.plot(pedestrian_trajectory[t:, 0], pedestrian_trajectory[t:, 1], color='red', label='Agent Real Traj')
				ax_in.plot(enc_seq[t][0,0::args.input_dim], enc_seq[t][0,1::args.input_dim], color='g', marker='o', label='Robot Prev Real Traj')
				ax_in.plot(enc_seq[t][0,4::args.input_dim], enc_seq[t][0,5::args.input_dim], color='red', marker='o', label='Agent Prev Traj')
				ax_in.plot(pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 0], pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 1], color='blue',marker='o', label='Agent Real')
				x0 = pedestrian_trajectory[t+args.prev_horizon,0]
				y0 = pedestrian_trajectory[t+args.prev_horizon,1]
				pred_vel = np.zeros((args.prediction_horizon,args.output_dim))
				for i in range(args.prediction_horizon):
					idx = i * args.output_dim
					idy = idx + 1
					vx = predictions[t][0,idx]
					vy = predictions[t][0,idy]
					pred_vel[i, :] = [vx, vy]

				traj_pred = sup.path_from_vel(initial_pos=np.array([x0 , y0]),pred_vel=pred_vel, dt=args.dt)
				ax_in.plot(traj_pred[:, 0], traj_pred[:, 1], color='blue', marker='x', label='Agent Pred Traj')
				ax_in.set_xlim(x_lim)
				ax_in.set_ylim(y_lim)

				ax_in.legend()
				fig.canvas.draw()
				pl.show(block=False)
				#time.sleep(t)
				writer.grab_frame()

#def plot_scenario_vel(enc_seq, dec_seq,traj, predictions, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	"""	metadata = dict(title='Movie Test', artist='Matplotlib')
	writer = FFMpegWriter(fps=1, metadata=metadata)

	fig = pl.figure("Trajectory Predictions")
	ax_in = pl.subplot()

	ax_in.clear()

	#x_lim = [-global_grid.map_size[0]/2 ,global_grid.map_size[0]/2]
	#y_lim = [-global_grid.map_size[1]/2,global_grid.map_size[1]/2]

	x_lim = [-5 ,5]
	y_lim = [-5,5]

	ax_in.clear()
	pedestrian_trajectory = np.zeros((len(traj),2))
	for t in range(len(traj)):
		pedestrian_trajectory[t,:] = traj[t]["pedestrian_state"]["position"]
	with writer.saving(fig, args.data_path + "/movie.mp4",100):
		for t in range(len(enc_seq)):
				ax_in.clear()
				# plot scenario grid
				#sup.plot_grid(ax_in, np.array([0, 0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
				ax_in.plot(pedestrian_trajectory[t:, 0], pedestrian_trajectory[t:, 1], color='red', label='Agent Real Traj')
				ax_in.plot(enc_seq[t][:,0], enc_seq[t][:,1], color='g', marker='o', label='Robot Prev Real Traj')
				ax_in.plot(enc_seq[t][:,4], enc_seq[t][:,5], color='red', marker='o', label='Agent Prev Traj')
				ax_in.plot(pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 0], pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 1], color='blue',marker='o', label='Agent Real')
				x0 = pedestrian_trajectory[t+args.prev_horizon,0]
				y0 = pedestrian_trajectory[t+args.prev_horizon,1]
				pred_vel = predictions[t]
				traj_pred = sup.path_from_vel(initial_pos=np.array([x0 , y0]),pred_vel=pred_vel, dt=0.1)
				ax_in.plot(traj_pred[:, 0], traj_pred[:, 1], color='k', marker='x', label='Agent Pred Traj')
				ax_in.set_xlim(x_lim)
				ax_in.set_ylim(y_lim)

				ax_in.legend()
				fig.canvas.draw()
				pl.show(block=False)
				#time.sleep(t)
				writer.grab_frame()
"""

def plot_scenario_vel_OpenCV_simdata(trajectories, batch_target, all_predictions, args, exp_num=0,display = True):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
	"""
	homography_file = os.path.join(args.data_path + args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')
	scenario = args.scenario
	video_file = args.model_path +"/"  + args.model_name+ "_" + str(exp_num) + '.avi'
	img_file = args.data_path + args.scenario +'/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = ~cv2.imread(img_file)
		img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
		height, width, channels = img.shape

		scale_percent = 400  # percent of original size
		width = int(img.shape[1] * scale_percent / 100)
		height = int(img.shape[0] * scale_percent / 100)
		dim = (width, height)
		# resize image
		resized_img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		out = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width, height))

		# Adding legend
		font = cv2.FONT_HERSHEY_SIMPLEX
		cv2.line(resized_img, (width - 150, 30), (width - 130, 30), (0, 255, 0), 4)
		cv2.putText(resized_img, "Query Agent", (width - 120, 30), font, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 50), (width - 130, 50), (255, 0, 0), 4)
		cv2.putText(resized_img, "Other Agents", (width - 120, 50), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 70), (width - 130, 70), (255, 255, 0), 4)
		cv2.putText(resized_img, "Query Agent traj", (width - 120, 70), font, 0.5, (255, 255, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 90), (width - 130, 90), (128, 128, 128), 4)
		cv2.putText(resized_img, "Pred traj 1", (width - 120, 90), font, 0.5, (128, 128, 128), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 110), (width - 130, 110), (0, 128, 255), 4)
		cv2.putText(resized_img, "Pred traj 2", (width - 120, 110), font, 0.5, (0, 128, 255), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 130), (width - 130, 130), (255, 0, 127), 4)
		cv2.putText(resized_img, "Pred traj 3", (width - 120, 130), font, 0.5, (255, 0, 127), 2, cv2.LINE_AA)

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]
			ped_vel = batch_target[exp_idx][0]

			for t in range(len(predictions)):
				pedestrian_trajectory = traj.pose_vec[:,:2]
				overlay_ellipses = resized_img.copy()

				# Plot Query Agent Positions
				query_agent_pos = traj.pose_vec[t,:2]
				query_agent = np.expand_dims(np.array((query_agent_pos[0], query_agent_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, query_agent)*int(scale_percent / 100)
				sigma_x = 0.2
				sigma_y = 0.2
				axis = (
					int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100), int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
				center = (obsv_XY[0, 1], obsv_XY[0, 0])
				cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (0, 255, 0), -1)

				# Plot other pedestrians positions
				agent_pos = traj.other_agents_positions[t]
				for agent_id in range(agent_pos.shape[0]):
					ped_pos = np.expand_dims(np.array((agent_pos[agent_id, 0], agent_pos[agent_id, 1])), axis=0).astype(float)
					obsv_XY = sup.to_image_frame(Hinv, ped_pos)*int(scale_percent / 100)
					axis = (
						int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100),
						int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (255, 0, 0), 1)

				# Query-agent real trajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:, :])*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 0, 0), 1)  # bgr convention

				# Query-agent predicted real trajectory from positions
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:t+args.prediction_horizon, :]) * int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 1)  # bgr convention

				#Query-agent Predicted Real TRajectory
				x0 = traj.pose_vec[t+1+args.prev_horizon,0]
				y0 = traj.pose_vec[t+1+args.prev_horizon,1]
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.truncated_backprop_length):
					if args.normalize_data:
						pedestrian_velocity[t_idx, 0] = ped_vel[t, t_idx * 2] / args.sx_vel + args.min_vel_x
						pedestrian_velocity[t_idx, 1] = ped_vel[t, t_idx * 2+1] / args.sy_vel + args.min_vel_y
					else:
						pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_traj = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pedestrian_velocity, dt=args.dt)
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_traj)*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 1)  # bgr convention

				# Query-agent real trajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:t+args.prediction_horizon, :])*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (0, 0, 0), 1)  # bgr convention

				# Plot Predictions
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				for sample_id in range(len(predictions[t])):
					prediction_sample = predictions[t][sample_id]
					for mix_id in range(args.n_mixtures):
						for pred_step in range(args.prediction_horizon):
							idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
							if args.normalize_data:
								pred_vel[pred_step, 0] = prediction_sample[0][idx] / args.sx_vel + args.min_vel_x
								pred_vel[pred_step, 1] = prediction_sample[0][idx+args.n_mixtures] / args.sy_vel + args.min_vel_y
							else:
								pred_vel[pred_step, :] = np.array((prediction_sample[0][idx],prediction_sample[0][idx+args.n_mixtures]
							                                   ))
							if args.output_pred_state_dim > args.output_dim:
								sigmax[pred_step, :] = prediction_sample[0][idx + 2*args.n_mixtures]
								sigmay[pred_step, :] = prediction_sample[0][idx + 3*args.n_mixtures]
						try:
							if args.predict_positions:
								traj_pred = pred_vel
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						except:
							traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
						sup.line_cv(overlay_ellipses, obsv_XY, colors[mix_id], 1)  # bgr convention

						if args.output_pred_state_dim > args.output_dim:
							sigma_x = np.square(sigmax[0]) * (args.dt ** 2)
							sigma_y = np.square(sigmay[0]) * (args.dt ** 2)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
							axis = (
								int(np.sqrt(sigma_x) / args.submap_resolution * scale_percent / 100), int(np.sqrt(sigma_y) / args.submap_resolution * scale_percent / 100))
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							try:
								cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], 1)
							except:
								print("failed to plot ellipse")

							for pred_step in range(1,args.prediction_horizon):
								sigma_x += np.square(sigmax[pred_step]) * (args.dt ** 2)
								sigma_y += np.square(sigmay[pred_step]) * (args.dt ** 2)
								axis = (
									int(np.sqrt(sigma_x) / args.submap_resolution * scale_percent / 100), int(np.sqrt(sigma_y) / args.submap_resolution * scale_percent / 100))
								center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
								try:
									cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], 1)
								except:
									print("failed to plot ellipse")

				#image_new = cv2.addWeighted(resized_img, 0.2, overlay_ellipses, 0.8, 0)
				if display:
					cv2.imshow("image", overlay_ellipses)
					cv2.waitKey(100)
				out.write(overlay_ellipses)

		# When everything done, release the video capture and video write objects
		out.release()

	print("Done recording...")

def plot_local_scenario_vel_OpenCV(trajectories, batch_target, all_predictions, args, exp_num=0,display = True):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	homography_file = os.path.join(args.data_path + args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')
	scenario = args.scenario
	video_file = args.model_path +"/"  + args.model_name+ "_" + str(exp_num) + '_local.avi'
	img_file = args.data_path + args.scenario +'/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = ~cv2.imread(img_file)

		height, width, channels = img.shape

		scale_percent = 200  # percent of original size
		width = int(img.shape[1] * scale_percent / 100)
		height = int(img.shape[0] * scale_percent / 100)
		dim = (width, height)
		# resize image
		resized_img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		out = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (600, 600))

		# Adding legend
		font = cv2.FONT_HERSHEY_SIMPLEX
		cv2.line(resized_img, (width - 150, 30), (width - 130, 30), (0, 255, 0), 4)
		cv2.putText(resized_img, "Query Agent", (width - 120, 30), font, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 50), (width - 130, 50), (255, 0, 0), 4)
		cv2.putText(resized_img, "Other Agents", (width - 120, 50), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 70), (width - 130, 70), (255, 255, 0), 4)
		cv2.putText(resized_img, "Query Agent traj", (width - 120, 70), font, 0.5, (255, 255, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 90), (width - 130, 90), (128, 128, 128), 4)
		cv2.putText(resized_img, "Pred traj 1", (width - 120, 90), font, 0.5, (128, 128, 128), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 110), (width - 130, 110), (0, 128, 255), 4)
		cv2.putText(resized_img, "Pred traj 2", (width - 120, 110), font, 0.5, (0, 128, 255), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 130), (width - 130, 130), (255, 0, 127), 4)
		cv2.putText(resized_img, "Pred traj 3", (width - 120, 130), font, 0.5, (255, 0, 127), 2, cv2.LINE_AA)

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]
			ped_vel = batch_target[exp_idx][0]

			for t in range(len(predictions)):
				pedestrian_trajectory = traj.pose_vec[:,:2]
				overlay_ellipses = resized_img.copy()

				# Plot Query Agent Positions
				query_agent_pos = traj.pose_vec[t,:2]
				query_agent = np.expand_dims(np.array((query_agent_pos[0], query_agent_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, query_agent)*int(scale_percent / 100)
				sigma_x = 0.4
				sigma_y = 0.4
				axis = (
					int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100), int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
				agent_center = (obsv_XY[0, 1], obsv_XY[0, 0])
				cv2.ellipse(overlay_ellipses, agent_center, axis, 0, 0, 360, (0, 255, 0), -1)

				# Plot ped real position
				agent_pos = traj.other_agents_positions[t]
				for agent_id in range(agent_pos.shape[0]):
					ped_pos = np.expand_dims(np.array((agent_pos[agent_id, 0], agent_pos[agent_id, 1])), axis=0).astype(float)
					obsv_XY = sup.to_image_frame(Hinv, ped_pos)*int(scale_percent / 100)
					axis = (
						int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100),
						int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (255, 0, 0), -1)

				# Query-agent real trajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:, :])*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 0, 0), 1)  # bgr convention

				#Query-agent Predicted Real TRajectory
				x0 = traj.pose_vec[t+1+args.prev_horizon,0]
				y0 = traj.pose_vec[t+1+args.prev_horizon,1]
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.truncated_backprop_length):
					if args.normalize_data:
						pedestrian_velocity[t_idx, 0] = ped_vel[t, t_idx * 2] / args.sx_vel + args.min_vel_x
						pedestrian_velocity[t_idx, 1] = ped_vel[t, t_idx * 2+1] / args.sy_vel + args.min_vel_y
					else:
						pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_trajectory = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pedestrian_velocity, dt=args.dt)
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory)*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 1)  # bgr convention


				# Plot Predictions
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				for sample_id in range(len(predictions[t])):
					prediction_sample = predictions[t][sample_id]
					for mix_id in range(args.n_mixtures):
						for pred_step in range(args.prediction_horizon):
							idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
							if args.normalize_data:
								pred_vel[pred_step, 0] = prediction_sample[0][idx] / args.sx_vel + args.min_vel_x
								pred_vel[pred_step, 1] = prediction_sample[0][idx+args.n_mixtures] / args.sy_vel + args.min_vel_y
							else:
								pred_vel[pred_step, :] = np.array((prediction_sample[0][idx],prediction_sample[0][idx+args.n_mixtures]))
							if args.output_pred_state_dim > args.output_dim:
								sigmax[pred_step, :] = prediction_sample[0][idx + 2*args.n_mixtures]
								sigmay[pred_step, :] = prediction_sample[0][idx + 3*args.n_mixtures]
						try:
							if args.predict_positions:
								traj_pred = pred_vel
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						except:
							traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
						sup.line_cv(overlay_ellipses, obsv_XY, colors[mix_id], 1)  # bgr convention

						if args.output_pred_state_dim > args.output_dim:
							sigma_x = np.square(sigmax[0]) * (args.dt ** 2)
							sigma_y = np.square(sigmay[0]) * (args.dt ** 2)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
							axis = (
								int(np.sqrt(sigma_x) / args.submap_resolution * scale_percent / 100), int(np.sqrt(sigma_y) / args.submap_resolution * scale_percent / 100))
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], 1)

							for pred_step in range(1,args.prediction_horizon):
								sigma_x += np.square(sigmax[pred_step]) * (args.dt ** 2)
								sigma_y += np.square(sigmay[pred_step]) * (args.dt ** 2)
								axis = (
									int(np.sqrt(sigma_x) / args.submap_resolution * scale_percent / 100), int(np.sqrt(sigma_y) / args.submap_resolution * scale_percent / 100))
								center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
								try:
									cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], 1)
								except:
									print("failed to plot ellipse")

				crop_img = cv2.getRectSubPix(overlay_ellipses,(600,600),agent_center)
				if display:
					cv2.imshow("image", crop_img)
					cv2.waitKey(100)
				out.write(crop_img)

		# When everything done, release the video capture and video write objects
		out.release()

	print("Done recording...")

def plot_scenario_vel_OpenCV(trajectories, batch_target, all_predictions, args, exp_num=0,display = False):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	homography_file = os.path.join(args.data_path+ args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')
	scenario = args.scenario
	video_file = args.model_path +"/"  + args.model_name+ "_" + str(exp_num) + '.avi'
	map_resolution = 1 / Hinv[0, 0]
	print("Recording to: " + video_file)

	img_file = args.data_path+ args.scenario + '/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = cv2.imread(img_file)

		#height, width, channels = img.shape
		width = 600
		height = 600
		dim = (width, height)

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		out = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width, height))

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]
			ped_vel = batch_target[exp_idx][0]

			pedestrian_trajectory = np.zeros((len(traj), 2))
			for t in range(len(traj)):
				pedestrian_trajectory[t, :] = traj[t]["pedestrian_state"]["position"]

			for t in range(len(predictions)):
				overlay_ellipses = img.copy()

				# Plot Robot Positions
				robot_pos = traj[t]["robot_state"]
				robot = np.expand_dims(np.array((robot_pos[0], robot_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, robot)
				sigma_x = 0.4
				sigma_y = 0.4
				axis_agent = (
					int(np.sqrt(sigma_x) / map_resolution), int(np.sqrt(sigma_y) / map_resolution))
				center = (obsv_XY[0, 0], obsv_XY[0, 1])
				cv2.ellipse(overlay_ellipses, center, axis_agent, 0, 0, 360, (0, 0, 0), -1)

				# Plot ped real position
				agent_pos = traj[t]["pedestrian_state"]["position"]
				ped_pos = np.expand_dims(np.array((agent_pos[0], agent_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, ped_pos)

				center = (obsv_XY[0, 0], obsv_XY[0, 1])
				cv2.ellipse(overlay_ellipses, center, axis_agent, 0, 0, 360, (255, 0, 0), -1)

				# Pedestrian real trajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:, :])
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 0, 0), 3)  # bgr convention

				#Pedestrian Predicted Real TRajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t + args.prev_horizon:t + args.prev_horizon + args.prediction_horizon, :])
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 3)  # bgr convention

				#Pedestrian Predicted Real TRajectory
				x0 = pedestrian_trajectory[t, 0]
				y0 = pedestrian_trajectory[t, 1]
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.truncated_backprop_length):
					pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_traj = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pedestrian_velocity, dt=args.dt)
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_traj)
				sup.line_cv(overlay_ellipses, obsv_XY, (0, 0, 255), 3)  # bgr convention

				# Plot Predictions
				x0 = pedestrian_trajectory[t, 0]
				y0 = pedestrian_trajectory[t, 1]
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				for sample_id in range(len(predictions[t])):
					prediction_sample = predictions[t][sample_id]
					for mix_id in range(args.n_mixtures):
						for pred_step in range(args.prediction_horizon):
							idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
							pred_vel[pred_step, :] = np.array((prediction_sample[0][idx],prediction_sample[0][idx+args.n_mixtures]
							                                   ))
							if args.output_pred_state_dim > args.output_dim:
								sigmax[pred_step, :] = prediction_sample[0][idx + 2*args.n_mixtures]
								sigmay[pred_step, :] = prediction_sample[0][idx + 3*args.n_mixtures]
						traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						obsv_XY = sup.to_image_frame(Hinv, traj_pred)
						sup.line_cv(overlay_ellipses, obsv_XY, colors[mix_id], 3)  # bgr convention

						if args.output_pred_state_dim > args.output_dim:
							sigma_x = np.square(np.exp(sigmax[0])) * (args.dt ** 2)
							sigma_y = np.square(np.exp(sigmay[0])) * (args.dt ** 2)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)
							axis = (
								int(np.sqrt(sigma_x) / map_resolution), int(np.sqrt(sigma_y) / map_resolution))
							#axis = (int(axis[0] + axis_agent[0]),int(axis[1] + axis_agent[1]))
							center = (obsv_XY[0, 0], obsv_XY[0, 1])
							cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], -1)

							for pred_step in range(1,args.prediction_horizon):
								sigma_x += np.square(np.exp(sigmax[pred_step])) * (args.dt ** 2)
								sigma_y += np.square(np.exp(sigmay[pred_step])) * (args.dt ** 2)
								axis = (
									int(np.sqrt(sigma_x) / map_resolution), int(np.sqrt(sigma_y) / map_resolution))
								#axis = (int(axis[0] + axis_agent[0]), int(axis[1] + axis_agent[1]))
								center = (obsv_XY[pred_step, 0], obsv_XY[pred_step, 1])
								cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], -1)

				image_new = cv2.addWeighted(img, 0.2, overlay_ellipses, 0.8, 0)

				# resize image
				resized_img = cv2.resize(image_new, dim, interpolation=cv2.INTER_AREA)
				if display:
					cv2.imshow("image", resized_img)
					cv2.waitKey(100)
				out.write(resized_img)

	# When everything done, release the video capture and video write objects
	out.release()

	print("Done recording...")

def plot_scenario_vel(trajectories, all_predictions, args, grid = None):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	metadata = dict(title='Movie Test', artist='Matplotlib')
	writer = FFMpegWriter(fps=2, metadata=metadata)

	fig = pl.figure("Trajectory Predictions")
	ax_in = pl.subplot()

	ax_in.clear()

	x_lim = [-5 ,5]
	y_lim = [-5,5]

	ax_in.clear()

	with writer.saving(fig, args.data_path + "/movie.mp4",100):

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]

			pedestrian_trajectory = np.zeros((len(traj), 2))
			for t in range(len(traj)):
				pedestrian_trajectory[t, :] = traj[t]["pedestrian_state"]["position"]

			for t in range(len(predictions)):
				ax_in.clear()
				# plot scenario grid
				#sup.plot_grid(ax_in, np.array([0, 0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
				ax_in.plot(pedestrian_trajectory[t:, 0], pedestrian_trajectory[t:, 1], color='red', label='Agent Real Traj')
				robot_pos = traj[t]["robot_state"]
				ax_in.plot(robot_pos[0], robot_pos[1],
				           color='g', marker='o', label='Robot')
				agent_pos = traj[t]["pedestrian_state"]["position"]
				ax_in.plot(agent_pos[0], agent_pos[1],
				           color='red', marker='o', label='Agent')

				if grid == None:
					gridmap = np.zeros((int((x_lim[1] - x_lim[0]) / map_resolution),
					                   int((y_lim[1] - y_lim[0]) / map_resolution)))

				sup.add_other_agents_to_grid(gridmap,robot_pos,args)

				sup.plot_grid(ax_in, np.array([0.0, 0.0]), gridmap, map_resolution,
				              np.array((10,10)))

				ax_in.plot(pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 0], pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 1], color='blue',marker='o', label='Agent Real')

				x0 = pedestrian_trajectory[t+args.prev_horizon,0]
				y0 = pedestrian_trajectory[t+args.prev_horizon,1]
				pred_vel = np.zeros((args.prediction_horizon,args.output_dim))
				for i in range(args.prediction_horizon):
					pred_vel[i,:] = predictions[t][i*args.output_pred_state_dim:i*args.output_pred_state_dim+2]
				traj_pred = sup.path_from_vel(initial_pos=np.array([x0 , y0]),pred_vel=pred_vel, dt=args.dt)
				ax_in.plot(traj_pred[:, 0], traj_pred[:, 1], color='k', marker='x', label='Agent Pred Traj')

				sigma_x = np.square(predictions[t][2]) * (args.dt **2)
				sigma_y = np.square(predictions[t][3]) * (args.dt **2)

				e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
				             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
				e1.set_alpha(0.5)
				ax_in.add_patch(e1)

				for pred_step in range(args.prediction_horizon):
					for mix_idx in range(args.n_mixtures):
						sigma_x += np.square(predictions[t][args.n_mixtures * pred_step * args.output_pred_state_dim +
						                                          mix_idx * args.output_pred_state_dim +2]) * (args.dt **2)
						sigma_y += np.square(predictions[t][args.n_mixtures * pred_step * args.output_pred_state_dim +
						                                          mix_idx * args.output_pred_state_dim +3])* (args.dt **2)

						e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
						             height=np.sqrt(sigma_y) / 2,
						             angle=0 / np.pi * 180)
						ax_in.add_patch(e1)

				ax_in.set_xlim(x_lim)
				ax_in.set_ylim(y_lim)

				ax_in.legend()
				fig.canvas.draw()
				pl.show(block=False)
				time.sleep(0.1)
				writer.grab_frame()

def plot_batch(batch_x, batch_grid, batch_ped_grid, batch_y, other_agents_pos,args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""

	fig = pl.figure("Batch Data")
	ax_in = pl.subplot()

	ax_in.clear()

    #x_lim = [-args.map_size[0]*args.resolution/2 ,args.map_size[0]*args.resolution/2]
	#y_lim = [-args.map_size[1]*args.resolution/2,args.map_size[1]*args.resolution/2]
	x_lim = [-args.grid_width*args.submap_resolution/2 ,args.grid_width*args.submap_resolution/2]
	y_lim = [-args.grid_height*args.submap_resolution/2,args.grid_height*args.submap_resolution/2]
	x_lim = [-5,5]
	y_lim = [-5, 5]
	ax_in.clear()

	for batch_id in range(batch_x.shape[0]):
		for t_step in range(batch_x.shape[1]):
			ax_in.clear()
			# plot scenario grid
			sup.plot_grid(ax_in, np.array([0, 0]), batch_grid[batch_id,t_step,:,:],
			              args.submap_resolution, np.array((args.grid_width,args.grid_height))*args.submap_resolution)

			robot_pos = other_agents_pos[batch_id,t_step,:]
			agent_pos = batch_x[batch_id, t_step, :2]
			ax_in.plot(robot_pos[0]-agent_pos[0], robot_pos[1]-agent_pos[1],
				           color='g', marker='o', label='Robot')

			ax_in.plot(0, 0,
				           color='red', marker='o', label='Agent')

			ax_in.set_xlim(x_lim)
			ax_in.set_ylim(y_lim)

			ax_in.legend()
			fig.canvas.draw()
			pl.show(block=False)

def plot_batch_OpenCV(step,batch_x, batch_grid, batch_ped_grid, batch_y, other_agents_pos,_model_prediction, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	homography_file = os.path.join(args.data_path+ args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')

	video_file = args.model_path +"/"  + args.model_name+ "_" + str(step) + '_training.avi'
	img_file = args.data_path+ args.scenario + '/map.png'
	font = cv2.FONT_HERSHEY_SIMPLEX
	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = cv2.imread(img_file)

		height, width, channels = img.shape

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		recorder = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width, height))

		# Adding legend
		cv2.putText(img, str(step), (30, 30), font, 1, (0, 0, 0), 2, cv2.LINE_AA)
		cv2.line(img, (850, 30), (900, 30), (0, 0, 0), 4)
		cv2.putText(img, "Robot Pos", (910, 30), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)
		cv2.line(img, (850, 50), (900, 50), (255, 0, 0), 4)
		cv2.putText(img, "Pedestrian", (910, 50), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
		cv2.line(img, (850, 70), (900, 70), (255, 255, 0), 4)
		cv2.putText(img, "Pedestrian traj", (910, 70), font, 0.5, (255, 255, 0), 2, cv2.LINE_AA)
		cv2.line(img, (850, 90), (900, 90), (128,128,128), 4)
		cv2.putText(img, "Pred traj 1", (910, 90), font, 0.5, (128,128,128), 2, cv2.LINE_AA)
		cv2.line(img, (850, 110), (900, 110), (0,128,255), 4)
		cv2.putText(img, "Pred traj 2", (910, 110), font, 0.5, (0,128,255), 2, cv2.LINE_AA)
		cv2.line(img, (850, 130), (900, 130), (255,0,127), 4)
		cv2.putText(img, "Pred traj 3", (910, 130), font, 0.5, (255,0,127), 2, cv2.LINE_AA)

		for batch_idx in range(batch_x.shape[0]):
			if "keras" in args.model_name:
				predictions = _model_prediction[batch_idx]
			else:
				predictions = np.transpose(_model_prediction,(1,0,2))[batch_idx]
			ped_vel = batch_y[batch_idx]

			for t in range(predictions.shape[0]):
				overlay_ellipses = img.copy()

				# Plot Robot Positions
				robot_pos = batch_x[batch_idx,t,:2]
				robot = np.expand_dims(np.array((robot_pos[0], robot_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, robot)
				sigma_x = 0.4
				sigma_y = 0.4
				axis = (
					int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
				center = (obsv_XY[0, 1], obsv_XY[0, 0])
				cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (0, 0, 0), -1)

				# Plot ped real position
				agent_pos = other_agents_pos[batch_idx][t]
				for agent_id in range(agent_pos.shape[0]):
					ped_pos = np.expand_dims(np.array((agent_pos[agent_id,0], agent_pos[agent_id,1])), axis=0).astype(float)
					obsv_XY = sup.to_image_frame(Hinv, ped_pos)
					axis = (
						int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (255, 0, 0), -1)

				#Pedestrian Predicted Real TRajectory
				x0 = batch_x[batch_idx, t, 0]
				y0 = batch_x[batch_idx, t, 1]
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.truncated_backprop_length):
					pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_trajectory = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pedestrian_velocity, dt=args.dt)
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 3)  # bgr convention

				# Plot Predictions
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				# orange grey and purple
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				for mix_id in range(args.n_mixtures):
					for pred_step in range(args.prediction_horizon):
						idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
						pred_vel[pred_step, :] = np.array((predictions[t][idx],predictions[t][idx+args.n_mixtures]
						                                   ))
						if args.output_pred_state_dim > args.output_dim:
							sigmax[pred_step, :] = predictions[t][idx + 2*args.n_mixtures]
							sigmay[pred_step, :] = predictions[t][idx + 3*args.n_mixtures]
					traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
					obsv_XY = sup.to_image_frame(Hinv, traj_pred)
					sup.line_cv(overlay_ellipses, obsv_XY, colors[mix_id], 3)  # bgr convention

					if args.output_pred_state_dim > args.output_dim:
						#Propagate uncertainty and plot
						sigma_x = np.square(sigmax[0]) * (args.dt ** 2)
						sigma_y = np.square(sigmay[0]) * (args.dt ** 2)

						obsv_XY = sup.to_image_frame(Hinv, traj_pred)
						axis = (
							int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
						center = (obsv_XY[0, 1], obsv_XY[0, 0])
						cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], -1)

						for pred_step in range(1,args.prediction_horizon):
							sigma_x += np.square(sigmax[pred_step]) * (args.dt ** 2)
							sigma_y += np.square(sigmay[pred_step]) * (args.dt ** 2)
							axis = (
								int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
							center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
							cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], -1)

				image_new = cv2.addWeighted(img, 0.2, overlay_ellipses, 0.8, 0)
				width = 1200
				height = 1200
				dim = (width, height)
				# resize image
				resized_img = cv2.resize(image_new, dim, interpolation=cv2.INTER_AREA)
				if False:
					cv2.imshow("image",resized_img)
					cv2.waitKey(100)
				recorder.write(resized_img)
		recorder.release()

def plot_local_scenario_vel_OpenCV2(trajectories, batch_target, all_predictions, args, exp_num=0, display = True):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	homography_file = os.path.join(args.data_path + args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')
	scenario = args.scenario
	video_file = args.model_path +"/"  + args.model_name+ "_" + str(exp_num) + '_local.avi'
	img_file = args.data_path + args.scenario +'/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = ~cv2.imread(img_file)

		scale_percent = 100  # percent of original size
		width = 300
		height = 300
		dim = (width, height)
		# resize image
		resized_img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		out = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width, height))

		# Adding legend
		font = cv2.FONT_HERSHEY_SIMPLEX

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]
			ped_vel = batch_target[exp_idx][0]

			for t in range(len(predictions)):
				pedestrian_trajectory = traj.pose_vec[:,:2]

				# Plot Query Agent Positions
				query_agent_pos = traj.pose_vec[t,:2]
				query_agent = np.expand_dims(np.array((query_agent_pos[0], query_agent_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, query_agent)*int(scale_percent / 100)
				sigma_x = 0.4
				sigma_y = 0.4
				axis = (
					int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
				query_agent_center = np.array([obsv_XY[0, 0], obsv_XY[0, 1]])
				axis_pos = query_agent_center - np.array([args.submap_width/args.submap_resolution/2,args.submap_height/args.submap_resolution/2])
				overlay_ellipses = resized_img.copy()
				# Crop image for obtain surrounding static environment
				crop_img = cv2.getRectSubPix(overlay_ellipses,(int(args.submap_height / args.submap_resolution),int(args.submap_width / args.submap_resolution)),(query_agent_center[0],query_agent_center[1]))

				# resize image
				resized_crop_img = cv2.resize(crop_img, dim, interpolation=cv2.INTER_AREA)

				cv2.line(resized_crop_img, (width - 150, 30), (width - 130, 30), (0, 255, 0), 4)
				cv2.putText(resized_crop_img, "Query Agent", (width - 120, 30), font, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 50), (width - 130, 50), (255, 0, 0), 4)
				cv2.putText(resized_crop_img, "Other Agents", (width - 120, 50), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 70), (width - 130, 70), (255, 255, 0), 4)
				cv2.putText(resized_crop_img, "Query Agent traj", (width - 120, 70), font, 0.5, (255, 255, 0), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 90), (width - 130, 90), (128, 128, 128), 4)
				cv2.putText(resized_crop_img, "Pred traj 1", (width - 120, 90), font, 0.5, (128, 128, 128), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 110), (width - 130, 110), (0, 128, 255), 4)
				cv2.putText(resized_crop_img, "Pred traj 2", (width - 120, 110), font, 0.5, (0, 128, 255), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 130), (width - 130, 130), (255, 0, 127), 4)
				cv2.putText(resized_crop_img, "Pred traj 3", (width - 120, 130), font, 0.5, (255, 0, 127), 2, cv2.LINE_AA)

				# Plot Query Agent Positions
				center = (dim[0]//2, dim[1]//2)
				cv2.ellipse(resized_crop_img, center, axis, 0, 0, 360, (0, 255, 0), -1)

				# Plot other pedestrians
				agent_pos = traj.other_agents_positions[t]
				for agent_id in range(agent_pos.shape[0]):
					ped_pos = np.expand_dims(np.array((agent_pos[agent_id, 0], agent_pos[agent_id, 1])), axis=0).astype(float)
					obsv_XY = sup.to_image_frame(Hinv, ped_pos)*int(scale_percent / 100)
					axis = (
						int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100),
						int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
					center = (np.array([obsv_XY[0, 0], obsv_XY[0, 1]]) - axis_pos)
					if (center[0]>0) & (center[0]<width) & (center[1]>0) & (center[1]<height):
						agent_center = (int(center[0]),int(center[1]))
						cv2.ellipse(resized_crop_img, agent_center, axis, 0, 0, 360, (0, 0, 0), -1)

				# Query-agent real trajectory
				obsv_XY = (sup.to_image_frame(Hinv, pedestrian_trajectory[t:, :])- axis_pos)
				obsv_XY_filtered = []
				for i in range(obsv_XY.shape[0]):
					if (obsv_XY[i,0]>0) & (obsv_XY[i,0]<dim[0]) & (obsv_XY[i,1]>0) & (obsv_XY[i,1]<dim[1]):
						obsv_XY_filtered.append(obsv_XY[i])
				sup.line_cv(resized_crop_img, np.asarray(obsv_XY_filtered), (255, 0, 0), 1)  # bgr convention

				#Query-agent Predicted Real TRajectory
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.prediction_horizon):
					pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_trajectory = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=pedestrian_velocity, dt=args.dt)/args.submap_resolution
				obsv_XY = (sup.to_image_frame(Hinv, pedestrian_trajectory)- axis_pos)*int(scale_percent / 100)
				obsv_XY_filtered = []
				for i in range(obsv_XY.shape[0]):
					if (obsv_XY[i, 0] > 0) & (obsv_XY[i, 0] < dim[0]) & (obsv_XY[i, 1] > 0) & (obsv_XY[i, 1] < dim[1]):
						obsv_XY_filtered.append(obsv_XY[i])
				sup.line_cv(resized_crop_img, np.asarray(obsv_XY_filtered), (255, 255, 0), 1)  # bgr convention

				# Plot Predictions
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				""""""
				for mix_id in range(args.n_mixtures):
					for pred_step in range(args.prediction_horizon):
						idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
						pred_vel[pred_step, :] = np.array((predictions[t][idx],predictions[t][idx+args.n_mixtures]
						                                   ))
						if args.output_pred_state_dim > args.output_dim:
							sigmax[pred_step, :] = predictions[t][idx + 2*args.n_mixtures]
							sigmay[pred_step, :] = predictions[t][idx + 3*args.n_mixtures]
					traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=pred_vel, dt=args.dt)/args.submap_resolution
					obsv_XY = (sup.to_image_frame(Hinv, traj_pred)- axis_pos)*int(scale_percent / 100)
					obsv_XY_filtered = []
					for i in range(obsv_XY.shape[0]):
						if (obsv_XY[i, 0] > 0) & (obsv_XY[i, 0] < dim[0]) & (obsv_XY[i, 1] > 0) & (obsv_XY[i, 1] < dim[1]):
							obsv_XY_filtered.append(obsv_XY[i])
					sup.line_cv(resized_crop_img, np.asarray(obsv_XY_filtered), colors[mix_id], 1)  # bgr convention

					if args.output_pred_state_dim > args.output_dim:
						sigma_x = np.square(sigmax[0]) * (args.dt ** 2)/ args.submap_resolution
						sigma_y = np.square(sigmay[0]) * (args.dt ** 2)/ args.submap_resolution

						obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
						axis = (
							int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
						if (obsv_XY[0, 0] > 0) & (obsv_XY[0, 0] < dim[0]) & (obsv_XY[0, 1] > 0) & (obsv_XY[0, 1] < dim[1]):
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							cv2.ellipse(resized_crop_img, center, axis, 0, 0, 360, colors[mix_id], -1)

						for pred_step in range(1,args.prediction_horizon):
							sigma_x += np.square(sigmax[pred_step]) * (args.dt ** 2)/ args.submap_resolution
							sigma_y += np.square(sigmay[pred_step]) * (args.dt ** 2)/ args.submap_resolution
							axis = (
								int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
							if (obsv_XY[pred_step, 0] > 0) & (obsv_XY[pred_step, 0] < dim[0]) & (obsv_XY[pred_step, 1] > 0) & (obsv_XY[pred_step, 1] < dim[1]):
								center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
								cv2.ellipse(resized_crop_img, center, axis, 0, 0, 360, colors[mix_id], -1)

				if display:
					cv2.imshow("image", resized_crop_img)
					cv2.waitKey(10)
				out.write(resized_crop_img)

		# When everything done, release the video capture and video write objects
		out.release()

	print("Done recording...")

def plot_dataset(trajs, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
	"""
	homography_file = os.path.join(args.data_path + args.scenario, 'H.txt')
	print("Looking for homography at: " + homography_file)
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')

	img_file = args.data_path + args.scenario +'/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = cv2.imread(img_file)
		map_resolution = 1/Hinv[0,1]
		height, width, channels = img.shape

		# Adding legend

		font = cv2.FONT_HERSHEY_SIMPLEX

		for exp_idx in range(len(trajs)):
			overlay_img = img.copy()
			save_img_to_file = args.data_path + "/" + args.scenario + "/" + str(exp_idx) + '.png'

			cv2.line(overlay_img, (width - 150, 30), (width - 130, 30), (0, 0, 0), 4)
			cv2.putText(overlay_img, "Agent", (width - 120, 30), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)
			cv2.line(overlay_img, (width - 150, 50), (width - 130, 50), (255, 0, 0), 4)
			cv2.putText(overlay_img, "Robot", (width - 120, 50), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)

			sigma_x = 1.0
			sigma_y = 1.0
			axis = (
				int(np.sqrt(sigma_x) / map_resolution), int(np.sqrt(sigma_y) / map_resolution))

			for t in range(len(trajs[exp_idx])):
				dict = trajs[exp_idx][t]
				robot_pos = np.expand_dims(dict["robot_state"][:2],axis=0).astype(float)
				ped_pos = np.expand_dims(dict["pedestrian_state"]["position"], axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, robot_pos)
				center = (obsv_XY[0, 0], obsv_XY[0,1])
				cv2.ellipse(overlay_img, center, axis, 0, 0, 360, (0, 0, 0), -1)

				obsv_XY = sup.to_image_frame(Hinv, ped_pos)
				center = (obsv_XY[0, 0], obsv_XY[0, 1])
				cv2.ellipse(overlay_img, center, axis, 0, 0, 360, (255, 0, 0), -1)

			if True:
				cv2.imshow("image",overlay_img)
				cv2.waitKey(100)
			cv2.imwrite(save_img_to_file,overlay_img)

def matplot_dataset(trajs, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
	"""
	fig = pl.figure('Training Performance ' + args.model_name)
	ax = pl.subplot()
	for exp_idx in range(len(trajs)):
		ax.clear()
		save_img_to_file = args.data_path + "/" + args.scenario + "/" + str(exp_idx) + '.png'

		sigma_x = 1.0
		sigma_y = 1.0
		ax.set_aspect('equal')
		ax.set_xlabel('x [m]')
		ax.set_ylabel('y [m]')
		robot_traj = np.zeros((len(trajs[exp_idx]),2))
		ped_traj = np.zeros((len(trajs[exp_idx]), 2))
		for t in range(len(trajs[exp_idx])):
			dict = trajs[exp_idx][t]
			robot_pos = dict["robot_state"][:2]
			ped_pos = dict["pedestrian_state"]["position"]
			robot_traj[t, :] = robot_pos
			ped_traj[t,:] = ped_pos

			e1 = Ellipse(xy=(robot_pos[0], robot_pos[1]), width=sigma_x / 2, height=sigma_y / 2,
				             angle=0 / np.pi * 180)
			e1.set_alpha(0.5)

			ax.add_patch(e1)

			e1 = Ellipse(xy=(ped_pos[0], ped_pos[1]), width=sigma_x / 2, height=sigma_y / 2,
				             angle=0 / np.pi * 180,facecolor='orange')
			e1.set_alpha(0.5)

			ax.add_patch(e1)
			pl.show(block=False)
		min_x = min(min(robot_traj[:,0]),min(ped_traj[:,0]))-1.0
		max_x = max(max(robot_traj[:, 0]), max(ped_traj[:, 0])) + 1.0
		min_y = min(min(robot_traj[:, 1]), min(ped_traj[:, 1])) - 1.0
		max_y = max(max(robot_traj[:, 1]), max(ped_traj[:, 1])) + 1.0
		ax.set_xlim([min_x, max_x])
		ax.set_ylim([min_y, max_y])
		fig.savefig(save_img_to_file)

build/lib/data_utils/plot_utils.py
eliatrevisan/IA-MPPI-LBM
import os
import csv
import sys
import pickle as pkl
import argparse

sys.path.append('../src/data_utils')
if sys.version_info[0] < 3:
	import Support as sup
else:
	import src.data_utils.Support as sup

sys.path.append('../src/data_utils')
class bcolors:
	HEADER = '\033[95m'
	OKBLUE = '\033[94m'
	OKGREEN = '\033[92m'
	WARNING = '\033[93m'
	FAIL = '\033[91m'
	ENDC = '\033[0m'
	BOLD = '\033[1m'
	UNDERLINE = '\033[4m'


def write_summary(training_loss, args):
	if not os.path.isfile("trained_models_summary.csv"):
		with open("trained_models_summary.csv", 'w') as csvfile:
			# Write header
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			writer.writerow(["Model name", "Dataset", "Ped Info Size",
			                 "Batch size", "Training loss", "MSE", "FDE",
			                 "N epochs", "Learning rate", "Grad clip", "N mixtures","Test Dataset"])
			writer.writerow(
				[args.model_name, args.scenario, args.pedestrian_vector_dim, args.batch_size,
				 training_loss, 0, 0, args.n_epochs,
				 args.learning_rate_init, args.grads_clip, args.n_mixtures])
	else:
		with open("trained_models_summary.csv", 'r') as readFile:
			reader = csv.reader(readFile)
			lines = list(reader)
		with open("trained_models_summary.csv", 'w') as csvfile:
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			if len(lines) > args.exp_num + 1:
				for i in range(len(lines)):
					if i == args.exp_num + 1:
						writer.writerow(
							[args.model_name + "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size,
							 training_loss, 0, 0, args.n_epochs,
							 args.learning_rate_init, args.grads_clip, args.n_mixtures])
					else:
						writer.writerow(lines[i])
			else:
				writer.writerows(lines)
				writer.writerow(
					[args.model_name, args.scenario, args.pedestrian_vector_dim, args.batch_size,
					 training_loss, 0, 0, args.n_epochs,
					 args.learning_rate_init, args.grads_clip, args.n_mixtures])


def write_results_summary(mse, fde, avg_div, args, test_args, cvmse, cvfde):
	if not os.path.isfile(test_args.model_name + "_summary.csv"):
		with open(test_args.model_name + "_summary.csv", 'w') as csvfile:
			# Write header
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			writer.writerow(["Model name", "Dataset", "Ped Info Size",
			                 "Batch size", "MSE", "FDE","Diversity",
			                 "dt", "Prediction Horizon", "Previous Steps", "tbpt","Test Dataset","Others Info", "CV-MSE", "CV-FDE"])
			writer.writerow(
				[args.model_name+ "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size,
				 mse, fde,avg_div, args.dt,
				 args.prediction_horizon, args.prev_horizon, args.truncated_backprop_length,test_args.scenario,args.others_info, cvmse, cvfde])

	else:
		with open(test_args.model_name + "_roboat_summary.csv", 'r') as readFile:
			reader = csv.reader(readFile)
			lines = list(reader)
		with open(test_args.model_name + "_roboat_summary.csv", 'w') as csvfile:
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			if len(lines) > 3000:#args.exp_num + 1:
				for i in range(len(lines)):
					if i == args.exp_num + 1:
						writer.writerow(
							[args.model_name + "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size,
							 mse, fde, avg_div,args.dt,
							 args.prediction_horizon, args.prev_horizon, args.truncated_backprop_length, test_args.scenario,args.others_info, cvmse, cvfde])
					else:
						writer.writerow(lines[i])
			else:
				writer.writerows(lines)
				writer.writerow(
					[args.model_name + "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size,
					 mse, fde,avg_div, args.dt,
					 args.prediction_horizon, args.prev_horizon, args.truncated_backprop_length, test_args.scenario,args.others_info, cvmse, cvfde])

def write_keras_results_summary(loss, args):
	if not os.path.isfile(args.model_name + "_summary.csv"):
		with open(args.model_name + "_summary.csv", 'w') as csvfile:
			# Write header
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			writer.writerow(["Model name", "Dataset", "Ped Info Size", "Batch size", "Number of steps", "Loss"])
			writer.writerow( [args.model_name+ "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size, args.n_steps, loss] )

	else:
		with open(args.model_name + "_summary.csv", 'r') as readFile:
			reader = csv.reader(readFile)
			lines = list(reader)
		with open(args.model_name + "_summary.csv", 'w') as csvfile:
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			if len(lines) > 3000:#args.exp_num + 1:
				for i in range(len(lines)):
					if i == args.exp_num + 1:
						writer.writerow( [args.model_name+ "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size, args.n_steps, loss] )
					else:
						writer.writerow(lines[i])
			else:
				writer.writerows(lines)
				writer.writerow( [args.model_name+ "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size, args.n_steps, loss] )


def convert_dataset_python_version(pwd,pwd_out):
	print("Loading data from: '{}'".format(pwd))
	file = open(pwd, 'rb')
	tmp_self = pkl.load(file, encoding='latin1')
	file.close()
	file = open(pwd_out, 'wb')
	pkl.dump(tmp_self, file, protocol=2)


def model_selector(args):
	if "modelKerasRNN_arbitraryAgents" in args.model_name:
		if sys.version_info[0] < 3:
			from modelKerasRNN_arbitraryAgents import NetworkModel
		else:
			from src.models.modelKerasRNN_arbitraryAgents import NetworkModel
	elif "IntNet" in args.model_name:
		if sys.version_info[0] < 3:
			from IntNet import NetworkModel
		else:
			from src.models.IntNet import NetworkModel
	elif "Grid" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_Grid import NetworkModel
		else:
			from src.models.VGDNN_Grid import NetworkModel
	elif "no_grid" in args.model_name:
		if sys.version_info[0] < 3:
			from RNN_no_grid_goal import NetworkModel
		else:
			from src.models.RNN_no_grid_goal import NetworkModel
	elif "goal" in args.model_name:
		if sys.version_info[0] < 3:
			from RNN_goal import NetworkModel
		else:
			from src.models.RNN_goal import NetworkModel
	elif "simple" in args.model_name:
		if sys.version_info[0] < 3:
			from RNN_simple import NetworkModel
		else:
			from src.models.RNN_simple import NetworkModel
	elif "RNN" in args.model_name:
		if sys.version_info[0] < 3:
			from RNN import NetworkModel
		else:
			from src.models.RNN import NetworkModel
	elif "KL" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_KL import NetworkModel
		else:
			from src.models.VGDNN_KL import NetworkModel
	elif "diversity" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_attention_diversity import NetworkModel
		else:
			from src.models.VGDNN_attention_diversity import NetworkModel
	elif "ped" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_ped_attention import NetworkModel
		else:
			from src.models.VGDNN_ped_attention import NetworkModel
	elif "attention" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_attention import NetworkModel
		else:
			from src.models.VGDNN_attention import NetworkModel
	elif "multihead" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_multihead_attention import NetworkModel
		else:
			from src.models.VGDNN_multihead_attention import NetworkModel
	elif "simple" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_simple import NetworkModel
		else:
			from src.models.VGDNN_simple import NetworkModel
	elif "pos" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_pos import NetworkModel
		else:
			from src.models.VGDNN_pos import NetworkModel
	else:
		if sys.version_info[0] < 3:
			from VGDNN import NetworkModel
		else:
			from src.models.VGDNN import NetworkModel

	return NetworkModel(args)

def parse_args(defaults, parse_type):
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name', help='Path to directory that comprises the model (default="model_name").', type=str, default=defaults['model_name'])
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=defaults['exp_num'])
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").', type=str, default=defaults['scenario'])
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str, default=defaults['data_path'])

	if parse_type == "Train" or parse_type == "train":
		parser.add_argument('--input1_type', help='Data type for the first input', type=str, default=defaults['input1_type'])
		parser.add_argument('--input2_type', help='Data type for the second input', type=str, default=defaults['input2_type'])
		parser.add_argument('--model_path', help='Path to directory to save the model (default=""../trained_models/"+model_name").', type=str, default=defaults['model_path'])
		parser.add_argument('--pretrained_convnet_path', help='Path to directory that comprises the pre-trained convnet model (default=" ").', type=str, default=defaults['pretrained_convnet_path'])
		parser.add_argument('--log_dir', help='Path to the log directory of the model (default=""../trained_models/"+model_name").', type=str, default=defaults['log_dir'])
		parser.add_argument('--scenario_val', help='Dataset pkl file', type=str, default= defaults['scenario_val'])
		parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool, default=defaults['real_world_data'])
		parser.add_argument('--dataset', help='Dataset pkl file', type=str, default= defaults['scenario'] + '.pkl')
		parser.add_argument('--dataset_val', help='Dataset pkl file', type=str, default= defaults['scenario_val'] + '.pkl')
		parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str, default='LSTM')
		parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool, default=defaults['warmstart_model'])
		parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool, default=defaults['warm_start_convnet'])
		parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float, default=defaults['dt'])
		parser.add_argument('--n_steps', help='Number of epochs (default=10000).', type=int, default=defaults['n_steps'])
		parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=defaults['batch_size'])
		parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float, default=defaults['regularization_weight'])
		parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float, default=defaults['keep_prob'])
		parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float, default=defaults['learning_rate_init'])
		parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float, default=defaults['beta_rate_init'])
		parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool, default=defaults['dropout'])
		parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float, default=defaults['grads_clip'])
		parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).', type=int, default=defaults['truncated_backprop_length'])
		parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int, default=defaults['prediction_horizon'])
		parser.add_argument('--prev_horizon', help='Previous seq length.', type=int, default=defaults['prev_horizon'])
		parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int, default=defaults['rnn_state_size'])
		parser.add_argument('--rnn_state_size_lstm_ped', help='Number of RNN / LSTM units of the pedestrian lstm layer (default=32).', type=int, default=defaults['rnn_state_size_lstm_ped'])
		parser.add_argument('--rnn_state_size_bilstm_ped', help='Number of RNN / LSTM units of the pedestrian bidirectional lstm layer (default=32).',type=int, default=defaults['rnn_state_size_bilstm_ped'])
		parser.add_argument('--rnn_state_size_lstm_grid', help='Number of RNN / LSTM units of the grid lstm layer (default=32).', type=int, default=defaults['rnn_state_size_lstm_grid'])
		parser.add_argument('--rnn_state_size_lstm_concat', help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).', type=int, default=defaults['rnn_state_size_lstm_concat'])
		parser.add_argument('--prior_size', help='prior_size', type=int, default=defaults['prior_size'])
		parser.add_argument('--latent_space_size', help='latent_space_size', type=int, default=defaults['latent_space_size'])
		parser.add_argument('--x_dim', help='x_dim', type=int, default=defaults['x_dim'])
		parser.add_argument('--fc_hidden_unit_size', help='Number of fully connected layer units after LSTM layer (default=64).', type=int, default=defaults['fc_hidden_unit_size'])
		parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int, default=defaults['input_state_dim'])
		parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float, default=defaults['input_dim'])
		parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float, default=defaults['output_dim'])
		parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int, default=defaults['output_pred_state_dim'])
		parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int, default=defaults['cmd_vector_dim'])
		parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int, default=defaults['n_mixtures'])
		parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int, default=defaults['pedestrian_vector_dim'])
		parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int, default=defaults['pedestrian_vector_state_dim'])
		parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float, default=defaults['max_range_ped_grid'])
		parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float, default=defaults['pedestrian_radius'])
		parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=int, default=0)
		parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int, default=defaults['print_freq'])
		parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).', type=int, default=defaults['save_freq'])
		parser.add_argument('--patience', help='Patience to stop the training loop early', type=int, default=defaults['patience'])
		parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool, default=False)
		parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool, default=defaults['agents_on_grid'])
		parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool, default=defaults['normalize_data'])
		parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool, default=defaults['rotated_grid'])
		parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool, default=defaults['centered_grid'])
		parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float, default=0)
		parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=defaults['submap_width'])
		parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=defaults['submap_height'])
		parser.add_argument('--submap_resolution', help='Map resolution.', type=float, default=defaults['submap_resolution'])
		parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
		parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
		parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
		parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool, default=False)
		parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool, default=defaults['predict_positions'])
		parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool, default=True)
		parser.add_argument('--relative_info', help='Use relative info for other agents.', type=sup.str2bool, default=True)
		parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool, default=defaults['regulate_log_loss'])
		parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool, default=defaults['diversity_update'])
		parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str, default='../config/topics.yaml')
		parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
		parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
		parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
		parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
		parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
		parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
		parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
		parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
		parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
		parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
		parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
		parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	else:
		parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
		parser.add_argument('--n_samples', help='Number of samples', type=int, default=1)
		parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool, default=True)
		parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool, default=False)
		parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float, default=0)
		parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float, default=5)
		parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool, default=False)
		parser.add_argument('--update_state', help='update_state', type=sup.str2bool, default=False)
		parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool, default=False)
		parser.add_argument('--freeze_other_agents', help='Freeze other agents.', type=sup.str2bool, default=False)
		parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,default=False)

	args = parser.parse_args()

	return args

if __name__ == '__main__':
	convert_dataset_python_version("/home/bbrito/Code/I-LSTM/data/2_agents_swap/trajs/GA3C-CADRL-10.pkl",
	                               "/home/bbrito/Code/I-LSTM/data/2_agents_swap/trajs/GA3C-CADRL-10-py27.pkl")

build/lib/data_utils/utils.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
# import tensorflow_probability as tfp
import numpy as np
import os

def scaled_dot_product_attention(q, k, v, mask):
	"""Calculate the attention weights.
	q, k, v must have matching leading dimensions.
	k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
	The mask has different shapes depending on its type(padding or look ahead)
	but it must be broadcastable for addition.

	Args:
		q: query shape == (..., seq_len_q, depth)
		k: key shape == (..., seq_len_k, depth)
		v: value shape == (..., seq_len_v, depth_v)
		mask: Float tensor with shape broadcastable
					to (..., seq_len_q, seq_len_k). Defaults to None.

	Returns:
		output, attention_weights
	"""

	matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

	# scale matmul_qk
	dk = tf.cast(tf.shape(k)[-1], tf.float32)
	scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

	# add the mask to the scaled tensor.
	if mask is not None:
		scaled_attention_logits += (mask * -1e9)

	# softmax is normalized on the last axis (seq_len_k) so that the scores
	# add up to 1.
	attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

	output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

	return output, attention_weights

class MultiHeadAttention(tf.keras.layers.Layer):
	def __init__(self, d_model, num_heads):
		super(MultiHeadAttention, self).__init__()
		self.num_heads = num_heads
		self.d_model = d_model

		assert d_model % self.num_heads == 0

		self.depth = d_model // self.num_heads

		self.wq = tf.keras.layers.Dense(d_model)
		self.wk = tf.keras.layers.Dense(d_model)
		self.wv = tf.keras.layers.Dense(d_model)

		self.dense = tf.keras.layers.Dense(d_model)

	def split_heads(self, x, batch_size):
		"""Split the last dimension into (num_heads, depth).
		Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
		"""

		x = tf.reshape(x, (batch_size, 1, self.num_heads, self.depth))
		return tf.transpose(x, perm=[0, 2, 1, 3])

	def call(self, v, k, q, mask):
		batch_size = tf.shape(q)[0]

		q = self.wq(q)  # (batch_size, seq_len, d_model)
		k = self.wk(k)  # (batch_size, seq_len, d_model)
		v = self.wv(v)  # (batch_size, seq_len, d_model)


		q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
		k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
		v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

		# scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
		# attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
		scaled_attention, attention_weights = scaled_dot_product_attention(
			q, k, v, mask)

		scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

		concat_attention = tf.reshape(scaled_attention,
																	(batch_size, 1, self.d_model))  # (batch_size, seq_len_q, d_model)

		output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

		return output, attention_weights
build/lib/external/attention.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_h = h_dim
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_x, state_is_tuple=True)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state

            #   Eq 5
            x_in = x[:,:x.shape[1]-self.args.pedestrian_vector_dim]
            ped_info = x[:,x.shape[1]-self.args.pedestrian_vector_dim:x.shape[1]]
            with tf.variable_scope("Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.exp(linear(prior_hidden, self.n_z))

            with tf.variable_scope("phi_x"):
                x_1 = tf.nn.relu(linear(x_in, self.n_x_1))
                #x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Encoder"):
                with tf.variable_scope("hidden"):
                    enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    #enc_hidden = linear(tf.concat(axis=1, values=(x_1, h)), self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.exp(linear(enc_hidden, self.n_z))
            eps = tf.random_normal(tf.shape(enc_mu), prior_mu, prior_sigma, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("phi_z"):
                z_1 = tf.nn.relu(linear(z, self.n_z_1))

            # propagate hidden state eq(7)
            output, state2 = self.lstm(tf.concat(axis=1, values=(x_1, z_1,ped_info)), state)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
build/lib/external/irnn_cell.py
eliatrevisan/IA-MPPI-LBM
import os
import time
import matplotlib.pyplot as plt
import copy
import torch
import numpy as np
import torch.nn as nn
import torch.optim as opt
from tqdm import tqdm, trange
from itertools import chain
from torch.autograd import Variable
from torch.utils.data import DataLoader

# SOCIALWAYS: ====== training hyper-parameters ======
# Unrolled GAN
n_unrolling_steps = 0
# Info GAN
use_info_loss = True
loss_info_w = 0.5
n_latent_codes = 2
# L2 GAN
use_l2_loss = False
use_variety_loss = False
loss_l2_w = 0.5  # WARNING for both l2 and variety
# Learning Rate
lr_g = 1E-3
lr_d = 1E-4
# FIXME: ====== Network Size ===================
# Batch size
batch_size = 256
# LSTM hidden size
hidden_size = 64
n_epochs = 100000
num_social_features = 3
social_feature_size = hidden_size
noise_len = hidden_size // 2
n_lstm_layers = 1
use_social = False
n_next = 12
# ================================================

# Augment tensors of positions into positions+velocity
def get_traj_4d(obsv_p, pred_p):
    obsv_v = obsv_p[:, 1:] - obsv_p[:, :-1]
    obsv_v = torch.cat([obsv_v[:, 0].unsqueeze(1), obsv_v], dim=1)
    obsv_4d = torch.cat([obsv_p, obsv_v], dim=2)
    if len(pred_p) == 0: return obsv_4d
    pred_p_1 = torch.cat([obsv_p[:, -1].unsqueeze(1), pred_p[:, :-1]], dim=1)
    pred_v = pred_p - pred_p_1
    pred_4d = torch.cat([pred_p, pred_v], dim=2)
    return obsv_4d, pred_4d


# Evaluate the error between the model prediction and the true path
def calc_error(pred_hat, pred):
    N = pred.size(0)
    T = pred.size(1)
    err_all = torch.pow((pred_hat - pred) / ss, 2).sum(dim=2).sqrt()  # N x T
    FDEs = err_all.sum(dim=0).item() / N
    ADEs = torch.cumsum(FDEs)
    for ii in range(T):
        ADEs[ii] /= (ii + 1)
    return ADEs.data.cpu().numpy(), FDEs.data().cpu().numpy()


class AttentionPooling(nn.Module):
    def __init__(self, h_dim, f_dim):
        super(AttentionPooling, self).__init__()
        self.f_dim = f_dim
        self.h_dim = h_dim
        self.W = nn.Linear(h_dim, f_dim, bias=True)

    def forward(self, f, h, sub_batches):
        Wh = self.W(h)
        S = torch.zeros_like(h)
        for sb in sub_batches:
            N = sb[1] - sb[0]
            if N == 1: continue

            for ii in range(sb[0], sb[1]):
                fi = f[ii, sb[0]:sb[1]]
                sigma_i = torch.bmm(fi.unsqueeze(1), Wh[sb[0]:sb[1]]. unsqueeze(2))
                sigma_i[ii-sb[0]] = -1000

                attentions = torch.softmax(sigma_i.squeeze(), dim=0)
                S[ii] = torch.mm(attentions.view(1, N), h[sb[0]:sb[1]])

        return S


class EmbedSocialFeatures(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(EmbedSocialFeatures, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.fc = nn.Sequential(nn.Linear(input_size, 32), nn.ReLU(),
                                nn.Linear(32, 64), nn.ReLU(),
                                nn.Linear(64, hidden_size))

    def forward(self, ftr_list, sub_batches):
        embedded_features = self.fc(ftr_list)
        return embedded_features


def DCA(xA_4d, xB_4d):
    dp = xA_4d[:2] - xB_4d[:2]
    dv = xA_4d[2:] - xB_4d[2:]
    ttca = torch.dot(-dp, dv) / (torch.norm(dv) ** 2 + 1E-6)
    # ttca = torch.max(ttca, 0)
    dca = torch.norm(dp + ttca * dv)
    return dca


def Bearing(xA_4d, xB_4d):
    dp = xA_4d[:2] - xB_4d[:2]
    v = xA_4d[2:]
    cos_theta = torch.dot(dp, v) / (torch.norm(dp) * torch.norm(v) + 1E-6)
    return cos_theta


def DCA_MTX(x_4d, D_4d):
    Dp = D_4d[:, :, :2]
    Dv = D_4d[:, :, 2:]
    DOT_Dp_Dv = torch.mul(Dp[:,:,0], Dv[:,:,0]) + torch.mul(Dp[:,:,1], Dv[:,:,1])
    Dv_sq = torch.mul(Dv[:,:,0], Dv[:,:,0]) + torch.mul(Dv[:,:,1], Dv[:,:,1]) + 1E-6
    TTCA = -torch.div(DOT_Dp_Dv, Dv_sq)
    DCA = torch.zeros_like(Dp)
    DCA[:, :, 0] = Dp[:, :, 0] + TTCA * Dv[:, :, 0]
    DCA[:, :, 1] = Dp[:, :, 1] + TTCA * Dv[:, :, 1]
    DCA = torch.norm(DCA, dim=2)
    return DCA


def BearingMTX(x_4d, D_4d):
    Dp = D_4d[:, :, :2]  # NxNx2
    v = x_4d[:, 2:].unsqueeze(1).repeat(1, x_4d.shape[0], 1)  # => NxNx2
    DOT_Dp_v = Dp[:, :, 0] * v[:, :, 0] + Dp[:, :, 1] * v[:, :, 1]
    COS_THETA = torch.div(DOT_Dp_v, torch.norm(Dp, dim=2) * torch.norm(v, dim=2) + 1E-6)
    return COS_THETA


def SocialFeatures(x, sub_batches):
    N = x.shape[0]  # x is NxTx4 tensor

    x_ver_repeat = x[:, -1].unsqueeze(0).repeat(N, 1, 1)
    x_hor_repeat = x[:, -1].unsqueeze(1).repeat(1, N, 1)
    Dx_mat = x_hor_repeat - x_ver_repeat

    l2_dist_MTX = Dx_mat[:, :, :2].norm(dim=2)
    bearings_MTX = BearingMTX(x[:, -1], Dx_mat)
    dcas_MTX = DCA_MTX(x[:, -1], Dx_mat)
    sFeatures_MTX = torch.stack([l2_dist_MTX, bearings_MTX, dcas_MTX], dim=2)

    return sFeatures_MTX   # directly return the Social Features Matrix


# LSTM path encoding module
class EncoderLstm(nn.Module):
    def __init__(self, hidden_size, n_layers=2):
        # Dimension of the hidden state (h)
        self.hidden_size = hidden_size
        super(EncoderLstm, self).__init__()
        # Linear embedding 4xh
        self.embed = nn.Linear(4, self.hidden_size)
        # The LSTM cell.
        # Input dimension (observations mapped through embedding) is the same as the output
        self.lstm = nn.LSTM(self.hidden_size, self.hidden_size, num_layers=n_layers, batch_first=True)
        self.lstm_h = []
        # init_weights(self)

    def init_lstm(self, h, c):
        # Initialization of the LSTM: hidden state and cell state
        self.lstm_h = (h, c)

    def forward(self, obsv):
        # Batch size
        bs = obsv.shape[0]
        # Linear embedding
        obsv = self.embed(obsv)
        # Reshape and applies LSTM over a whole sequence or over one single step
        y, self.lstm_h = self.lstm(obsv.view(bs, -1, self.hidden_size), self.lstm_h)
        return y


class Discriminator(nn.Module):
    def __init__(self, n_next, hidden_dim, n_latent_code):
        super(Discriminator, self).__init__()
        self.lstm_dim = hidden_dim
        self.n_next = n_next
        # LSTM Encoder for the observed part
        self.obsv_encoder_lstm = nn.LSTM(4, hidden_dim, batch_first=True)
        # FC sub-network: input is hidden_dim, output is hidden_dim//2. This ouput will be part of
        # the input of the classifier.
        self.obsv_encoder_fc = nn.Sequential(nn.Linear(hidden_dim, hidden_dim // 2), nn.LeakyReLU(0.2),
                                             nn.Linear(hidden_dim // 2, hidden_dim // 2))
        # FC Encoder for the predicted part: input is n_next*4 (whole predicted trajectory), output is
        # hidden_dim//2. This ouput will also be part of the input of the classifier.
        self.pred_encoder = nn.Sequential(nn.Linear(n_next * 4, hidden_dim // 2), nn.LeakyReLU(0.2),
                                          nn.Linear(hidden_dim // 2, hidden_dim // 2))
        # Classifier: input is hidden_dim (concatenated encodings of observed and predicted trajectories), output is 1
        self.classifier = nn.Sequential(nn.Linear(hidden_dim, hidden_dim // 2), nn.LeakyReLU(0.2),
                                        nn.Linear(hidden_dim // 2, 1))
        # Latent code inference: input is hidden_dim (concatenated encodings of observed and predicted trajectories), output is n_latent_code (distribution of latent codes)
        self.latent_decoder = nn.Sequential(nn.Linear(hidden_dim, hidden_dim // 2), nn.LeakyReLU(0.2),
                                            nn.Linear(self.lstm_dim // 2, n_latent_code))

    def forward(self, obsv, pred):
        bs = obsv.size(0)
        lstm_h_c = (torch.zeros(1, bs, self.lstm_dim).cuda(),
                    torch.zeros(1, bs, self.lstm_dim).cuda())
        # Encoding of the observed sequence trhough an LSTM cell
        obsv_code, lstm_h_c = self.obsv_encoder_lstm(obsv, lstm_h_c)
        # Further encoding through a FC layer
        obsv_code = self.obsv_encoder_fc(obsv_code[:, -1])
        # Encoding of the predicted/next part of the sequence through a FC layer
        pred_code = self.pred_encoder(pred.view(-1, self.n_next * 4))
        both_codes = torch.cat([obsv_code, pred_code], dim=1)
        # Applies classifier to the concatenation of the encodings of both parts
        label = self.classifier(both_codes)
        # Inference on the latent code
        code_hat = self.latent_decoder(both_codes)
        return label, code_hat

    def load(self, backup):
        for m_from, m_to in zip(backup.modules(), self.modules()):
            if isinstance(m_to, nn.Linear):
                m_to.weight.data = m_from.weight.data.clone()
                if m_to.bias is not None:
                    m_to.bias.data = m_from.bias.data.clone()


# FC path decoding module
class DecoderFC(nn.Module):
    def __init__(self, hidden_dim):
        super(DecoderFC, self).__init__()
        # Fully connected sub-network. Input is hidden_dim, output is 2.
        self.fc1 = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim), nn.LeakyReLU(0.2),
                                       # torch.nn.Linear(64, 64), nn.LeakyReLU(0.2),
                                       torch.nn.Linear(hidden_dim, hidden_dim // 2), nn.LeakyReLU(0.2),
                                       torch.nn.Linear(hidden_dim // 2, hidden_dim // 4),
                                       torch.nn.Linear(hidden_dim // 4, 2))

    def forward(self, h, s, z):
        # For each sample in the batch, concatenate h (hidden state), s (social term) and z (noise)
        inp = torch.cat([h, s, z], dim=1)
        # Applies the fully connected layer
        out = self.fc1(inp)
        return out


# LSTM path decoding module
class DecoderLstm(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(DecoderLstm, self).__init__()
        # Decoding LSTM
        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True)
        # Fully connected sub-network. Input is hidden_size, output is 2.
        self.fc = nn.Sequential(torch.nn.Linear(hidden_size, 64), nn.Sigmoid(),
                                torch.nn.Linear(64, 64), nn.LeakyReLU(0.2),
                                torch.nn.Linear(64, 32), nn.LeakyReLU(0.2),
                                torch.nn.Linear(32, 2))

        # init_weights(self)
        self.lstm_h = []

    def init_lstm(self, h, c):
        # Initialization of the LSTM: hidden state and cell state
        self.lstm_h = (h, c)

    def forward(self, h, s, z):
        # Batch size
        bs = z.shape[0]
        # For each sample in the batch, concatenate h (hidden state), s (social term) and z (noise)
        inp = torch.cat([h, s, z], dim=1)
        # Applies a forward step.
        out, self.lstm_h = self.lstm(inp.unsqueeze(1), self.lstm_h)
        # Applies the fully connected layer to the LSTM output
        out = self.fc(out.squeeze())
        return out


# ===================================================


# =============== Training Loop ==================
def train():
    tic = time.clock()
    # Evaluation metrics (ADE/FDE)
    train_ADE, train_FDE = 0, 0
    batch_size_accum = 0;
    sub_batches = []
    # For all the training batches
    for ii, batch_i in enumerate(train_batches):
        batch_size_accum += batch_i[1] - batch_i[0]
        sub_batches.append(batch_i)

        # FIXME: Just keep it for toy dataset
        # sub_batches = the_batches
        # batch_size_accum = sub_batches[-1][1]
        # ii = train_size-1

        if ii >= train_size - 1 or \
                batch_size_accum + (the_batches[ii + 1][1] - the_batches[ii + 1][0]) > batch_size:
            # Observed partial paths
            obsv = dataset_obsv[sub_batches[0][0]:sub_batches[-1][1]]
            # Future partial paths
            pred = dataset_pred[sub_batches[0][0]:sub_batches[-1][1]]
            sub_batches = sub_batches - sub_batches[0][0]
            # May have to fill with 0
            filling_len = batch_size - int(batch_size_accum)
            #obsv = torch.cat((obsv, torch.zeros(filling_len, n_past, 2).cuda()), dim=0)
            #pred = torch.cat((pred, torch.zeros(filling_len, n_next, 2).cuda()), dim=0)

            bs = batch_size_accum

            # Completes the positional vectors with velocities (to have dimension 4)
            obsv_4d, pred_4d = get_traj_4d(obsv, pred)
            zeros = Variable(torch.zeros(bs, 1) + np.random.uniform(0, 0.1), requires_grad=False).cuda()
            ones = Variable(torch.ones(bs, 1) * np.random.uniform(0.9, 1.0), requires_grad=False).cuda()
            noise = torch.FloatTensor(torch.rand(bs, noise_len)).cuda()

            # ============== Train Discriminator ================
            for u in range(n_unrolling_steps + 1):
                # Zero the gradient buffers of all parameters
                D.zero_grad()
                with torch.no_grad():
                    pred_hat_4d = predict(obsv, noise, n_next, sub_batches)

                fake_labels, code_hat = D(obsv_4d, pred_hat_4d)  # classify fake samples
                # Evaluate the MSE loss: the fake_labels should be close to zero
                d_loss_fake = mse_loss(fake_labels, zeros)
                d_loss_info = mse_loss(code_hat.squeeze(), noise[:, :n_latent_codes])
                # Evaluate the MSE loss: the real should be close to one
                real_labels, code_hat = D(obsv_4d, pred_4d)  # classify real samples
                d_loss_real = mse_loss(real_labels, ones)

                #  FIXME: which loss functinos to use for D?
                d_loss = d_loss_fake + d_loss_real
                if use_info_loss:
                    #
                    d_loss += loss_info_w * d_loss_info
                d_loss.backward()  # update D
                D_optimizer.step()

                if u == 0 and n_unrolling_steps > 0:
                    backup = copy.deepcopy(D)

            # =============== Train Generator ================= #
            # Zero the gradient buffers of all the discriminator parameters
            D.zero_grad()
            # Zero the gradient buffers of all the generator parameters
            predictor_optimizer.zero_grad()
            # Applies a forward step of prediction
            pred_hat_4d = predict(obsv, noise, n_next, sub_batches)

            # Classify the generated fake sample
            gen_labels, code_hat = D(obsv_4d, pred_hat_4d)
            # L2 loss between the predicted paths and the true ones
            g_loss_l2 = mse_loss(pred_hat_4d[:, :, :2], pred)
            # Adversarial loss (classification labels should be close to one)
            g_loss_fooling = mse_loss(gen_labels, ones)
            # Information loss
            g_loss_info = mse_loss(code_hat.squeeze(), noise[:, :n_latent_codes])

            #  FIXME: which loss functions to use for G?
            #
            g_loss = g_loss_fooling
            # If using the info loss
            if use_info_loss:
                g_loss += loss_info_w * g_loss_info
            # If using the L2 loss
            if use_l2_loss:
                g_loss += loss_l2_w * g_loss_l2
            if use_variety_loss:
                KV = 20
                all_20_losses = []
                for k in range(KV):
                    pred_hat_4d = predict(obsv, noise, n_next, sub_batches)
                    loss_l2_k = mse_loss(pred_hat_4d[k, :, :2], pred[k])
                all_20_losses.append(loss_l2_k.unsqueeze(0))
                all_20_losses = torch.cat(all_20_losses)
                variety_loss, _ = torch.min(all_20_losses, dim=0)
                g_loss += loss_l2_w * variety_loss

            g_loss.backward()
            predictor_optimizer.step()

            if n_unrolling_steps > 0:
                D.load(backup)
                del backup

            # calculate error
            with torch.no_grad():  # TODO: use the function above
                err_all = torch.pow((pred_hat_4d[:, :, :2] - pred) / ss, 2)
                err_all = err_all.sum(dim=2).sqrt()
                e = err_all.sum().item() / n_next
                train_ADE += e
                train_FDE += err_all[:, -1].sum().item()

            batch_size_accum = 0;
            sub_batches = []

    train_ADE /= n_train_samples
    train_FDE /= n_train_samples
    toc = time.clock()
    print(" Epc=%4d, Train ADE,FDE = (%.3f, %.3f) | time = %.1f" \
          % (epoch, train_ADE, train_FDE, toc - tic))


def test(n_gen_samples=1, linear=False, write_to_file=None, just_one=False):
    # =========== Test error ============
    plt.close()
    ade_avg_12, fde_avg_12 = 0, 0
    ade_min_12, fde_min_12 = 0, 0
    counter = 0
    for ii, batch_i in enumerate(test_batches):
        if ii > 20: break
        obsv = dataset_obsv[batch_i[0]:batch_i[1]]
        pred = dataset_pred[batch_i[0]:batch_i[1]]
        current_t = dataset_t[batch_i[0]]
        bs = int(batch_i[1] - batch_i[0])
        with torch.no_grad():
            all_20_errors = []
            all_20_preds = []

            linear_preds = predict_cv(obsv, n_next)
            if linear and not write_to_file:
                all_20_preds.append(linear_preds.unsqueeze(0))
                err_all = torch.pow((linear_preds[:, :, :2] - pred) / ss, 2).sum(dim=2, keepdim=True).sqrt()
                all_20_errors.append(err_all.unsqueeze(0))
            else:
                for kk in range(n_gen_samples):
                    noise = torch.FloatTensor(torch.rand(bs, noise_len)).cuda()
                    pred_hat_4d = predict(obsv, noise, n_next)
                    all_20_preds.append(pred_hat_4d.unsqueeze(0))
                    err_all = torch.pow((pred_hat_4d[:, :, :2] - pred) / ss, 2).sum(dim=2, keepdim=True).sqrt()
                    all_20_errors.append(err_all.unsqueeze(0))

            all_20_errors = torch.cat(all_20_errors)
            if write_to_file:
                file_name = os.path.join(write_to_file, str(epoch) + '-' + str(current_t) + '.npz')
                print('saving to ', file_name)
                np_obsvs = scale.denormalize(obsv[:, :, :2].data.cpu().numpy())
                np_preds_our = scale.denormalize(torch.cat(all_20_preds)[:, :, :, :2].data.cpu().numpy())
                np_preds_gtt = scale.denormalize(pred[:, :, :2].data.cpu().numpy())
                np_preds_lnr = scale.denormalize(linear_preds[:, :, :2].data.cpu().numpy())
                np.savez(file_name, timestamp=current_t,
                         obsvs=np_obsvs, preds_our=np_preds_our, preds_gtt=np_preds_gtt, preds_lnr=np_preds_lnr)

            # =============== Prediction Errors ================
            dummy = all_20_errors.mean(2).mean(0, keepdim=True)
            fde_min_12_i, _ = all_20_errors[:, :, -1].min(0, keepdim=True)
            ade_min_12_i, _ = all_20_errors.mean(2).min(0, keepdim=True)
            fde_min_12 += fde_min_12_i.sum().item()
            ade_min_12 += ade_min_12_i.sum().item()
            fde_avg_12 += all_20_errors[:, :, -1].mean(0, keepdim=True).sum().item()
            ade_avg_12 += all_20_errors.mean(2).mean(0, keepdim=True).sum().item()
            counter +=dummy.shape[1]
            # ==================================================
        if just_one: break

    ade_avg_12 /= counter
    fde_avg_12 /= counter
    ade_min_12 /= counter
    fde_min_12 /= n_test_samples
    print('Avg ADE,FDE (12)= (%.3f, %.3f) | Min(20) ADE,FDE (12)= (%.3f, %.3f)' \
          % (ade_avg_12, fde_avg_12, ade_min_12, fde_min_12))


# =======================================================
# ===================== M A I N =========================
# =======================================================
build/lib/external/socialways.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_h = h_dim
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_dec_hidden, state_is_tuple=True,reuse=tf.AUTO_REUSE)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state

            #   Eq 5
            with tf.variable_scope("Storn_Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.exp(linear(prior_hidden, self.n_z))

            with tf.variable_scope("Storn_phi_x"):
                x_1 = tf.nn.relu(linear(x, self.n_x_1))
                #x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Storn_Encoder"):
                with tf.variable_scope("hidden"):
                    enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    #enc_hidden = linear(tf.concat(axis=1, values=(x_1, h)), self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.exp(linear(enc_hidden, self.n_z))
            eps = tf.random_normal(tf.shape(enc_mu), 0.0, 1.0, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("Storn_phi_z"):
                z_1 = tf.nn.relu(linear(z, self.n_z_1))

            # propagate hidden state eq(7)
            concat = tf.concat(axis=1, values=(x_1, z_1))
            output, state2 = self.lstm(z_1, state)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
build/lib/external/storn_cell.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_h = h_dim
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_x, state_is_tuple=True,reuse=tf.AUTO_REUSE)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state

            #   Eq 5
            with tf.variable_scope("Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.exp(linear(prior_hidden, self.n_z)) # exp() ensures the std.dev. is always positive

            with tf.variable_scope("phi_x"):
                x_1 = tf.nn.relu(linear(x, self.n_x_1))
                #x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Encoder"):
                with tf.variable_scope("hidden"):
                    enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    #enc_hidden = linear(tf.concat(axis=1, values=(x_1, h)), self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.exp(linear(enc_hidden, self.n_z)) # exp() ensures the std.dev. is always positive
            eps = tf.random.normal(tf.shape(enc_mu), prior_mu, prior_sigma, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("phi_z"):
                z_1 = tf.nn.relu(linear(z, self.n_z_1))

            # propagate hidden state eq(7)
            output, state2 = self.lstm(tf.concat(axis=1, values=(x_1, z_1)), state) # tf.concat(axis=1, values=(x_1, z_1)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
build/lib/external/vrnn_cell.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_x, state_is_tuple=True,reuse=tf.AUTO_REUSE)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state
            x_in = x[:,:self.args.rnn_state_size]
            #   Eq 5
            with tf.variable_scope("Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.exp(linear(prior_hidden, self.n_z))

            with tf.variable_scope("phi_x"):
                x_1 = tf.nn.relu(linear(x, self.n_x_1))
                #x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Encoder"):
                with tf.variable_scope("hidden"):
                    enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    #enc_hidden = linear(tf.concat(axis=1, values=(x_1, h)), self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.exp(linear(enc_hidden, self.n_z))
            eps = tf.random.normal(tf.shape(enc_mu), prior_mu, prior_sigma, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("phi_z"):
                z_1 = tf.nn.relu(linear(z, self.n_z_1))

            # propagate hidden state eq(7)
            output, state2 = self.lstm(z_1, state) # tf.concat(axis=1, values=(x_1, z_1)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
build/lib/external/vrnn_cell2.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_h = h_dim
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_x, state_is_tuple=True)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state
            x_in = x[:,:self.args.rnn_state_size]
            x_ped = x[:,self.args.rnn_state_size:]
            #   Eq 5
            with tf.variable_scope("Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.nn.softplus(linear(prior_hidden, self.n_z))

            with tf.variable_scope("phi_x"):
                #x_1 = tf.nn.relu(linear(x, self.n_x_1))
                x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Encoder"):
                with tf.variable_scope("hidden"):
                    #enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    enc_hidden = linear(x_in, self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.nn.softplus(linear(enc_hidden, self.n_z))
            eps = tf.random.normal(tf.shape(enc_mu), prior_mu, prior_sigma, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("phi_z"):
                z_1 = linear(tf.concat(axis=1,values=(z,x_ped)), self.n_z_1)

            # propagate hidden state eq(7)
            output, state2 = self.lstm(z_1, state) # tf.concat(axis=1, values=(x_1, z_1)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
build/lib/external/vrnn_cell3.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
from cv2 import log
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
import json
import importlib
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils import DataHandlerRoboat as dhroboat
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
from datetime import datetime
import pickle as pkl
import time
from copy import deepcopy
from multiprocessing.pool import ThreadPool
import colorama
from colorama import Fore, Style
import matplotlib.pyplot as plt

# Model directories
def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
											help='Path to directory that comprises the model (default="model_name").',
											type=str, default="VGDNN")
	parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=247)
	parser.add_argument('--n_samples', help='Number of samples', type=int, default=10)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
											type=str, default="datasets/ewap_dataset/seq_eth")
	parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool,
											default=True)
	parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool,
											default=True)
	parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float,
											default=0)
	parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float,
											default=5)
	parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool,
											default=False)
	parser.add_argument('--update_state', help='update_state', type=sup.str2bool,
											default=False)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
											default=False)
	parser.add_argument('--freeze_other_agents', help='FReeze other agents.', type=sup.str2bool,
											default=False)
	parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,
											default=False)
	parser.add_argument('--constant_velocity', help='Run CV comparison', type=sup.str2bool,
											default=False)
	args = parser.parse_args()

	return args


test_args = parse_args()

if test_args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

cwd = os.getcwd()

model_path = os.path.normpath(cwd + '/../') + '/trained_models/' + test_args.model_name + "/" + str(test_args.exp_num)

print("Loading data from: '{}'".format(model_path))
file = open(model_path + '/model_parameters.pkl', 'rb')
if sys.version_info[0] < 3:
	model_parameters = pkl.load(file)  # ,encoding='latin1')
else:
	model_parameters = pkl.load(file, encoding='latin1')
file.close()
args = model_parameters["args"]

with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

print("Training steps: ", args.total_training_steps)

# change some args because we are doing inference
truncated_backprop_length = args.truncated_backprop_length
args.truncated_backprop_length = 8
args.batch_size = 16
args.keep_prob = 1.0

training_scenario = args.scenario
args.scenario = 'herengracht'
args.real_world_data = test_args.real_world_data
args.dataset = '/' + args.scenario + '.pkl'
# Specify the number of datasets / scenarios 
n_scenarios = 4

# Create Datahandler class for Herengracht experiment
original_scenario = args.scenario
data_prep = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.081,
				"map_size": np.array([1715,881]),
				"map_center": np.array([-78,-40])}
data_prep.processData(**map_args)

# Create Datahandler class for Prinsengracht scenario
args.scenario = "prinsengracht"
data_prep_2 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.091,
				"map_size": np.array([816,816]),
				"map_center": np.array([-37.1, -37.1])} 
data_prep_2.processData(**map_args)

args.scenario = "bloemgracht"
data_prep_3 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.22,
				"map_size": np.array([1469,928]),
				"map_center": np.array([-161.59,-102.08])} 
data_prep_3.processData(**map_args)

args.scenario = "open_crossing"
data_prep_4 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.6,
				"map_size": np.array([100,100]),
				"map_center": np.array([-30,-30])} 
data_prep_4.processData(**map_args)

args.scenario = original_scenario

# Initialize main data handler
dpreps = [data_prep, data_prep_2, data_prep_3, data_prep_4]
data_roboat = dhroboat.DataHandlerRoboat(args)

# Import Deep Learning model
#module = importlib.import_module("src.models."+args.model_name)
#globals().update(module.__dict__)

# Create Model Graph
#model = NetworkModel(args)

#config = tf.ConfigProto()

def visualizeBatch(dpreps, data_log):
	dset = data_log[0]
	traj_idx = data_log[1]

	batch_size = 16
	for i in range(batch_size):

		fig, ax = plt.subplots()
		sup.plot_grid_roboat(ax, dpreps[dset[i]].agent_container.occupancy_grid.center, dpreps[dset[i]].agent_container.occupancy_grid.gridmap, dpreps[dset[i]].agent_container.occupancy_grid.resolution, dpreps[dset[i]].agent_container.occupancy_grid.map_size)
		_, traj = dpreps[dset[i]].trajectory_set[traj_idx[i]]
		otherv = traj.other_agents_positions
		ax.plot(traj.pose_vec[:,0], traj.pose_vec[:,1])

		for t in range(len(otherv)-1):
			for v in range(len(otherv[t])):
				ax.scatter(otherv[t][v][0], otherv[t][v][1])

		plt.show()
		print("Scenario: " + i + " / " + batch_size)
	
	print("done")

def plotBatch(dpreps, data_log):
	dset = data_log[1]
	agent_ids = data_log[0]

	batch_size = 16
	for i in range(batch_size):

		fig, ax = plt.subplots()
		sup.plot_grid_roboat(ax, dpreps[dset[i]].agent_container.occupancy_grid.center, dpreps[dset[i]].agent_container.occupancy_grid.gridmap, dpreps[dset[i]].agent_container.occupancy_grid.resolution, dpreps[dset[i]].agent_container.occupancy_grid.map_size)
		traj = dpreps[dset[i]].agent_container.getAgentTrajectories(agent_ids[i])[0] 
		otherv = traj.other_agents_positions
		ax.plot(traj.pose_vec[:,0], traj.pose_vec[:,1])

		for t in range(len(otherv)-1):
			for v in range(len(otherv[t])):
				ax.scatter(otherv[t][v][0], otherv[t][v][1])

		plt.show()
		print("Scenario: " + str(i) + " / " + str(batch_size))


#rainLog = [[1, 3, 3, 1, 1, 3, 0, 3, 0, 2, 2, 0, 3, 3, 2, 2], [209, 222, 212, 223, 210, 217, 205, 220, 201, 202, 218, 207, 206, 204, 221, 211]]
#ValidationLog = [[1, 1, 2, 3, 0, 3, 3, 1, 0, 1, 2, 1, 0, 2, 3, 3], [3558, 3513, 2532, 3549, 2550, 1955, 3477, 3567, 1962, 3576, 2556, 3585, 3531, 3540, 2993, 1969]]

#TrainLog = [[1410, 1132, 3391, 1571, 3260, 1369, 3511, 2476, 1769, 2167, 2772, 1312, 1615, 602, 946, 2301], 3]

# Outlier 1
#ValidationLog = [[1155, 858, 103, 1830, 2998, 243, 275, 1327, 2286, 1954, 110, 457, 552, 2466, 2421, 811], [1, 2, 0, 1, 1, 2, 3, 2, 2, 3, 3, 0, 3, 1, 1, 2]]

# Outlier 2
ValidationLog = [[1968, 1037, 1138, 2527, 39, 244, 3737, 588, 3606, 356, 547, 1266, 3217, 2023, 3364, 2903], [0, 3, 3, 1, 3, 3, 1, 3, 1, 1, 1, 0, 2, 1, 1, 1]]


# Get Batch of Data
#train_dict = data_roboat.getPredefinedBatch(dpreps, TrainLog)
#validation_dict = data_roboat.getPredefinedBatch(dpreps, ValidationLog)

#visualizeBatch(dpreps, TrainLog)
#visualizeBatch(dpreps, ValidationLog)
plotBatch(dpreps, ValidationLog)





build/lib/inspect_outliers.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os
from colorama import Fore, Style

class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    self.grid_width, self.grid_height], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, 12,
		                                                        self.pedestrian_vector_dim*self.args.n_other_agents], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.diversity_placeholder = tf.placeholder(dtype=tf.float32,
		                                            shape=[self.batch_size, self.truncated_backprop_length,
		                                                   self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')
		self.seq_length = tf.placeholder(tf.int32, [None])
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped_fw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped_fw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_ped_bw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped_bw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_cell_state_current_lstm_ped_fw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped_fw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_ped_bw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped_bw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		self.test_natch_ped_grid = np.zeros([args.batch_size, 12, self.pedestrian_vector_dim*self.args.n_other_agents])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped_fw = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped_fw')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped_fw = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped_fw')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped_fw = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped_fw, self.hidden_state_lstm_ped_fw)
		self.cell_state_lstm_ped_bw = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped_bw')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped_bw = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped_bw')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped_bw = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped_bw, self.hidden_state_lstm_ped_bw)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)
		outputs_diversity_series = tf.unstack(self.diversity_placeholder, axis=1)
		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[ped_grid_out_dim*2, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
					                                         trainable=True)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped_fw = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell_fw', trainable=True)
					self.cell_ped_bw = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell_fw',
					                                           trainable=True)
					ped_grid_feature_series , self.current_state_lstm_ped = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.cell_ped_fw,
					                                                                                        cell_bw=self.cell_ped_bw,
					                                                                                        inputs=tf.stack(ped_grid_series,axis=1),
					                                                                                        dtype=tf.float32,
					                                                                                        sequence_length=self.seq_length,
					                                                                                        initial_state_fw=self.init_state_tuple_lstm_ped_fw,
					                                                                                        initial_state_bw=self.init_state_tuple_lstm_ped_bw)

					# Embedding layer of pedestrian grid
					self.cell_outputs_series_lstm_ped = self.process_pedestrian_grid(tf.unstack(tf.stack(ped_grid_feature_series,axis=2),axis=1))

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state,
				                                               self.cell_outputs_series_lstm_grid,
				                                               self.cell_outputs_series_lstm_ped):
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

					# Concatenated LSTM layer
					with tf.variable_scope('generation_model') as scope:

						# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
						cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
						                             args.latent_space_size)
						self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
							cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

						# Output reshape
						enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

						# Decoder layers
						dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
						dec_output_layer = []
						dec_mux_series = []
						dec_muy_series = []
						dec_sigmax_series = []
						dec_sigmay_series = []
						dec_pi_series = []
						# Reusing weights for each truncated back propagation step
						for t_idx in range(len(dec_fc_layer)):
							dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, 50))
							dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100))
						for t_idx in range(len(dec_output_layer)):
							m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1)
							dec_mux_series.append(m_x)
							dec_muy_series.append(m_y)
							dec_sigmax_series.append(s_x)
							dec_sigmay_series.append(s_y)
						dec_pi_series = [dense for dense in dec_pi_series]

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
				                                                decay_rate=0.9, staircase=False)

				self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

				prediction_loss_list = []
				loss_list = []
				kl_loss_list = []
				self.div_loss = []
				self.prediction = []
				self.likelihood = []
				pred = []
				div_loss_over_truncated_back_prop = []

				# Compute Loss Function
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out, out_div in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series, dec_pi_series,
						    outputs_series, outputs_diversity_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(dec_pi, axis=1)

					self.likelihood.append(pi)
					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
						x_diversity = tf.reshape(out_div[:, idx_x], shape=(-1, 1))
						y_diversity = tf.reshape(out_div[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                 sigmax[prediction_step], sigmay[prediction_step],
						                                 x_data, y_data, pi)

						diversity_loss = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                        sigmax[prediction_step], sigmay[prediction_step],
						                                       x_diversity, y_diversity, pi)

						div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss_list.append(kl_loss)
					div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			self.prediction = tf.stack(self.prediction, axis=1)
			self.likelihood = tf.stack(self.likelihood, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
			self.total_loss = tf.reduce_mean(loss_list, axis=0)+args.diversity_update*self.div_loss + tf.reduce_mean(kl_loss_list, axis=0)*self.beta
			self.reconstruction_loss = tf.reduce_mean(loss_list, axis=0)
			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			self.loss_summary = tf.summary.scalar('loss', self.total_loss)
			self.kl_loss_summary = tf.summary.scalar('kl-loss', self.kl_loss)
			self.div_loss_summary = tf.summary.scalar('div_loss', self.div_loss)
			self.reconstruction_loss_summary = tf.summary.scalar('reconstruction_loss', self.reconstruction_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)

			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
			self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
			                                                       graph=tf.Session().graph)

			self.summary = tf.summary.merge_all()

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
			self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		n_other_agents = np.zeros([self.args.batch_size])
		for i,other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.step: kwargs['step'],
		        self.seq_length: n_other_agents,
		        self.diversity_placeholder: kwargs['batch_div'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
		        self.cell_state_lstm_ped_fw: self.cell_state_current_lstm_ped_fw,
		        self.hidden_state_lstm_ped_fw: self.hidden_state_current_lstm_ped_fw,
		        self.cell_state_lstm_ped_bw: self.cell_state_current_lstm_ped_bw,
		        self.hidden_state_lstm_ped_bw: self.hidden_state_current_lstm_ped_bw,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }


	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		n_other_agents = np.zeros([self.args.batch_size])
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		self.test_natch_ped_grid = np.roll(self.test_natch_ped_grid,shift=1,axis=1)
		self.test_natch_ped_grid[:,0] = kwargs["batch_ped_grid"][:, step]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
		        self.input_ped_grid_placeholder: self.test_natch_ped_grid,
		        self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
		        self.step: 0,
		        self.seq_length: n_other_agents,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
		        self.cell_state_lstm_ped_fw: self.test_cell_state_current_lstm_ped_fw,
		        self.hidden_state_lstm_ped_fw: self.test_hidden_state_current_lstm_ped_fw,
		        self.cell_state_lstm_ped_bw: self.test_cell_state_current_lstm_ped_bw,
		        self.hidden_state_lstm_ped_bw: self.test_hidden_state_current_lstm_ped_bw,
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def feed_val_dic(self, **kwargs):
		n_other_agents = np.zeros([self.args.batch_size])
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: kwargs["batch_vel"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.diversity_placeholder: kwargs['batch_div'],
		        self.output_placeholder: kwargs['batch_y'],
		        self.seq_length: n_other_agents,
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), kwargs["state_noise"]),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), kwargs["state_noise"]),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
		        self.cell_state_lstm_ped_fw: np.random.normal(self.test_cell_state_current_lstm_ped_fw.copy(), kwargs["ped_noise"]),
		        self.hidden_state_lstm_ped_fw: np.random.normal(self.test_hidden_state_current_lstm_ped_fw.copy(), kwargs["ped_noise"]),
		        self.cell_state_lstm_ped_bw: np.random.normal(self.test_cell_state_current_lstm_ped_bw.copy(),
		                                                   kwargs["ped_noise"]),
		        self.hidden_state_lstm_ped_bw: np.random.normal(self.test_hidden_state_current_lstm_ped_bw.copy(),
		                                                     kwargs["ped_noise"]),
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped_fw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped_fw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_ped_bw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped_bw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])


	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_cell_state_current_lstm_ped_fw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped_fw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_ped_bw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped_bw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def update_test_hidden_state(self):
		self.test_cell_state_current = self.cell_state_current.copy()
		self.test_hidden_state_current = self.hidden_state_current.copy()
		self.test_cell_state_current_lstm_grid = self.cell_state_current_lstm_grid.copy()
		self.test_hidden_state_current_lstm_grid = self.hidden_state_current_lstm_grid.copy()
		self.test_cell_state_current_lstm_ped_fw = self.cell_state_current_lstm_ped_fw.copy()
		self.test_hidden_state_current_lstm_ped_fw = self.hidden_state_current_lstm_ped_fw.copy()
		self.test_cell_state_current_lstm_ped_bw = self.cell_state_current_lstm_ped_bw.copy()
		self.test_hidden_state_current_lstm_ped_bw = self.hidden_state_current_lstm_ped_bw.copy()
		self.test_cell_state_current_lstm_concat = self.cell_state_current_lstm_concat.copy()
		self.test_hidden_state_current_lstm_concat = self.hidden_state_current_lstm_concat.copy()

	def validation_step(self, sess, feed_dict_validation,update = True):
		batch_loss, _model_prediction,_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, summary = sess.run([self.reconstruction_loss,
		                                                self.prediction,
		                                                self.current_state,
		                                                self.current_state_lstm_grid,
		                                                self.current_state_lstm_ped,
		                                                self.current_state_lstm_concat,
		                                                self.loss_summary],
		                                                feed_dict=feed_dict_validation)

		if update:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[0]
			self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[1]
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, summary, _model_prediction

	def train_step(self, sess, feed_dict_train, step=0):

		# Diversity Training
		if step > 2000 and self.args.diversity_update:
			dict = {"batch_vel": feed_dict_train[self.input_state_placeholder],
			        "batch_grid": feed_dict_train[self.input_grid_placeholder],
			        "batch_ped_grid": feed_dict_train[self.input_ped_grid_placeholder],
			        "batch_y": feed_dict_train[self.output_placeholder],
			        "batch_div": feed_dict_train[self.diversity_placeholder],
			        "step": step,
			        "state_noise": 0.5,
			        "grid_noise": 0.0,
			        "ped_noise": 0.5
			        }
			self.update_test_hidden_state()
			feed_test_dic = self.feed_val_dic(**dict)
			loss, _, y_model_pred = self.validation_step(sess, feed_test_dic, False)
			batch_y_varible = np.zeros((self.args.batch_size, self.args.truncated_backprop_length,
			                            self.args.prediction_horizon * self.args.output_dim))
			# Grount truth
			batch_y = feed_dict_train[self.output_placeholder]

			diversity_distance = 0
			diff_traj_index = 0
			# Find most different trajectory index
			for batch_id in range(self.args.batch_size):
				for t_id in range(self.args.truncated_backprop_length):
					for mix_id in range(self.args.n_mixtures):
						for pred_id in range(self.args.prediction_horizon):
							new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures+ mix_id]
							new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + mix_id]
							batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
								[new_batch_x, new_batch_y])
						if diversity_distance < np.linalg.norm(batch_y_varible[batch_id, t_id]-batch_y[batch_id, t_id]):
							diversity_distance = np.linalg.norm(batch_y_varible[batch_id, t_id] - batch_y[batch_id, t_id])
							diff_traj_index = mix_id
					# Saving most different traj
					for pred_id in range(self.args.prediction_horizon):
						new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + diff_traj_index]
						new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + diff_traj_index]
						batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
							[new_batch_x, new_batch_y])

			# Assemble feed dict for training
			feed_dict_train[self.diversity_placeholder] = batch_y_varible

		_, batch_loss, kl_loss, recons_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		 _current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
		                                                                                                  self.total_loss,
		                                                                                                  self.kl_loss,
		                                                                                                  self.reconstruction_loss,
		                                                                                                  self.current_state,
		                                                                                                  self.current_state_lstm_grid,
		                                                                                                  self.current_state_lstm_ped,
		                                                                                                  self.current_state_lstm_concat,
		                                                                                                  self.prediction,
		                                                                                                  self.summary,
		                                                                                                  self.learning_rate,
		                                                                                                  self.beta,
		                                                                                                  self.deconv1,
		                                                                                                  self.div_loss,
		                                                                                                  self.autoencoder_loss],
		                                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[0]
		self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[1]
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		# Train autoencoder if loss larger than threshold
		""""""
		if np.mean(autoencoder_loss) > 0.05:
			self.run_autoencoder(sess, feed_dict_train)
			self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

		out = {"batch_loss": recons_loss,
		       "model_predictions": _model_prediction,
		       "summary": _summary_str,
		       "learning_rate": lr,
		       "autoencoder_loss": autoencoder_loss,
		       "log loss": batch_loss,
		       "kl_loss": kl_loss,
		       "output_decoder": output_decoder
		       }

		return out


	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)


	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, _likelihood = sess.run([self.current_state,
		                                              self.current_state_lstm_grid,
		                                              self.current_state_lstm_ped,
		                                              self.current_state_lstm_concat,
		                                              self.prediction,
		                                              self.deconv1,
                                                  self.likelihood],
		                                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[0]
			self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[1]
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, _likelihood


	def warmstart_model(self, args, sess):
		# Restore whole model
		try:
			print('Loading session from "{}"'.format(args.model_path))
			ckpt = tf.train.get_checkpoint_state(args.model_path)
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.full_saver.restore(sess, ckpt.model_checkpoint_path)
		except:
			print(Fore.RED + "Model not found..."+Style.RESET_ALL)
			sys.exit()


	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print(Fore.RED +"Error. Encoder not found...")
			sys.exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True


	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)


	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)],
				                   1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss


	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result


	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y, pi):

		# gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		# gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		loss = tf.multiply(z + tf.log(sxsy), pi)

		return tf.reduce_mean(loss)


	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):
		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #


	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors


	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector


	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations


	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation


	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

build/lib/models/BiVRNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from irnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.irnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os


class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')

		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')

		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)


		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			with tf.variable_scope('encoder_model') as scope:

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state,
				                                                     ped_grid_series):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1,
					                             name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
				cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
				                             args.latent_space_size)
				self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
					cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# Output reshape
				enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
				enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
				dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
				prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
				prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

				# Decoder layers
				dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
				dec_fc_layer2 = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_fc_layer])))
				dec_output_layer = []
				dec_mux_series = []
				dec_muy_series = []
				dec_sigmax_series = []
				dec_sigmay_series = []
				dec_pi_series = []
				for t_idx in range(len(dec_fc_layer2)):
					dec_output_layer.append(linear(dec_fc_layer2[t_idx], self.output_vec_dim, t_idx))
				for t_idx in range(len(dec_output_layer)):
					if args.output_pred_state_dim > 4:
						m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
						dec_mux_series.append(m_x)
						dec_muy_series.append(m_y)
						dec_sigmax_series.append(s_x)
						dec_sigmay_series.append(s_y)
						dec_pi_series.append(likelihood)
					else:
						m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
						dec_mux_series.append(m_x)
						dec_muy_series.append(m_y)
						dec_sigmax_series.append(s_x)
						dec_sigmay_series.append(s_y)

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
				                                                decay_rate=0.9, staircase=False)

				self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

			prediction_loss_list = []
			loss_list = []
			kl_loss_list = []
			self.div_loss = []
			self.prediction = []
			pred = []
			div_loss_over_truncated_back_prop = []
			if args.output_pred_state_dim > 4:
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series, dec_pi_series,
						    outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
							pred = tf.concat([pred, pi[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
							pred = tf.concat([pred, pi[prediction_step]], 1)

						lossfunc = self.get_mixture_lossfunc(mux[prediction_step], muy[prediction_step],
						                                     sigmax[prediction_step], sigmay[prediction_step],
						                                     pi[prediction_step],
						                                     x_data, y_data)

						diversity_loss = self.get_diversity_loss(mux[prediction_step], muy[prediction_step],
						                                         sigmax[prediction_step], sigmay[prediction_step])

						div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
					kl_loss_list.append(kl_loss)
					div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			else:
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series,
						    outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

					sigmax = tf.maximum(tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1)), 0.2)
					sigmay = tf.maximum(tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1)), 0.2)

					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                 sigmax[prediction_step], sigmay[prediction_step],
						                                 x_data, y_data)

						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
					kl_loss_list.append(kl_loss)
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			# Reduce mean in all dimensions
			# self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
			# self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
			self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)
			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
			self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

	def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: batch_ped_grid,
		        self.step: step,
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, batch_x, batch_grid, batch_ped_grid, noise=0):
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: batch_ped_grid,
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def run(self, sess, feed_dict_train):
		_, batch_loss, kl_loss, _current_state, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, div_loss = sess.run([self.update,
		                                                                                                  self.total_loss,
		                                                                                                  self.kl_loss,
		                                                                                                  self.current_state,
		                                                                                                  self.current_state_lstm_concat,
		                                                                                                  self.prediction,
		                                                                                                  self.summary,
		                                                                                                  self.learning_rate,
		                                                                                                  self.beta,
		                                                                                                  self.div_loss],
		                                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state

		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, kl_loss, _model_prediction, _summary_str, lr, beta, 0, div_loss, 0

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state,	_current_state_lstm_concat, \
		_model_prediction = sess.run([self.current_state,
		                                              self.current_state_lstm_concat,
		                                              self.prediction],
		                                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, 0

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y):

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		# Normalization constant
		denom = tf.math.log(2 * np.pi * sxsy)
		# Final PDF calculation
		result = denom + z
		gaussian = tf.reduce_sum(result, axis=1)  # do inner summation

		return gaussian

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

build/lib/models/CVGDNN.py
eliatrevisan/IA-MPPI-LBM
import numpy as np

class ConstantVelocity():

    def __init__(self, args):
        self.args = args

    def predict(self, input_list):
        # Initialize array to save constant velocity predictions
        all_pred = []        
        for i in range(0, len(input_list)):
            input = input_list[i]
            pred_list = np.zeros([input.shape[0]], dtype=np.ndarray)
            for step in range(input.shape[0]):
                pred = self.cv(input[step,0], input[step,1], input[step,2], input[step,3]) 
                pred_list[step] = pred
            all_pred.append(pred_list)
        return np.array(all_pred)

    def cv(self, posx, posy, vx, vy):
        pred = []

        for i in range(self.args.prediction_horizon):
            posx = posx + vx*self.args.dt
            posy = posy + vy*self.args.dt
            pred.append([posx, posy])
        return np.array(pred)




build/lib/models/CV_model.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
	import src.data_utils.Support as sup
import numpy as np
import os


def length(sequence):
	used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
	length = tf.reduce_sum(used, 1)
	length = tf.cast(length, tf.int32)
	return length


class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_ped_size = args.rnn_state_ped_size
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.n_peds = self.args.n_other_agents
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    self.grid_width, self.grid_height], name='input_grid')
		if self.args.relative_info:
			self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
			                                                 shape=[self.batch_size, self.truncated_backprop_length,
			                                                        self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
		if "prev" in self.args.others_info:
			self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
			                                                 shape=[self.batch_size, self.truncated_backprop_length,
			                                                        self.n_peds,
			                                                        self.pedestrian_vector_dim * (self.prev_horizon + 1)],
			                                                 name='ped_grid')
		else:
			self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
			                                                 shape=[self.batch_size, self.truncated_backprop_length,
			                                                        self.n_peds,
			                                                        self.pedestrian_vector_dim * self.prediction_horizon],
			                                                 name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32, shape=[], name='global_step')
		self.seq_len = tf.placeholder(tf.int32, [None])

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.ped_gru_states_current = np.zeros([args.batch_size, self.n_peds, self.rnn_state_ped_size])
		self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_ped_gru_states_current = np.zeros([args.batch_size, self.n_peds, self.rnn_state_ped_size])
		self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.ped_gru_states = tf.placeholder(dtype=tf.float32,
		                                     shape=[self.batch_size, self.n_peds, self.rnn_state_ped_size],
		                                     name='cell_state_lstm_ped')
		self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                    name='cell_state_lstm_ped')
		self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                    name='cell_state_lstm_ped')
		# self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		# self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		# self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
					                                         trainable=train_encoder)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

				# LSTM cells to process pedestrian
				self.cell_outputs_series_lstm_ped = []
				self.current_state_lstm_ped = []
				self.other_peds_states = []
				self.other_peds_gru_outputs = []
				with tf.variable_scope('lstm_ped', reuse=tf.AUTO_REUSE) as scope:
					# TODO: try to share weights
					for i, ped in enumerate(ped_grid_series):
						# TODO: ADAPT for ped input sequence
						ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell' + str(i),
						                                  trainable=True)
						ped_input = tf.unstack(ped, axis=1)
						ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell, ped_input, dtype=tf.float32,
						                                                          initial_state=self.ped_gru_states[:, i, :])
						self.other_peds_states.append(ped_gru_state)

						if i > 0:
							for j, ped_out in enumerate(ped_gru_output):
								out_ped_concat[j] = tf.concat([out_ped_concat[j], ped_gru_output[j]], axis=1)
						else:
							out_ped_concat = ped_gru_output
					self.other_peds_gru_outputs = out_ped_concat

					self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
					                                          trainable=True)
					self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
					                                          trainable=True)

					outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
					                                                  cell_bw=self.fw_cell_ped,
					                                                  inputs=tf.stack(self.other_peds_gru_outputs, axis=1),
					                                                  dtype=tf.float32,
					                                                  initial_state_fw=self.fw_ped_states,
					                                                  initial_state_bw=self.bw_ped_states)

					self.cell_outputs_series_lstm_ped = tf.unstack(
						tf.transpose(tf.concat([outputs[0], outputs[1]], axis=2), [1, 0, 2]), axis=0)
					self.current_state_lstm_ped = states
					self.other_peds_states = tf.stack(self.other_peds_states, axis=1)

					rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
					                                                   shape=[self.rnn_state_size_lstm_ped * 2,
					                                                          self.rnn_state_size_lstm_ped],
					                                                   trainable=self.is_training)
					b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
					                                 trainable=self.is_training)

					self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
					              self.cell_outputs_series_lstm_ped]

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.fused):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out',
				                                 shape=[self.fc_hidden_unit_size, self.output_dim * self.prediction_horizon],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
				                               trainable=self.is_training)

				# Second concatenated LSTM layer
				with tf.variable_scope('lstm_concat') as scope:
					cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# FC layers
				fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
				           self.cell_outputs_series_lstm_concat]
				# fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
				#           fc_fuse]
				# fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
				# Model prediction
				self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

			# Exponential / stepwise decrease of learning rate
			global_step = tf.Variable(0, trainable=False)
			self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
			                                                decay_rate=0.9, staircase=False)

			# Loss function
			loss_list = []
			lossfunc = []
			# Iterate through tbp steps
			for pred, out in zip(self.prediction,
			                     outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
				prediction_loss_list = []
				val_prediction_loss_list = []
				for prediction_step in range(self.prediction_horizon):
					# Compute Euclidean error for each prediction step within each tbp step

					idx_x = prediction_step * self.output_dim
					idx_y = idx_x + 1

					# Calculate loss for the current ped
					lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
					                          tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

					prediction_loss_list.append(lossfunc)  # sum over number of batches

				loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

			self.prediction = tf.stack(self.prediction, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0)  # + l2 # mean over tbp steps

			# Optimizer specification
			self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			# self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

		# Only save / restore the parameters responsible for occupancy grid processing
		self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
		self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)

	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		return {self.input_state_placeholder: batch_x,
		        self.input_grid_placeholder: kwargs['batch_grid'],
		        self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
		        self.step: kwargs['step'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
		        self.ped_gru_states: self.ped_gru_states_current,
		        self.fw_ped_states: self.fw_ped_states_current,
		        self.bw_ped_states: self.bw_ped_states_current,
		        self.seq_len: [100],
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
		        self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
		        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
		        self.ped_gru_states: self.test_ped_gru_states_current,
		        self.fw_ped_states: self.test_fw_ped_states_current,
		        self.bw_ped_states: self.test_bw_ped_states_current,
		        self.seq_len: [100],
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.ped_gru_states_current = np.zeros([self.args.batch_size, self.n_peds, self.args.rnn_state_ped_size])
					self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
					self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_ped_gru_states_current = np.zeros([self.args.batch_size, self.n_peds, self.args.rnn_state_ped_size])
					self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
					self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def run(self, sess, feed_dict_train):
		_, batch_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, _other_peds_states, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
		                                                                                  self.total_loss,
		                                                                                  self.current_state,
		                                                                                  self.current_state_lstm_grid,
		                                                                                  self.current_state_lstm_ped,
		                                                                                  self.other_peds_states,
		                                                                                  self.current_state_lstm_concat,
		                                                                                  self.prediction,
		                                                                                  self.summary,
		                                                                                  self.learning_rate,
		                                                                                  self.deconv1,
		                                                                                  self.autoencoder_loss],
		                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
		self.ped_gru_states_current = _other_peds_states
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, 0, _model_prediction, _summary_str, lr, 0, output_decoder, 0, 0, 0, 0

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, _other_peds_states, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, \
		_cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
		_cell_outputs_series_lstm_concat = sess.run([self.current_state,
		                                             self.current_state_lstm_grid,
		                                             self.current_state_lstm_ped,
		                                             self.other_peds_states,
		                                             self.current_state_lstm_concat,
		                                             self.prediction,
		                                             self.deconv1,
		                                             self.cell_outputs_series_state,
		                                             self.cell_outputs_series_lstm_ped,
		                                             self.cell_outputs_series_lstm_concat
		                                             ],
		                                            feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_fw_ped_states_current, self.test_bw_ped_states_current = _current_state_lstm_ped
			self.test_ped_gru_states_current = _other_peds_states
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, output_decoder, _cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		try:
			if self.args.regulate_log_loss:
				likelihood_loss = -1.0 * tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)), 4.0)
			else:
				likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		except:
			likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

build/lib/models/IntNet.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.prediction_horizon,
                                                              self.pedestrian_vector_dim],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []

        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_series = tf.unstack(ped, axis=1)
            ped_states = []
            # iterate over truncated back propagation
            for j, ped_t in enumerate(ped_series):
              ped_input = tf.unstack(ped_t, axis=1)
              ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
              ped_states.append(ped_gru_state)
            self.other_peds_states.append(ped_states)
          a = tf.transpose(tf.concat(tf.unstack(tf.stack(self.other_peds_states),axis=0),axis=2),[1,0,2])
          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=a,
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.squeeze(tf.stack(self.other_peds_states,axis=2),axis=0)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder, 0, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntNet2.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_size_lstm_ped])
    self.ped_dec_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_size_lstm_ped])
    self.test_ped_dec_gru_states_current = np.zeros([args.batch_size, self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.ped_dec_gru_states = tf.placeholder(dtype=tf.float32,
                                         shape=[self.batch_size, self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon],
                                         name='cell_dec_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='fw_cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='bw_cell_state_lstm_ped')
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)

    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:
      with tf.variable_scope('encoder_model') as scope:
        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_dec_states = []
        self.other_peds_gru_outputs = []
        out_ped_concat = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          w_enc_state_lstm_ped = self.get_weight_variable(name='w_enc_state_lstm_ped',
                                                          shape=[self.rnn_state_size_lstm_ped,
                                                                 self.rnn_state_ped_size],
                                                          trainable=self.is_training)
          b_enc_state_lstm_ped = self.get_bias_variable(name='b_enc_state_lstm_ped', shape=[1, self.rnn_state_ped_size],
                                                        trainable=self.is_training)

          b_dec_state_lstm_ped = self.get_bias_variable(name='b_dec_state_lstm_ped',
                                                        shape=[1, self.rnn_state_size_lstm_ped],
                                                        trainable=self.is_training)
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_enc_cell = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='ped_enc_lstm_cell'+str(i), trainable=True)
            ped_dec_cell = tf.contrib.rnn.GRUCell(self.pedestrian_vector_dim * self.prediction_horizon, name='ped_dec_lstm_cell' + str(i),trainable=True)

            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_enc_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            ped_gru_output_fc = tf.add(tf.matmul(ped_gru_output[0], w_enc_state_lstm_ped), b_enc_state_lstm_ped)

            ped_dec_fc = tf.add(tf.matmul(ped_gru_output_fc, tf.transpose(w_enc_state_lstm_ped)), b_dec_state_lstm_ped)

            ped_dec_gru_output, ped_dec_gru_state = tf.contrib.rnn.static_rnn(ped_dec_cell,[ped_dec_fc],dtype=tf.float32,initial_state=self.ped_dec_gru_states[:,i,:])

            self.other_peds_dec_states.append(ped_dec_gru_state)

            ped_loss = tf.reduce_mean(tf.sqrt(tf.square(tf.subtract(ped_dec_gru_output[0], ped_input[0]))),axis=1)

            out_ped_concat.append(ped_gru_output_fc)

          self.ped_enc_dec_loss = tf.reduce_mean(ped_loss)

          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)
          self.other_peds_dec_states = tf.stack(self.other_peds_dec_states, axis=1)

          w_birnn = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, w_birnn), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.fused):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        #fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
        #           fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction, axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      #self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)
      self.ped_dec_enc_gradients = tf.gradients(self.ped_enc_dec_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
      self.ped_dec_enc_gradients_clipped, _ = tf.clip_by_global_norm(self.ped_dec_enc_gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      self.update_ped_dec_enc = self.optimizer.apply_gradients(zip(self.ped_dec_enc_gradients_clipped, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('Model Loss', self.total_loss)
      tf.summary.scalar('Pedestrian Encode Decoder Loss', self.ped_enc_dec_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      batch_x = kwargs['batch_vel']
      batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.ped_dec_gru_states: self.ped_dec_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.ped_dec_gru_states: self.test_ped_dec_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_size_lstm_ped])
          self.ped_dec_gru_states_current = np.zeros([self.args.batch_size, self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_size_lstm_ped])
          self.test_ped_dec_gru_states_current = np.zeros([self.args.batch_size, self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_ped,_other_peds_states, \
    _other_peds_dec_states, _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, _ped_dec_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.other_peds_dec_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.ped_enc_dec_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.ped_dec_gru_states_current = _other_peds_dec_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, 0,0, 0, 0,_ped_dec_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)
  def run_ped_encoder(self,sess,feed_dict_train):
    _ = sess.run([self.update_ped_dec_enc],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_ped,_other_peds_states, \
    _other_peds_dec_states ,_current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.other_peds_dec_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_fw_ped_states_current, self.test_bw_ped_states_current = _current_state_lstm_ped
      self.test_ped_gru_states_current = _other_peds_states
      self.test_ped_dec_gru_states_current = _other_peds_dec_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntNet3.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')

    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)
        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        #fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
        #           fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction, axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      #self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('Model Loss', self.total_loss)
      tf.summary.scalar('Grid Encode Loss', self.autoencoder_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      tf.summary.image('Pedestrian Grid', self.deconv1)

      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      batch_x = kwargs['batch_vel']
      batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :, :], axis=0),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state,_cell_outputs_series_lstm_grid, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_grid,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_grid, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntNetGrid.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.pos_output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    pos_outputs_series = tf.unstack(self.pos_output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      pos_loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out, pos_out in zip(self.prediction,outputs_series,pos_outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        pos = 0.0
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          pos += pred[:,idx_x:idx_x+self.output_dim] * self.args.dt

          prediction_loss_list.append(lossfunc)  # sum over number of batches
        pos_loss = tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(pos,pos_out)),axis=1))
        pos_loss_list.append(tf.reduce_mean(pos_loss))  # mean over prediction horizon
        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0) +tf.reduce_mean(pos_loss_list, axis=0)*0.5 #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y, batch_pos_final_backup):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.pos_output_placeholder: batch_pos_final_backup,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntNet_end_pos_loss.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.math.abs(tf.reduce_mean(loss_list, axis=0)-0.05)+0.05
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      batch_x = kwargs['batch_x'][:,:,2:]
      batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntNet_flood.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*self.prediction_horizon], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    elif "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction,axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    batch_x = kwargs['batch_y_pred']
    batch_x[:,:,:2] = kwargs['batch_vel'][:,:,:2]
    batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntNet_future.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.fused):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        #fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
        #           fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          error_x = tf.maximum(tf.abs(tf.subtract(pred[:, idx_x], out[:, idx_x])),0.03)
          error_y = tf.maximum(tf.abs(tf.subtract(pred[:, idx_y], out[:, idx_y])),0.03)
          lossfunc = tf.sqrt(tf.add(tf.square(error_x),  # error in x
                                    tf.square(error_y)))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction, axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      #self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      batch_x = kwargs['batch_vel']
      batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntNet_minloss.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * (self.prev_horizon+1)],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder, 0, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntNetprev.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell3 import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell3 import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    #self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                     shape=[self.batch_size, self.truncated_backprop_length,
                                                            self.args.n_other_agents,
                                                            self.pedestrian_vector_dim * self.prediction_horizon],
                                                     name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    #self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    #self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    #self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    #self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    #self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    #self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    #conv_grid_feature_series = []
    #for output in self.process_grid(self.occ_grid_series):
    #  conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        #with tf.variable_scope('lstm_grid') as scope:
        #  self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
        #                                           trainable=True)
        #  self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
        #                                                                                               conv_grid_feature_series,
        #                                                                                               dtype=tf.float32,
        #                                                                                              initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped', reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.args.rnn_state_ped_size, name='ped_basic_lstm_cell' + str(i),
                                              trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell, ped_input, dtype=tf.float32,
                                                                      initial_state=self.ped_gru_states[:, i, :])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j], ped_gru_output[j]], axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                            cell_bw=self.fw_cell_ped,
                                                            inputs=tf.stack(self.other_peds_gru_outputs, axis=1),
                                                            dtype=tf.float32,
                                                            initial_state_fw=self.fw_ped_states,
                                                            initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(
            tf.transpose(tf.concat([outputs[0], outputs[1]], axis=2), [1, 0, 2]), axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states, axis=1)
        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args,args.rnn_state_size_lstm_concat,self.output_vec_dim, args.prior_size, args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step],muy[prediction_step],sigmax[prediction_step],sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    #self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    #self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            #self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            #self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,noise=0):
    return {self.input_state_placeholder: batch_x,
            #self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            #self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            #self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          #self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          #self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          #self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          #self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          #self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          #self.deconv1,
                                                          #self.autoencoder_loss
                                                          ],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    #self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, 0,0, 0, 0,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction,\
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped,\
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              #self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              #self.deconv1
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.dec_series],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      #self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, 0, _cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/IntVGDNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
  from utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  from src.data_utils.utils import *
import numpy as np
import os
from datetime import datetime
from colorama import Fore, Style

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)

    self.input_grid = tf.expand_dims(input=self.occ_grid_series[0], axis=3)

    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction, axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)

      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Tensorboard data
      self.loss_summary = tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)

      # Grid input
      sliced_grid = tf.slice(self.input_grid, begin=[0,0,0,0], size=[1, self.grid_width, self.grid_height, 1], name="slice_in_grid")
      tf.summary.image('input_grid', sliced_grid, max_outputs=1)

      # Grid output
      sliced_grid = tf.slice(self.deconv1, begin=[0,0,0,0], size=[1, self.grid_width, self.grid_height, 1], name="slice_out_grid")
      tf.summary.image('output_grid', sliced_grid, max_outputs=1)

      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
      self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'), graph=tf.Session().graph)

      self.summary = tf.summary.merge_all()

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    return {self.input_state_placeholder: kwargs['batch_vel'],
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: kwargs['batch_y'],
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_val_dic(self, **kwargs):
    return {self.input_state_placeholder: kwargs['batch_vel'],
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.output_placeholder: kwargs['batch_y'],
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, **kwargs):
    step = kwargs["step"]
    return {self.input_state_placeholder:  np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
            self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :,:], axis=1),
            self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),kwargs["state_noise"]),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),kwargs["state_noise"]),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),kwargs["grid_noise"]),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),kwargs["grid_noise"]),
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def validation_step(self,sess,feed_dict_validation):
    batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, summary = sess.run([self.total_loss,
                                           self.current_state,
                                           self.current_state_lstm_grid,
                                           self.current_state_lstm_ped,
                                           self.current_state_lstm_concat,
                                           self.loss_summary],
                                           feed_dict=feed_dict_validation)
    self.test_cell_state_current, self.test_hidden_state_current = _current_state
    self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, summary

  def train_step(self,sess,feed_dict_train,step=0):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    # Train autoencoder if loss larger than threshold
    if np.mean(autoencoder_loss) > 0.05:
      print(Fore.LIGHTMAGENTA_EX +"\n[{:s}] Training autoencoder. Model loss: {:01.2f}" .format(
          datetime.now().strftime("%d-%m-%y %H:%M:%S"), autoencoder_loss))

      self.run_autoencoder(sess, feed_dict_train)
      self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

    out = {"batch_loss": batch_loss,
            "model_predictions": _model_prediction,
            "summary": _summary_str,
            "learning_rate": lr,
            "autoencoder_loss": autoencoder_loss,
            "output_decoder": output_decoder
            }

    return out

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/RNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.goal_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, 2], name='goal')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    self.goal_series = tf.unstack(self.goal_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=True)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat+2, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(tf.concat([lstm_out,goal],axis=1), W_fc), b_fc)) for lstm_out,goal in
                   zip(self.cell_outputs_series_lstm_concat,self.goal_series)]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y,batch_goal):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.goal_placeholder: batch_goal,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_goal,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.goal_placeholder: batch_goal,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder, 0, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/RNN_goal.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))

    # Process Pedestrian occupancy grid, apply convolutional filters
    self.ped_occ_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
    ped_grid_feature_series = []
    with tf.variable_scope('ped_encoder_model') as scope:
      for output in self.process_ped_grid(self.ped_occ_grid_series):
        ped_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:

        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)

    self.ped_convnet_saver = tf.train.Saver(self.ped_encoder_var_list)
    self.ped_encoder_update = self.optimizer.minimize(self.ped_encoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder,output_ped_decoder,\
    autoencoder_loss,ped_encoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.ped_deconv1,
                                                          self.autoencoder_loss,
                                                          self.ped_encoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,output_ped_decoder, 0, autoencoder_loss,ped_encoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def run_ped_encoder(self,sess,feed_dict_train):
    _ = sess.run([self.ped_encoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)
    print('Loading Pedestrian convnet parameters from "{}"'.format(args.model_path+"/ped_encoder"))
    ckpt_conv = tf.train.get_checkpoint_state(args.model_path+"/ped_encoder")
    if ckpt_conv == None:
      print("Error. Pedestrian Convnet not found. Loading default")
      ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    print('Restoring Pedestrian convnet {}'.format(ckpt_conv.model_checkpoint_path))
    self.ped_convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope("auto_encoder") as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=True)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="auto_encoder")

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="auto_encoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors

  def process_ped_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope("auto_encoder") as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=True)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases", trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                 shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                       [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                       padding='VALID', trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                       [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                       padding='SAME', trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.ped_deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME', trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.ped_deconv1.get_shape()))

          self.ped_encoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="auto_encoder")

          self.ped_encoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.ped_deconv1)), name="auto_encoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/RNN_grid.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight

    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.goal_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, 2], name='goal')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])

    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    self.goal_series = tf.unstack(self.goal_placeholder, axis=1)
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat+2, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(tf.concat([lstm_out,goal],axis=1), W_fc), b_fc)) for lstm_out,goal in
                   zip(self.cell_outputs_series_lstm_concat,self.goal_series)]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y,batch_goal):
    return {self.input_state_placeholder: batch_x,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.goal_placeholder: batch_goal,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_goal,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.goal_placeholder: batch_goal,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, 0, 0, 0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction = sess.run([self.current_state,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, 0

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

build/lib/models/RNN_no_grid_goal.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os


class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight

		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.goal_placeholder = tf.placeholder(dtype=tf.float32,
		                                       shape=[self.batch_size, self.truncated_backprop_length, args.goal_size],
		                                       name='goal')

		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32, shape=[], name='global_step')

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)
		self.goal_series = tf.unstack(self.goal_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
				# Embedding layer of pedestrian grid
				ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped,
					                                                                                           ped_grid_feature_series,
					                                                                                           dtype=tf.float32,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat+args.goal_size, self.rnn_state_size_lstm_concat+args.goal_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.rnn_state_size_lstm_concat+args.goal_size],
				                              trainable=self.is_training)
				W_fc2 = self.get_weight_variable(name='w_fc2',
				                                shape=[self.rnn_state_size_lstm_concat+args.goal_size, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc2 = self.get_bias_variable(name='b_fc2', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out',
				                                 shape=[self.fc_hidden_unit_size, self.output_dim * self.prediction_horizon],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
				                               trainable=self.is_training)

				# Second concatenated LSTM layer
				with tf.variable_scope('lstm_concat') as scope:
					cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# Concatenate the goal information
				concat_goal = []
				for lstm_out, goal_feature in zip(self.cell_outputs_series_lstm_concat, self.goal_series):
					concat_goal.append(tf.concat([lstm_out, goal_feature], axis=1, name="concatenate_lstm_out_with_goal_info"))

				# FC layers
				fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
				           concat_goal]
				#fc_drop = [tf.layers.dropout(inputs=dense, rate=1 - self.args.keep_prob) for dense in
				 #          fc_fuse]
				fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
				# Model prediction
				self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse2]

			# Exponential / stepwise decrease of learning rate
			global_step = tf.Variable(0, trainable=False)
			self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
			                                                decay_rate=0.9, staircase=False)

			# Loss function
			loss_list = []
			lossfunc = []
			# Iterate through tbp steps
			for pred, out, goal in zip(self.prediction,
			                     outputs_series,self.goal_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
				prediction_loss_list = []
				val_prediction_loss_list = []
				position = 0
				for prediction_step in range(self.prediction_horizon):
					# Compute Euclidean error for each prediction step within each tbp step

					idx_x = prediction_step * self.output_dim
					idx_y = idx_x + 1

					# Calculate loss for the current ped
					lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
					                          tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

					prediction_loss_list.append(lossfunc)  # sum over number of batches

					position += pred[:, idx_x:idx_x+2]*args.dt

				pos_lossfunc = tf.norm(tf.norm(position)-goal[:,0])+tf.norm(tf.math.atan2(position[:, 1],position[:, 0])-goal[:,1])

				loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

			self.prediction = tf.stack(self.prediction, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0)  # + l2 # mean over tbp steps

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.goal_placeholder: kwargs["batch_goal"],
		        self.step: kwargs['step'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
	        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
	        self.goal_placeholder: np.expand_dims(kwargs["batch_goal"][:, step, :], axis=1),
	        self.step: 0,
	        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
	        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
	        # self.cell_state: self.cell_state_current,
	        # self.hidden_state: self.hidden_state_current,
	        # self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
	        # self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
	        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
	        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
	        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
	        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
	        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def run(self, sess, feed_dict_train):
		_, batch_loss, _current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr = sess.run([self.update,
		                                                self.total_loss,
		                                                self.current_state,
		                                                self.current_state_lstm_ped,
		                                                self.current_state_lstm_concat,
		                                                self.prediction,
		                                                self.summary,
		                                                self.learning_rate],
		                                               feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, 0, _model_prediction, _summary_str, lr, 0, 0, 0, 0, 0, 0

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update], feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction = sess.run([self.current_state,
		                              self.current_state_lstm_ped,
		                              self.current_state_lstm_concat,
		                              self.prediction],
		                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, 0, _current_state, _current_state_lstm_ped, _current_state_lstm_concat

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		try:
			if self.args.regulate_log_loss:
				likelihood_loss = -1.0 * tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)), 4.0)
			else:
				likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		except:
			likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

build/lib/models/RNN_plannet.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os


class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight

		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32, shape=[], name='global_step')

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
				# Embedding layer of pedestrian grid
				ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped,
					                                                                                           ped_grid_feature_series,
					                                                                                           dtype=tf.float32,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out',
				                                 shape=[self.fc_hidden_unit_size, self.output_dim * self.prediction_horizon],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
				                               trainable=self.is_training)

				# Second concatenated LSTM layer
				with tf.variable_scope('lstm_concat') as scope:
					cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# FC layers
				fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
				           self.cell_outputs_series_lstm_concat]
				fc_drop = [tf.layers.dropout(inputs=dense, rate=1 - self.args.keep_prob) for dense in
				           fc_fuse]
				# fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
				# Model prediction
				self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

			# Exponential / stepwise decrease of learning rate
			global_step = tf.Variable(0, trainable=False)
			self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
			                                                decay_rate=0.9, staircase=False)

			# Loss function
			loss_list = []
			lossfunc = []
			# Iterate through tbp steps
			for pred, out in zip(self.prediction,
			                     outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
				prediction_loss_list = []
				val_prediction_loss_list = []
				for prediction_step in range(self.prediction_horizon):
					# Compute Euclidean error for each prediction step within each tbp step

					idx_x = prediction_step * self.output_dim
					idx_y = idx_x + 1

					# Calculate loss for the current ped
					lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
					                          tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

					prediction_loss_list.append(lossfunc)  # sum over number of batches

				loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

			self.prediction = tf.stack(self.prediction, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0)  # + l2 # mean over tbp steps

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.step: kwargs['step'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
	        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
	        self.step: 0,
	        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
	        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
	        # self.cell_state: self.cell_state_current,
	        # self.hidden_state: self.hidden_state_current,
	        # self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
	        # self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
	        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
	        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
	        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
	        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
	        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def run(self, sess, feed_dict_train):
		_, batch_loss, _current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr = sess.run([self.update,
		                                                self.total_loss,
		                                                self.current_state,
		                                                self.current_state_lstm_ped,
		                                                self.current_state_lstm_concat,
		                                                self.prediction,
		                                                self.summary,
		                                                self.learning_rate],
		                                               feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, 0, _model_prediction, _summary_str, lr, 0, 0, 0, 0, 0, 0

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update], feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction = sess.run([self.current_state,
		                              self.current_state_lstm_ped,
		                              self.current_state_lstm_concat,
		                              self.prediction],
		                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, 0, _current_state, _current_state_lstm_ped, _current_state_lstm_concat

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		try:
			if self.args.regulate_log_loss:
				likelihood_loss = -1.0 * tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)), 4.0)
			else:
				likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		except:
			likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

build/lib/models/RNN_simple.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import numpy as np
import os
import sys
sys.path.append('./autoencoder_grid')
tfk = tf.keras
tfkl = tf.keras.layers

class EncoderState(tfk.Model):
  def __init__(self, rnn_state_size,truncated_backprop_length, input_state_dim,batch_size):
    super(EncoderState, self).__init__()

    self.rnn_state_size = rnn_state_size
    self.truncated_backprop_length = truncated_backprop_length
    self.input_state_dim = input_state_dim
    self.batch_size = batch_size
    self.lambda_ = 0.01

    self.agent_state_lstm = tfkl.GRU(self.rnn_state_size, input_shape=[self.truncated_backprop_length, self.input_state_dim],
                                 name='ped_lstm',return_sequences=True,
                                 return_state=True,
                                 kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))

  def call(self, x, hidden):
    output, state = self.agent_state_lstm(x, initial_state = hidden)
    return output, state

  def initialize_hidden_state(self):
    return tf.zeros((self.batch_size, self.rnn_state_size))

class EncoderPeds(tfk.Model):
  def __init__(self, rnn_state_size_lstm_ped, truncated_backprop_length, input_state_dim, batch_size):
    super(EncoderPeds, self).__init__()

    self.rnn_state_size_lstm_ped = rnn_state_size_lstm_ped
    self.truncated_backprop_length = truncated_backprop_length
    self.input_state_dim = input_state_dim
    self.lambda_ = 0.01
    self.batch_size = batch_size
    self.ped_embed = tfkl.Dense(self.rnn_state_size_lstm_ped, activation='relu')

    self.agent_other_peds_lstm = tfkl.GRU(self.rnn_state_size_lstm_ped,
                                     input_shape=[self.truncated_backprop_length, self.input_state_dim],
                                     name='other_peds_lstm', return_sequences=True,
                                     return_state=True,
                                     kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))

  def call(self, x, hidden):
    x = self.ped_embed(x)
    output, state = self.agent_other_peds_lstm(x, initial_state=hidden)
    return output, state

  def initialize_hidden_state(self):
    return tf.zeros((self.batch_size, self.rnn_state_size_lstm_ped))

class BahdanauAttention(tfkl.Layer):
  def __init__(self, units):
    super(BahdanauAttention, self).__init__()
    self.W1 = tf.keras.layers.Dense(units)
    self.W2 = tf.keras.layers.Dense(units)
    self.V = tf.keras.layers.Dense(1)

  def call(self, query, values):
    # hidden shape == (batch_size, hidden size)
    # hidden_with_time_axis shape == (batch_size, 1, hidden size)
    # we are doing this to perform addition to calculate the score
    hidden_with_time_axis = tf.expand_dims(query, 1)

    # score shape == (batch_size, max_length, 1)
    # we get 1 at the last axis because we are applying score to self.V
    # the shape of the tensor before applying self.V is (batch_size, max_length, units)
    score = self.V(tf.nn.tanh(
        self.W1(values) + self.W2(hidden_with_time_axis)))

    # attention_weights shape == (batch_size, max_length, 1)
    attention_weights = tf.nn.softmax(score, axis=1)

    # context_vector shape after sum == (batch_size, hidden_size)
    context_vector = attention_weights * values
    context_vector = tf.reduce_sum(context_vector, axis=1)

    return context_vector, attention_weights

def scaled_dot_product_attention(q, k, v, mask):
  """Calculate the attention weights.
  q, k, v must have matching leading dimensions.
  k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
  The mask has different shapes depending on its type(padding or look ahead)
  but it must be broadcastable for addition.

  Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable
          to (..., seq_len_q, seq_len_k). Defaults to None.

  Returns:
    output, attention_weights
  """

  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

  # scale matmul_qk
  dk = tf.cast(tf.shape(k)[-1], tf.float32)
  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

  # add the mask to the scaled tensor.
  if mask is not None:
    scaled_attention_logits += (mask * -1e9)

    # softmax is normalized on the last axis (seq_len_k) so that the scores
  # add up to 1.
  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

  return output, attention_weights

class MultiHeadAttention(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads):
    super(MultiHeadAttention, self).__init__()
    self.num_heads = num_heads
    self.d_model = d_model

    assert d_model % self.num_heads == 0

    self.depth = d_model // self.num_heads

    self.wq = tf.keras.layers.Dense(d_model)
    self.wk = tf.keras.layers.Dense(d_model)
    self.wv = tf.keras.layers.Dense(d_model)

    self.dense = tf.keras.layers.Dense(d_model)

  def split_heads(self, x, batch_size):
    """Split the last dimension into (num_heads, depth).
    Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
    """
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
    return tf.transpose(x, perm=[0, 2, 1, 3])

  def call(self, v, k, q, mask):
    batch_size = tf.shape(q)[0]

    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)

    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
    scaled_attention, attention_weights = scaled_dot_product_attention(
      q, k, v, mask)

    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention,
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

    return output, attention_weights

class Decoder(tfk.Model):
  def __init__(self, args):
    super(Decoder, self).__init__()
    self.batch_size = args.batch_size
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prediction_horizon = args.prediction_horizon
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.output_dim = args.output_dim
    self.lambda_ = 0.01

    self.gru = tf.keras.layers.GRU(self.rnn_state_size_lstm_concat,
                                   return_sequences=True,
                                   return_state=True,
                                   recurrent_initializer='glorot_uniform')
    self.fc1 = tfkl.Dense(self.fc_hidden_unit_size, activation='relu', name='concat_fc',
                                              kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))
    self.fco = tfkl.Dense(self.output_dim*self.prediction_horizon, activation='linear', name='out_fc',
                                              kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))

    # used for attention
    self.attention = BahdanauAttention(self.rnn_state_size_lstm_concat)

  def call(self, x, hidden,ped_info):

    # x shape after concatenation == (batch_size, 1, embedding_dim + hidden_size)
    concat_x = tf.concat([ped_info, x], axis=-1)

    context_vector, attention_weights = self.attention(hidden, concat_x)

    context_vector = tf.expand_dims(context_vector, 1)

    # passing the concatenated vector to the GRU
    output, state = self.gru(concat_x,initial_state=hidden)

    # two FCL
    output = self.fc1(output)
    output = self.fco(output)

    return output, state, attention_weights

  def initialize_hidden_state(self):
    return tf.zeros((self.batch_size, self.rnn_state_size_lstm_concat))

class RNN(tf.keras.Model):

  def __init__(self,args):
    super(RNN, self).__init__()
    self.args = args

    # Network Elements

    self.agent_state_encoder = EncoderState(self.args.rnn_state_size, self.args.truncated_backprop_length, self.args.input_state_dim, self.args.batch_size)
    self.ped_info_encoder = EncoderPeds(self.args.rnn_state_size_lstm_ped, self.args.truncated_backprop_length, self.args.input_state_dim, self.args.batch_size)
    self.decoder = Decoder(self.args)
    # Initialize States
    self.state_hidden = self.agent_state_encoder.initialize_hidden_state()
    self.ped_info_hidden = self.ped_info_encoder.initialize_hidden_state()
    self.dec_hidden = self.decoder.initialize_hidden_state()

    # Optimizer and Loss
    self.optimizer = tf.keras.optimizers.Adam()
    self.loss_object = tf.keras.losses.Huber()

  def call(self, x, ped_info):
    enc_state_output, self.state_hidden = self.agent_state_encoder(x, self.state_hidden)

    enc_ped_output, self.ped_info_hidden = self.ped_info_encoder(ped_info, self.ped_info_hidden)

    return self.decoder(enc_state_output, self.dec_hidden,enc_ped_output)

  def loss_function(self,real, pred):
    loss_ = self.loss_object(real, pred)

    return tf.reduce_mean(loss_)

  @tf.function
  def train_step(self,x, ped, targ):
    loss = 0

    with tf.GradientTape() as tape:
      enc_state_output, enc_state_hidden = self.agent_state_encoder(x, self.state_hidden)

      enc_ped_output, enc_ped_info_hidden = self.ped_info_encoder(ped, self.ped_info_hidden)

      predictions, dec_hidden, attention_weights = self.decoder(enc_state_output, self.dec_hidden,
                                                           enc_ped_output)

      # Teacher forcing - feeding the target as the next input
      for t in range(1, targ.shape[1]):
        loss += self.loss_function(targ[:, t], predictions[:, t])

    batch_loss = (loss / int(targ.shape[1]))

    gradients = tape.gradient(loss, self.trainable_variables)

    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

    return batch_loss, enc_state_hidden, enc_ped_info_hidden, dec_hidden, predictions

  @tf.function
  def predict(self,batch_x, other_agents_info):
    enc_state_output, enc_state_hidden = self.agent_state_encoder(batch_x, self.state_hidden)

    enc_ped_output, enc_ped_info_hidden = self.ped_info_encoder(other_agents_info, self.ped_info_hidden)

    predictions, dec_hidden, attention_weights = self.decoder(enc_state_output, self.dec_hidden,
                                                              enc_ped_output)

    return predictions
build/lib/models/modelKerasRNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
tfk = tf.keras
tfkm = tf.keras.models
tfkl = tf.keras.layers

class StateEncoder(tfkl.Layer):
    def __init__(self, args):
        super().__init__()
        
        self.lambda_ = 0.01
        self.rnn_state_size = args.rnn_state_size
        
        self.lstm_ped = tfkl.LSTM(self.rnn_state_size, name = 'lstm_state', return_sequences = False, kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))
        
    def call(self, x):
        x = self.lstm_ped(x)
        return x

class OtherPedsEncoder(tfkl.Layer):
    def __init__(self, args):
        super().__init__()

        self.lambda_ = 0.01
        self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
        self.rnn_state_size_bilstm = args.rnn_state_size_lstm_ped # TODO: Possibly add another argument to modify this state size

        self.lstm_other_peds = tfkl.LSTM(self.rnn_state_size_lstm_ped, name = 'lstm_other_peds', return_sequences = False, kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))
        self.bilstm = tfkl.Bidirectional(tfkl.LSTM(self.rnn_state_size_bilstm,  name = 'bilstm_other_peds', return_sequences = False, return_state = False, kernel_regularizer=tfk.regularizers.l2(l=self.lambda_)), merge_mode = 'ave') # Average of forward and backward LSTMs
        
    def call(self, x):
        other_peds = []
        for x_ped in x:
            other_peds.append(self.lstm_other_peds(x_ped))
        stacked_other_peds = tf.stack(other_peds, axis = 1)
        out = self.bilstm(stacked_other_peds)
        return out

class Decoder(tfkl.Layer):
    def __init__(self, args):
        super().__init__()
        
        self.output_dim = args.output_dim
    
        self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
        self.fc_hidden_unit_size = args.fc_hidden_unit_size

        self.lstm_concat = tfkl.LSTM(self.rnn_state_size_lstm_concat, name = 'lstm_decoder', return_sequences = True, return_state = False)
        self.fc1 = tfkl.TimeDistributed(tfkl.Dense(self.fc_hidden_unit_size, name = 'fc_decoder', activation = 'relu'))
        self.fco = tfkl.TimeDistributed(tfkl.Dense(self.output_dim, name = 'fc_out'))
                
    def call(self, x):        
        x = self.lstm_concat(x)
        x = self.fc1(x)
        x = self.fco(x)
        return x

class NetworkModel(tfk.Model):
    def __init__(self, args):
        super().__init__()
        
        future_steps = args.prediction_horizon
        
        # Define optimizer and loss
        self.loss_object = tfk.losses.MeanSquaredError()
        self.optimizer = tfk.optimizers.Adam()
        self.train_loss = tfk.metrics.MeanSquaredError(name='train_loss')
        self.val_loss = tfk.metrics.MeanSquaredError(name='val_loss')
        
        # Define architecture
        self.state_encoder = StateEncoder(args)
        self.other_peds_encoder = OtherPedsEncoder(args)
        self.concat = tfkl.Concatenate()
        self.repeat = tfkl.RepeatVector(future_steps)
        self.decoder = Decoder(args)
                
    def call(self, x):
        x1 = self.state_encoder(x[0])
        
        x2 = self.other_peds_encoder(x[1:])
        concat = self.concat([x1, x2])
        repeated = self.repeat(concat)
        out = self.decoder(repeated)
        return out
    
    @tf.function
    def train_step(self, X, Y):
        with tf.GradientTape() as tape:
            predictions = self(X)
            loss = self.loss_object(Y, predictions)
        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        self.train_loss(Y, predictions)

    @tf.function
    def val_step(self, X, Y):
        predictions = self(X)
        v_loss = self.loss_object(Y, predictions)

        self.val_loss(Y, predictions)



build/lib/models/modelKerasRNN_arbitraryAgents.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
# import tensorflow_probability as tfp
import numpy as np
import os


class NetworkModel():

	def __init__(self, args, batch_size=None, is_training=True):

		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_state_dim = args.output_dim
		self.output_pred_state_dim = args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.n_mixtures = args.n_mixtures
		self.output_sequence_length = args.prediction_horizon
		self.output_placeholder_dim = self.output_sequence_length * self.output_state_dim
		self.output_dim = self.output_sequence_length * self.output_pred_state_dim
		self.pedestrian_vector_dim = args.pedestrian_vector_dim

		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size

		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.keep_prob = args.keep_prob
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim], name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    int(args.submap_width/args.submap_resolution),
		                                                    int(args.submap_height/args.submap_resolution)], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')

		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, args.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, args.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		# Set up model with variables
		with tf.variable_scope('model') as scope:

			# CNN weights for static occupancy grid processing
			train_grid_encoder_conv = False
			train_grid_encoder_fc = False
			use_summary_convnet = False

			# Convolutional layer 1
			self.conv1_kernel_size = 5
			self.conv1_number_filters = 64
			self.conv1_stride_length = 2
			self.conv1_weights = self.get_weight_variable(name="conv1_weights",
			                                              shape=[self.conv1_kernel_size, self.conv1_kernel_size, 1,
			                                                     self.conv1_number_filters],
			                                              trainable=train_grid_encoder_conv, summary=use_summary_convnet)
			self.conv1_biases = self.get_bias_variable(name="conv1_biases", shape=[self.conv1_number_filters],
			                                           trainable=train_grid_encoder_conv, summary=use_summary_convnet)

			# Convolutional layer 2
			self.conv2_kernel_size = 3
			self.conv2_number_filters = 32
			self.conv2_stride_length = 2
			self.conv2_weights = self.get_weight_variable(name="conv2_weights",
			                                              shape=[self.conv2_kernel_size, self.conv2_kernel_size,
			                                                     self.conv1_number_filters, self.conv2_number_filters],
			                                              trainable=train_grid_encoder_conv, summary=use_summary_convnet)
			self.conv2_biases = self.get_bias_variable(name="conv2_biases", shape=[self.conv2_number_filters],
			                                           trainable=train_grid_encoder_conv, summary=use_summary_convnet)

			# Convolutional layer 3
			self.conv3_kernel_size = 3
			self.conv3_number_filters = 8
			self.conv3_stride_length = 2
			self.conv3_weights = self.get_weight_variable(name="conv3_weights",
			                                              shape=[self.conv3_kernel_size, self.conv3_kernel_size,
			                                                     self.conv2_number_filters, self.conv3_number_filters],
			                                              trainable=train_grid_encoder_conv, summary=use_summary_convnet)
			self.conv3_biases = self.get_bias_variable(name="conv3_biases", shape=[self.conv3_number_filters],
			                                           trainable=train_grid_encoder_conv, summary=use_summary_convnet)

			fc_grid_hidden_dim = 64
			self.fc_grid_weights = self.get_weight_variable(shape=[512, fc_grid_hidden_dim], name="fc_grid_weights",
			                                                trainable=train_grid_encoder_fc)
			self.fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
			                                             trainable=train_grid_encoder_fc)
			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
				                                                  trainable=self.is_training)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim],
				                                                trainable=self.is_training)

				###################### Assemble Neural Net #####################
				# Process occupancy grid, apply convolutional filters
				occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
				conv_grid_feature_series = self.process_grid(occ_grid_series, trainable=self.is_training)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
				# Embedding layer of pedestrian grid
				ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					cell_grid = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					cell_ped = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(cell_ped,
					                                                                                           ped_grid_feature_series,
					                                                                                           dtype=tf.float32,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					cell = tf.keras.layers.LSTMCell(self.rnn_state_size, name='basic_lstm_cell')
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(cell,
					                                                                               inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid,
				                                               self.cell_outputs_series_lstm_ped):
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			with tf.variable_scope('decoder_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim], trainable=self.is_training)

				# Second concatenated LSTM layer
				with tf.variable_scope('lstm_concat') as scope:
					cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# FC layers
				fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
				           self.cell_outputs_series_lstm_concat]
				fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.keep_prob) for dense in
				           fc_fuse]
				# fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
				# Model prediction
				self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

			# Loss function
			loss_list = []
			lossfunc = []
			# Iterate through tbp steps
			for pred, out in zip(self.prediction,
			                     outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
				prediction_loss_list = []
				val_prediction_loss_list = []
				for prediction_step in range(self.output_sequence_length):
					# Compute Euclidean error for each prediction step within each tbp step
					idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
					idx_x = prediction_step * self.output_state_dim
					idx_y = idx_x + 1

					# Calculate loss for the current ped
					lossfunc= tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
						               tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

					prediction_loss_list.append(lossfunc)  # sum over number of batches

				loss_list.append(tf.reduce_mean(prediction_loss_list)) # mean over prediction horizon

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0) #+ l2 # mean over tbp steps

			# Parameter optimization

			# Exponential / stepwise decrease of learning rate
			global_step = tf.Variable(0, trainable=False)
			self.learning_rate = tf.train.exponential_decay(args.learning_rate_init, global_step, decay_steps=5000,
			                                                decay_rate=0.975, staircase=True)
			self.beta = (tf.tanh((tf.to_float(self.step) - 3500) / 1000) + 1) / 2
			# Optimizer specification
			#self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(args.log_dir, 'train'), graph=tf.Session().graph)

			# Saver

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(var_list={'conv1_weights': self.conv1_weights,
			                                              'conv1_biases': self.conv1_biases,
			                                              'conv2_weights': self.conv2_weights,
			                                              'conv2_biases': self.conv2_biases,
			                                              'conv3_weights:': self.conv3_weights,
			                                              'conv3_biases': self.conv3_biases,
			                                              'fc_grid_weights': self.fc_grid_weights,
			                                              'fc_grid_biases': self.fc_grid_biases,
			                                              })
			# Retrieve just the LSTM variables.
			lstm_variables = [v for v in tf.all_variables()
			                  if "encoder" in v.name]
			self.encoder_saver = tf.train.Saver(var_list=lstm_variables)

	def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
		return {self.input_state_placeholder: batch_x[:, :, 2:],
		        self.input_grid_placeholder: batch_grid,
		        self.input_ped_grid_placeholder: batch_ped_grid,
		        self.step: step,
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
		        self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
						self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat
		        }

	def feed_test_dic(self, batch_x, batch_grid, batch_ped_grid,batch_target):
		return {self.input_state_placeholder: batch_x,
		        self.input_grid_placeholder: batch_grid,
		        self.input_ped_grid_placeholder: batch_ped_grid,
		        self.output_placeholder: batch_target,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
		        self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def run(self, sess, feed_dict_train):
		_, batch_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta = sess.run([self.update,
		                                                      self.total_loss,
		                                                      self.current_state,
		                                                      self.current_state_lstm_grid,
		                                                      self.current_state_lstm_ped,
		                                                      self.current_state_lstm_concat,
		                                                      self.prediction,
		                                                      self.summary,
		                                                      self.learning_rate,
		                                                      self.beta],
		                                                     feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, _model_prediction, _summary_str, lr, beta

	def predict(self, sess, feed_dict_train,update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, loss = sess.run([self.current_state,
		                              self.current_state_lstm_grid,
		                              self.current_state_lstm_ped,
		                              self.current_state_lstm_concat,
		                              self.prediction,
		                              self.total_loss],
		                             feed_dict=feed_dict_train)
		if update_state:
			self.cell_state_current, self.hidden_state_current = _current_state
			self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, loss

	def reset_cells(self,sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def update_states(self, current_state, current_state_lstm_grid, current_state_lstm_ped, current_state_lstm_concat):
		# Store the hidden states to re-initialize the in the next training iteration if necessary
		self.cell_state_current, self.hidden_state_current = current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = current_state_lstm_concat

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not finded...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

	def process_grid(self, occ_grid_series, trainable):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)
			conv1 = self.conv_layer(input=grid_batch, weights=self.conv1_weights, biases=self.conv1_biases,
			                        conv_stride_length=self.conv1_stride_length, name="conv1_grid")
			conv2 = self.conv_layer(input=conv1, weights=self.conv2_weights, biases=self.conv2_biases,
			                        conv_stride_length=self.conv2_stride_length, name="conv2_grid")
			conv3 = self.conv_layer(input=conv2, weights=self.conv3_weights, biases=self.conv3_biases,
			                        conv_stride_length=self.conv3_stride_length, name="conv3_grid")
			conv_grid_size = 8
			conv5_flat = tf.reshape(conv3, [-1, conv_grid_size * conv_grid_size * self.conv3_number_filters])
			fc_final = self.fc_layer(input=conv5_flat, weights=self.fc_grid_weights, biases=self.fc_grid_biases,
			                         use_activation=True, name="fc_grid")
			if idx == 0:
				tf.summary.histogram("fc_activations", fc_final)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, use_activation=False,
				              name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def fc_layer(self, input, weights, biases, use_activation=False, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if use_activation:
			activations = tf.nn.sigmoid(affine_result)
		else:
			activations = affine_result
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations




build/lib/models/modelRNN_v2.py
eliatrevisan/IA-MPPI-LBM
#!/usr/bin/env python
import time
import rospy
import rospkg
import cv2
from geometry_msgs.msg import Pose, PoseStamped
from pedsim_msgs.msg import TrackedPersons, TrackedPerson
from nav_msgs.msg import Path, OccupancyGrid
from visualization_msgs.msg import MarkerArray, Marker
import tf as transferfunction
import sys; print('Python %s on %s' % (sys.version, sys.platform))
sys.path.extend(['/media/bdebrito/7697ec91-468a-4763-b1c3-135caa7f5aed/home/code/I-LSTM/src/data_utils'])
sys.path.extend(['/media/bdebrito/7697ec91-468a-4763-b1c3-135caa7f5aed/home/code/I-LSTM/src/models'])
import os
import math
import json
if sys.version_info[0] < 3:
	import Support as sup
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils import Support as sup

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as pl
import threading as th
from lmpcc_msgs.msg import lmpcc_obstacle, lmpcc_obstacle_array, lmpcc_predict_point, lmpcc_predict_path
import pickle as pkl


model_name = "Plan_VGDNN"

# Import model
from RNN_plannet import NetworkModel

class PlanNet:
	def __init__(self):

		self.model_name = rospy.get_param('~model_name', 'RNN_plannet')
		self.model_id = rospy.get_param('~id', '8')
		self.robot_state_topic = rospy.get_param('~robot_state_topic', '/robot_state')
		self.goal_topic = rospy.get_param('~goal_topic', '/move_base_simple/goal')
		self.other_agents_topic = rospy.get_param('~other_agents_topic',"/pedsim_visualizer/tracked_persons")
		self.grid_topic = rospy.get_param('~grid_topic',"/move_base/local_costmap/costmap")

		# Load Model Parameters
		self.load_args()

		# State variables
		self.current_position_ = np.zeros([1,1,(self.model_args.prev_horizon+1)*2])
		self.current_velocity_ = np.zeros([1,1,(self.model_args.prev_horizon+1)*2])
		self.goal_ = np.zeros([1,1,2])
		self.rel_goal = np.zeros([1,1,self.model_args.goal_size])
		self.n_other_agents = 8
		self.other_agents = np.ones([1,1,self.model_args.pedestrian_vector_dim])
		self.other_agents_info = np.zeros([1, 1, self.model_args.pedestrian_vector_dim])

		self.n_peds = 8

		self.other_pedestrians = []
		for i in range(self.n_peds):
			self.other_pedestrians.append(TrackedPerson())

		self.load_model()

		self.width = int(self.model_args.submap_width / self.model_args.submap_resolution )
		self.height = int(self.model_args.submap_height / self.model_args.submap_resolution)
		self.grid = np.zeros([1, 1, self.width, self.height])
		self.fig_animate = pl.figure('Animation')
		self.fig_width = 12  # width in inches
		self.fig_height = 25  # height in inches
		self.fig_size = [self.fig_width, self.fig_height]
		self.fontsize = 9
		self.colors = ['r', 'g', 'y']
		self.params = {'backend': 'ps',
		          'axes.labelsize': self.fontsize,
		          'font.size': self.fontsize,
		          'xtick.labelsize': self.fontsize,
		          'ytick.labelsize': self.fontsize,
		          'figure.figsize': self.fig_size}
		pl.rcParams.update(self.params)
		#self.ax_pos = pl.subplot()
		pl.show(block=False)

		# ROS Subscribers
		rospy.Subscriber(self.robot_state_topic, PoseStamped, self.robot_state_CB, queue_size=1)
		rospy.Subscriber(self.goal_topic,PoseStamped,self.goal_CB, queue_size=1)
		rospy.Subscriber(self.other_agents_topic, TrackedPersons, self.other_agents_CB, queue_size=1)
		rospy.Subscriber(self.grid_topic, OccupancyGrid, self.grid_CB, queue_size=1)

		# ROS Publishers
		self.tf_listener = transferfunction.TransformListener()
		self.tf_transformer = transferfunction.TransformerROS()

		self.pub_viz = rospy.Publisher('plannet_trajectory', Path, queue_size=10)
		self.pub_static_collision_traj = rospy.Publisher('nn_predicted_trajectory', lmpcc_predict_path, queue_size=10)
		self.pub_goal_marker = rospy.Publisher('goal_marker', Marker, queue_size=1)

		# THread control
		self.lock = th.Lock()

	def load_args(self):
		cwd = os.getcwd()

		model_path = os.path.normpath(cwd + '/../') + '/trained_models/' + self.model_name + "/" + str(self.model_id)

		print("Loading data from: '{}'".format(model_path))
		file = open(model_path + '/model_parameters.pkl', 'rb')
		if sys.version_info[0] < 3:
			model_parameters = pkl.load(file)  # ,encoding='latin1')
		else:
			model_parameters = pkl.load(file, encoding='latin1')
		file.close()
		self.model_args = model_parameters["args"]

		if not(os.path.isfile(model_path + '/model_parameters.json')):
			json.dump(vars(self.model_args),open(model_path + '/model_parameters.json','w'))

		# change some args because we are doing inference
		self.model_args.truncated_backprop_length = 1
		self.model_args.batch_size = 1

	def load_model(self):

		self.model = NetworkModel(self.model_args)

		self.sess = tf.Session()

		self.model.warmstart_model(self.model_args, self.sess)
		try:
			self.model.warmstart_convnet(self.model_args, self.sess)
		except:
			print("No convnet")

		print("Model Initialized")

	def robot_state_CB(self, data):
		# Robot state callback
		# shift
		self.current_velocity_ = np.roll(self.current_velocity_, 2, axis=2)
		self.current_position_ = np.roll(self.current_position_, 2, axis=2)

		self.current_position_[0,0,0] = data.pose.position.x
		self.current_position_[0,0,1] = data.pose.position.y
		orientation = data.pose.orientation.z
		linear_velocity = data.pose.position.z

		self.current_velocity_[0,0,0] = linear_velocity*np.cos(orientation)
		self.current_velocity_[0,0,1] = linear_velocity*np.sin(orientation)

	def goal_CB(self, data):
		# Goal Callback
		self.goal_[0,0,0] = data.pose.position.x
		self.goal_[0,0,1] = data.pose.position.y

		rel_pos = self.goal_ - self.current_position_[0,0,:2]

		for t in range(self.model_args.prev_horizon+1):
			self.current_velocity_[0, 0, 2*t] = np.tanh(rel_pos[0,0,0]/self.model_args.dt)
			self.current_velocity_[0, 0, 2*t+1] = np.tanh(rel_pos[0,0,1]/self.model_args.dt)

		self.publish_goal_marker()

	def grid_CB(self, data):
		# scale the grid from 0-100 to 0-1 and invert
		print("Grid data size: " + str(len(data.data)))
		self.grid[0,0,:,:] = (np.asarray(data.data).reshape((self.width, self.height)).astype(float) / 100.0)
		self.grid[0,0,:,:] = np.flip(self.grid[0,0,:,:], 1)
		self.grid[0,0,:,:] = sup.rotate_grid_around_center(self.grid[0,0,:,:], 90)

		if False:
			self.ax_pos.clear()
			sup.plot_grid(self.ax_pos, np.array([0.0, 0.0]), self.grid, self.model_args.submap_resolution,
			              np.array([self.model_args.submap_width, self.model_args.submap_height]))
			self.ax_pos.set_xlim([-self.model_args.submap_width/2,self.model_args.submap_width/2])
			self.ax_pos.set_ylim([-self.model_args.submap_height/2,self.model_args.submap_height/2])
			self.ax_pos.set_aspect('equal')

	def other_agents_CB(self, data):
		if len(data.tracks) != self.n_peds:
			print("NUmber of peds do not match callback info")

		for person_it in range(self.n_peds):
			ped = TrackedPerson()
			ped.pose=data.tracks[person_it].pose
			ped.track_id = data.tracks[person_it].track_id
			ped.twist = data.tracks[person_it].twist
			self.other_pedestrians[person_it] = ped

	def fillBatchOtherAgents(self):

		other_poses_ordered = np.zeros((len(self.other_pedestrians), 6))

		for ag_id in range(len(self.other_pedestrians)):
			other_poses_ordered[ag_id, 0] = self.other_pedestrians[ag_id].pose.pose.position.x
			other_poses_ordered[ag_id, 1] = self.other_pedestrians[ag_id].pose.pose.position.y
			other_poses_ordered[ag_id, 2] = self.other_pedestrians[ag_id].twist.twist.linear.x
			other_poses_ordered[ag_id, 3] = self.other_pedestrians[ag_id].twist.twist.linear.y
			other_poses_ordered[ag_id, 4] = np.linalg.norm(other_poses_ordered[ag_id, :2] - self.current_position_[0,0,:2])
			other_poses_ordered[ag_id, 5] = self.other_pedestrians[ag_id].track_id

		#other_poses_ordered = other_poses_ordered[other_poses_ordered[:, 4].argsort()]
		heading = math.atan2(self.current_velocity_[0,0,1], self.current_velocity_[0,0,0])
		other_pos_local_frame = sup.positions_in_local_frame(self.current_position_[0,0,:2], heading, other_poses_ordered[:,:2])
		self.other_agents_info[0,0] = sup.compute_radial_distance_vector(self.model_args.pedestrian_vector_dim, other_pos_local_frame,
		                                                            max_range=self.model_args.max_range_ped_grid, min_angle=0,
		                                                            max_angle=2 * np.pi,
		                                                            normalize=True)

	# query feed the data into the net and calculates the trajectory
	def query(self):
		# save current position and orientation to calculate the prediction from

		self.fillBatchOtherAgents()

		rel_goal = self.goal_[0,0,:2] -self.current_position_[0,0,:2]

		#self.rel_goal[0,0, 0] = np.linalg.norm(rel_goal)
		#self.rel_goal[0,0, 1] = np.arctan2(rel_goal[1], rel_goal[0])
		self.rel_goal[0,0, :] = rel_goal/np.linalg.norm(rel_goal)

		dict = {"batch_vel": self.current_velocity_,
		        "batch_ped_grid": self.other_agents_info,
		        "batch_goal": self.rel_goal,
		        "step": 0
		        }
		feed_dict_ = self.model.feed_test_dic(**dict)

		outputs = self.model.predict(self.sess, feed_dict_, True)

		# publish the predicted trajectories
		local_trajectory, global_trajectory = self.calculate_trajectory(outputs[0])
		self.local_trajectory = local_trajectory
		self.global_trajectory = global_trajectory

	# add the velocity predictions together to form points and convert them to global coordinate frame
	def calculate_trajectory(self, y_model_pred):
		local_trajectory = np.empty([self.model_args.prediction_horizon, 2])
		global_trajectory = Path()
		global_trajectory.header.frame_id = "odom"
		global_trajectory.header.stamp = rospy.Time.now()
		x = 0.0
		y = 0.0
		x_global = self.current_position_[0,0,0]
		y_global = self.current_position_[0,0,1]
		mix_idx = 0
		for pred_step in range(self.model_args.prediction_horizon):
			pose = PoseStamped()
			idx = pred_step * self.model_args.output_pred_state_dim * self.model_args.n_mixtures + mix_idx
			x = x + y_model_pred[0][0][idx] * self.model_args.dt
			y = y + y_model_pred[0][0][idx + self.model_args.n_mixtures] * self.model_args.dt
			local_trajectory[pred_step, :] = np.array([x, y])
			x_global += y_model_pred[0][0][idx] * self.model_args.dt
			y_global += y_model_pred[0][0][idx + self.model_args.n_mixtures] * self.model_args.dt

			pose.pose.position.x = x_global
			pose.pose.position.y = y_global
			pose.pose.orientation.w = 1.0
			global_trajectory.poses.append(pose)

		return local_trajectory, global_trajectory

	# gather the trajectories from the different pedestrians and publish their visualisation
	def publish_trajectory_visualisation(self):

		self.pub_viz.publish(self.global_trajectory)


	def publish_goal_marker(self):
		marker = Marker()
		marker.header.frame_id = "odom"
		marker.header.stamp = rospy.Time.now()
		marker.ns = "goal_marker"
		marker.id = 0
		marker.type = 3
		marker.pose.position.z = 0.1
		marker.pose.orientation.x = 0.0
		marker.pose.orientation.y = 0.0
		marker.pose.orientation.z = 0.0
		marker.pose.orientation.w = 1.0
		marker.scale.x = 1.2
		marker.scale.y = 1.2
		marker.scale.z = 1.1
		marker.pose.position.x = self.goal_[0,0,0]
		marker.pose.position.y = self.goal_[0,0,1]
		marker.color.a = 1.0
		marker.color.r = 1.0
		marker.color.g = 0.0
		marker.color.b = 0.0
		marker.lifetime = rospy.rostime.Duration()
		self.pub_goal_marker.publish(marker)


def visualize_ped_grid(self, ped_grid):
	pedestrian_vector_dim = 72
	max_range_ped_grid = 1.0
	fig1 = plt.figure("Angular grid")
	ax_ped_grid = plt.subplot()
	ax_ped_grid.clear()
	grid = ped_grid[0]
	grid = np.expand_dims(grid, axis=1)
	grid_flipped = np.zeros_like(grid)
	grid_flipped[0:int(pedestrian_vector_dim / 2)] = grid[-int(pedestrian_vector_dim / 2):]
	grid_flipped[int(pedestrian_vector_dim / 2):] = grid[0:int(pedestrian_vector_dim / 2)]
	sup.plot_radial_distance_vector(ax_ped_grid, grid_flipped, max_range=1.0, min_angle=0.0,
	                                max_angle=2 * np.pi)
	ax_ped_grid.plot(30, 30, color='r', marker='o', markersize=4)
	ax_ped_grid.arrow(0, 0, 1, 0, head_width=0.1, head_length=max_range_ped_grid)  # agent poiting direction
	# x- and y-range only need to be [-1, 1] since the pedestrian grid is normalized
	ax_ped_grid.set_xlim([-max_range_ped_grid - 1, max_range_ped_grid + 1])
	ax_ped_grid.set_ylim([-max_range_ped_grid - 1, max_range_ped_grid + 1])

	fig1.canvas.draw()
	plt.show(block=False)


if __name__ == '__main__':
	rospy.init_node('PlanNet_node')
	planning_network = PlanNet()
	rospy.sleep(5.0)
	while not rospy.is_shutdown():
		start_time = time.time()

		planning_network.lock.acquire()
		planning_network.query()
		planning_network.publish_trajectory_visualisation()
		planning_network.lock.release()
		#planning_network.fig_animate.canvas.draw()
		#cv2.imshow("image", planning_network.grid)
		#cv2.waitKey(100)
		# wait around a bit if neccesairy
		now = time.time()
		if now - start_time < 0.05:  # args.dt:
			rospy.sleep(0.05 - (now - start_time))
		else:
			print("not keeping up to rate")
	planning_network.sess.close()

build/lib/plannet.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
import json
import importlib
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *

import pickle as pkl
import time
from copy import deepcopy
from multiprocessing.pool import ThreadPool
from progressbar import *

home = os.path.expanduser('~')

class Map:
    size = [1715, 881]
    data = pl.imread(home + '/I-LSTM/data/simulation/roboat/map.png')
    origin = [-78, -40, 0.0]
    resolution = 0.081

def plot_map(map, ax):
    ax.imshow(map.data,
               extent = (map.origin[0], map.origin[0] + map.size[0] * map.resolution,
                         map.origin[1], map.origin[1] + map.size[1] * map.resolution),
               cmap='gray_r')
    ax.set_facecolor('#262626')
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    ax.yaxis.set_label_coords(-0.08, .5)
    ax.xaxis.set_label_coords(0.5, -0.09)

#os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

pretrained_convnet_path = "../trained_models/autoencoder_with_ped"

#data_path = '../data/2_agents_swap/trajs/'
#scenario = 'GA3C-CADRL-10-py27'
#data_path = '../data/cyberzoo_experiments/'
#scenario = 'all_trajectories'
exp_num = 2
data_path = '../data/'
scenario = 'simulation/roboat'

# Hyperparameters
n_epochs = 2

batch_size = 1
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 10
prediction_horizon = 10
prev_horizon = 0

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 32
rnn_state_ped_size = 16
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 256
x_dim = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 1
tensorboard_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 3

print_freq = 2000
save_freq = 2000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = True
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
submap_resolution = 0.081
submap_width = 4.86
submap_height = 4.86
diversity_update = False
predict_positions = False
warm_start_convnet = True

train_set = 1.0

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default= "RNN")
	parser.add_argument('--model_path',
	                    help='Path to directory to save the model (default=""../trained_models/"+model_name").',
	                    type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
	                    help='Path to directory that comprises the pre-trained convnet model (default=" ").',
	                    type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
	                    help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
	                    type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
	                    type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
	                    default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
	                    default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
	                    default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
	                    default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
	                    default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
	                    default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
	                    default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
	                    default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
	                    default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
	                    default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
	                    default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
	                    default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
	                    help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
	                    type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
	                    type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
	                    type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
	                    help='Number of fully connected layer units after LSTM layer (default=64).',
	                    type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
	                    default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
	                    default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
	                    default=output_dim)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
	                    default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
	                    default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
	                    default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
	                    default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
	                    default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
	                    default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
	                    default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the netwprk.', type=int,
	                    default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=int, default=0)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
	                    default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
	                    type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
	                    default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
	                    default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
	                    default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
	                    default=0)
	parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=submap_width)
	parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=submap_height)
	parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
	                    default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--relative_info', help='Use relative info for other agents.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
	                    default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
	                    default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
	                    default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
	                    default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	args = parser.parse_args()

	return args


args = parse_args()

if args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf
tfk = tf.keras

args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)
args.log_dir = args.model_path + '/log'
if not os.path.exists(args.log_dir):
	os.makedirs(args.log_dir)
args.dataset = '/' + args.scenario + '.pkl'
model_parameters = {"args": args}
print(args)
# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()
with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

data_prep = dhlstm.DataHandlerLSTM(args)
# Only used to create a map from png
map_args = {"file_name": 'map.png',
	            "resolution": 0.1,
	            "map_size": np.array([30., 30.]), }

data_prep.processData(**map_args)

fig_animate = pl.figure('Dataset Trajectories')
fig_animate.tight_layout()
fig_width = 12  # width in inches
fig_height = 10  # height in inches
fig_size = [fig_width, fig_height]
fontsize = 26

params = {'backend': 'ps',
          'axes.labelsize': fontsize,
          'font.size': fontsize,
          'xtick.labelsize': fontsize,
          'ytick.labelsize': fontsize,
          'figure.figsize': fig_size}

plt_colors = [0.8500, 0.3250, 0.0980]

colors = ['r', 'g', 'b','y','m','c','k']
ax_pos = pl.subplot()

ax_pos.set_xlabel('x [m]')
ax_pos.set_ylabel('y [m]')
ax_pos.axis('on')
ax_pos.set_aspect('equal')

new_epoch = False

while not new_epoch:
	plt_colors = np.random.uniform(0,1,[3])

	batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch = data_prep.getBatch()
	mij = np.min(batch_ped_grid,axis=2)
	#print(np.min(mij,axis=1))
	if np.min(mij, axis=1) < 1.0:
		#ax_pos.plot(batch_pos_target[0,0,0::2], batch_pos_target[0,0,1::2], color=plt_colors)
		ax_pos.plot(batch_pos[0,:,0], batch_pos[0,:,1], color=plt_colors)
map = Map()
plot_map(map, ax_pos)
pl.show(block=True)

print(batch_pos[0,:,0])
print(batch_pos[0,:,1])


fig_animate.savefig(args.data_path + args.scenario + "/trajectories.jpg")

build/lib/plot_dataset.py
eliatrevisan/IA-MPPI-LBM
from cProfile import label
import sys
import os
from turtle import home

sys.path.append('../')

# With debugger:
#os.chdir(os.getcwd() + "/src/")

import numpy as np
import argparse
import pylab as pl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pickle as pkl
import importlib
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import random
import progressbar
import scipy.io as sio
import json
from copy import deepcopy
import colorama
from colorama import Fore, Style

if sys.version_info[0] < 3:
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from Performance import *
	from utils import *
	import Recorder as rec
else:
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.Performance import *
	from src.data_utils.utils import *
	from src.data_utils.Recorder import Recorder as rec
	from src.models.CV_model import ConstantVelocity

"""
exp_num = 33
model_name = "VGDNN" 
num_test_sequences = 100 
scenario = "simulation/roboat" 
record = False 
n_samples = 1 
unit_testing = False 
freeze_other_agents = False
"""

from random import seed
seed(5)

# Model directories
def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
											help='Path to directory that comprises the model (default="model_name").',
											type=str, default="VGDNN")
	parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=33)
	parser.add_argument('--n_samples', help='Number of samples', type=int, default=10)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
											type=str, default="datasets/ewap_dataset/seq_eth")
	parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool,
											default=True)
	parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool,
											default=True)
	parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float,
											default=0)
	parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float,
											default=5)
	parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool,
											default=False)
	parser.add_argument('--update_state', help='update_state', type=sup.str2bool,
											default=False)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
											default=False)
	parser.add_argument('--freeze_other_agents', help='FReeze other agents.', type=sup.str2bool,
											default=False)
	parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,
											default=False)
	parser.add_argument('--constant_velocity', help='Run CV comparison', type=sup.str2bool,
											default=False)
	args = parser.parse_args()

	return args


test_args = parse_args()

if test_args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

cwd = os.getcwd()

model_path = os.path.normpath(cwd + '/../') + '/trained_models/' + test_args.model_name + "/" + str(test_args.exp_num)

print("Loading data from: '{}'".format(model_path))
file = open(model_path + '/model_parameters.pkl', 'rb')
if sys.version_info[0] < 3:
	model_parameters = pkl.load(file)  # ,encoding='latin1')
else:
	model_parameters = pkl.load(file, encoding='latin1')
file.close()
args = model_parameters["args"]

with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

print("Training steps: ", args.total_training_steps)

# change some args because we are doing inference
truncated_backprop_length = args.truncated_backprop_length
args.truncated_backprop_length = 1
args.batch_size = 1
args.keep_prob = 1.0


training_scenario = args.scenario
args.scenario = test_args.scenario
args.real_world_data = test_args.real_world_data
args.dataset = '/' + args.scenario + '.pkl'
data_prep = dhlstm.DataHandlerLSTM(args)

# Only used to create a map from png
# Load Map Parameters
map_params = os.path.join(args.data_path+args.scenario, 'map.json')
#home = os.path.expanduser('~')
#map_params = home + "/I-LSTM/data/roboat/bloemgracht/" + 'map.json'
with open(map_params) as json_file:
	data = json.load(json_file)
map_args = {"file_name": data["file_name"],
						"resolution": data["resolution"],
						"map_size": np.array(data["map_size"]),
						"map_center": np.array(data["map_center"])}

data_prep.processData(**map_args)
if args.normalize_data:
	data_prep.compute_min_max_values()

# Import model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

model = NetworkModel(args)

# Lists for logging of the input / output data of the model
input_list = []
grid_list = []
goal_list = []
ped_grid_list = []
y_ground_truth_list = []
y_pred_list = []  # uses ground truth as input at every step
other_agents_list = []
all_predictions = []
all_traj_likelihood = []
trajectories = []
batch_y = []
batch_loss = []

config = tf.ConfigProto(
	device_count={'GPU': 0}
)


def clear_plot():
	axarr[0].clear()
	axarr[0].set_title('Training Sample')

if test_args.unit_testing:
	data_handler = dhlstm.DataHandlerLSTM(args)
	data_handler.unit_test_data_(map_args)

with tf.Session(config=config) as sess:
	model.warmstart_model(args, sess)
	try:
			model.warmstart_convnet(args, sess)
	except:
		print("")

	for exp_id in range(np.minimum(test_args.num_test_sequences,len(data_prep.trajectory_set)-1)):
		predictions = []
		traj_likelihood = []
		# sample a trajectory id for testing
		traj_id = 2670 #random.randint(0, len(data_prep.trajectory_set) - 1)

		batch_x, batch_vel, batch_pos,batch_goal, batch_grid, other_agents_info, batch_target,batch_end_pos, other_agents_pos, traj = data_prep.getTrajectoryAsBatch(
			traj_id,freeze = test_args.freeze_other_agents)  # trajectory_set random.randint(0, len(data_prep.dataset) - 1)
		
		#batch_x, batch_vel, batch_pos,batch_goal, batch_grid, other_agents_info, batch_target,batch_end_pos, other_agents_pos, traj = data_prep.getTestTrajectoryAsBatch(
		#	traj_id,freeze = test_args.freeze_other_agents)
		
		trajectories.append(traj)
		x_input_series = np.zeros([0, (args.prev_horizon + 1) * args.input_dim])
		goal_input_series = np.zeros([0, 2])
		grid_input_series = np.zeros(
			[0, int(args.submap_size), int(args.submap_size)])
		if args.others_info == "relative":
			ped_grid_series = np.zeros([0, args.n_other_agents,args.pedestrian_vector_dim])
		elif "sequence" in args.others_info:
			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim*args.prediction_horizon])
		elif args.others_info == "prev_sequence":
			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim *(args.prev_horizon+1)])
		elif args.others_info == "sequence2":
			ped_grid_series = np.zeros([0, args.n_other_agents, args.prediction_horizon,args.pedestrian_vector_dim])
		else:
			ped_grid_series = np.zeros([0, args.pedestrian_vector_dim])
		y_ground_truth_series = np.zeros([0, args.prediction_horizon * 2])
		y_pred_series = np.zeros([0, args.n_mixtures * args.prediction_horizon * args.output_pred_state_dim])

		batch_y.append(batch_target)
		model.reset_test_cells(np.ones((args.batch_size)))
		cell_state_list= []
		cell_ped_list = []
		cell_concat_list = []
		if "grid" in args.model_name:
			batch_ped_grid_backup = np.zeros_like(batch_grid)
			data_prep.add_other_agents_to_grid(batch_ped_grid_backup, batch_x, [other_agents_pos])

		for step in range(batch_x.shape[1]):
			samples = []
			# Assemble feed dict for training
			if "future" in args.others_info:
				if step == 0:
					batch_y_pred = deepcopy(batch_vel)
					for pred_step in range(1, args.prediction_horizon):
						batch_y_pred[:, :, pred_step * 2:pred_step * 2 + 2] = batch_vel[:, :, :2]
				else:
					batch_y_pred = deepcopy(batch_vel)
					batch_y_pred[:, :, 2:] = y_model_pred[:, :, 2:]
			"""
			dict = {"batch_x": batch_x,
							"batch_vel": batch_vel,
							"batch_pos": batch_pos,
							"batch_grid": batch_grid,
							"batch_ped_grid": other_agents_info,
							"step": step,
							"batch_goal": batch_goal,
							"state_noise": 0.0,
							"grid_noise": 0.0,
							"other_agents_pos": [other_agents_pos]
			}
		
			feed_dict_ = model.feed_test_dic(**dict)
			"""

			dict = {
							"batch_vel": batch_vel,
							"batch_grid": batch_grid,
							"batch_ped_grid": other_agents_info,
							"step": step,
							"other_agents_pos": [other_agents_pos]
			}

			feed_dict_ = model.feed_test_dic(
                **dict
            )
			ag1 = []
			ag2 = []
			ag3 = []
			for i in range(len(other_agents_pos)):
				ag1.append(other_agents_pos[i][0][:2])
				ag2.append(other_agents_pos[i][1][:2])
				ag3.append(other_agents_pos[i][2][:2])
			ag1 = np.array(ag1)
			ag2 = np.array(ag2)
			ag3 = np.array(ag3)
			
			fig, ax = plt.subplots()
			sup.plot_grid_roboat(ax, data_prep.agent_container.occupancy_grid.center ,data_prep.agent_container.occupancy_grid.gridmap, data_prep.agent_container.occupancy_grid.resolution, data_prep.agent_container.occupancy_grid.map_size)
			ax.plot(ag1[:,0],ag1[:,1], label='1')
			ax.plot(ag2[:,0], ag2[:,1], label='2')
			ax.plot(ag3[:,0], ag3[:,1], label='3')
			plt.legend()
			plt.show()

			# Append to logging series
			x_input_series = np.append(x_input_series, batch_x[:, step, :], axis=0)
			grid_input_series = np.append(grid_input_series, batch_grid[:, step, :, :], axis=0)
			goal_input_series = np.append(goal_input_series, batch_goal[:, step, :], axis=0)

			y_ground_truth_series = np.append(y_ground_truth_series, batch_target[:, step, :], axis=0)

			y_model_pred, likelihood = model.predict(sess, feed_dict_, True)

			# Backup cell states for later analysis
			cell_state_list.append(model.test_cell_state_current[0,:])
			#cell_ped_list.append(model.test_cell_state_current_lstm_ped[0, :])
			cell_concat_list.append(model.test_cell_state_current_lstm_concat[0, :])

			# Rotate predictions to global frame
			if args.rotated_grid:
				heading = math.atan2(batch_vel[0, step, 1], batch_vel[0, step, 0])
				rot_mat = np.array([[np.cos(heading), -np.sin(heading)], [np.sin(heading), np.cos(heading)]])
				for pred_step in range(args.prediction_horizon):
					y_model_pred[0][0, 2 * pred_step:2 * pred_step + 2] = np.dot(rot_mat, y_model_pred[0][0,2 * pred_step:2 * pred_step + 2])
				samples.append(y_model_pred[0])
			else:
				samples.append(y_model_pred[:,0,:])

			test = False
			# If sample more than one trajectory from the model
			for sample_id in range(test_args.n_samples - 1):
				test = True
				dict = {"batch_x": batch_x,
								"batch_vel": batch_vel,
								"batch_pos": batch_pos,
								"batch_grid": batch_grid,
								"batch_ped_grid": other_agents_info,
								"step": step,
								"batch_goal": batch_goal,
								"state_noise": 0.0,
								"grid_noise": 0.0
								}
				feed_dict_ = model.feed_test_dic(**dict)
				y_model_pred = model.predict(sess, feed_dict_, test_args.update_state)
				samples.append(y_model_pred[:,0,:])

			traj_likelihood.append(likelihood)
			predictions.append(samples)

		latent_analysis = False

		# Plot hidden states evolution over a trajectory
		if latent_analysis:

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			pl.show(block=False)
			x_state = np.arange(0, args.rnn_state_size, 1)
			x_ped = np.arange(0, args.rnn_state_size_lstm_ped * 2, 1)
			x_concat = np.arange(0, args.rnn_state_size_lstm_concat, 1)

			time = np.arange(0, len(cell_state_list), 1)
			R= np.array(cell_state_list)
			ax.clear()
			X, Y = np.meshgrid(x_state,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			save_img_to_file = model_path + "/figs"
			if not os.path.exists(save_img_to_file):
				os.makedirs(save_img_to_file)
			pl.savefig(save_img_to_file+"/cell_state" + str(exp_id) + '.png')
			pl.close()

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			time = np.arange(0, len(cell_ped_list), 1)
			R= np.array(cell_ped_list)
			#print(R)
			ax.clear()
			X, Y = np.meshgrid(x_ped,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			pl.savefig(save_img_to_file+"/cell_ped" + str(exp_id) + '.png')
			pl.close()

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			time = np.arange(0, len(cell_concat_list), 1)
			R= np.array(cell_concat_list)
			ax.clear()
			X, Y = np.meshgrid(x_concat,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			pl.savefig(save_img_to_file+"/cell_concat" + str(exp_id) + '.png')
			pl.close()

		all_predictions.append(predictions)
		all_traj_likelihood.append(traj_likelihood)
		input_list.append(x_input_series)
		goal_list.append(goal_input_series)
		grid_list.append(grid_input_series)
		y_ground_truth_list.append(y_ground_truth_series)
		other_agents_list.append(other_agents_pos)
		# update progress bar

sess.close()

# Save data in Matlab format for post-processing
if not os.path.exists(args.model_path + '/../results/'):
	os.makedirs(args.model_path + '/../results/')
results= {
	"trajectories": trajectories,
	"predictions": all_predictions
}
scenario = args.scenario.split('/')[-1]
results_file = args.model_path + '/'+ scenario + "_results.mat"
sio.savemat(results_file,results)

if test_args.constant_velocity:
	# Make constant velocity predictions
	CVModel = ConstantVelocity(args)
	cv_predictions = CVModel.predict(input_list)
else:
	cv_predictions = None


if test_args.record:
		recorder = rec(args, data_prep.agent_container.occupancy_grid)
		if ( "real_world" in test_args.scenario) and not test_args.unit_testing:
			print("Real data!!")
			recorder.plot_on_video(input_list, grid_list, all_predictions, y_ground_truth_list,
														 other_agents_list,
														 trajectories,all_traj_likelihood, test_args)
		else:
			#recorder.plot_on_image(input_list, grid_list, all_predictions, y_ground_truth_list, other_agents_list,
			#	                       trajectories,test_args)
			recorder.animate_local(input_list, grid_list, ped_grid_list, all_predictions, y_ground_truth_list, other_agents_list,
										 trajectories,test_args)

			recorder.animate_global(input_list, grid_list, all_predictions, y_ground_truth_list,
														 other_agents_list,
														 trajectories, all_traj_likelihood, cv_predictions, test_args)

			print("Recorder is done!")
else:
	print("Performance tests")
	pred_error, pred_error_summary_lstm = compute_trajectory_prediction_mse(args, trajectories, all_predictions)
	pred_fde, pred_error_summary_lstm_fde = compute_trajectory_fde(args, trajectories, all_predictions)
	if test_args.constant_velocity: 
		cv_pred_error = compute_ade_cv(args, trajectories, cv_predictions)
		cv_fde_error = compute_fde_cv(args, trajectories, cv_predictions)
		print(
		Fore.LIGHTBLUE_EX + "\nConstant Velocity: MSE: {:01.2f}, FDE: {:01.2f}".format(np.mean(cv_pred_error), np.mean(cv_fde_error))+Style.RESET_ALL)
	diversity, diversity_summary = compute_2_wasserstein(args, all_predictions)
	args.scenario = training_scenario
	args.truncated_backprop_length = truncated_backprop_length
	write_results_summary(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde), np.mean(diversity_summary), args, test_args, cv_pred_error, cv_fde_error)
	print(
		Fore.LIGHTBLUE_EX + "\nMSE: {:01.2f}, FDE: {:01.2f}, DIVERSITY: {:01.2f}".format(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde),np.mean(diversity_summary))+Style.RESET_ALL)


build/lib/scenario_check.py
eliatrevisan/IA-MPPI-LBM
# -*- coding: utf-8 -*-
"""seq2seq_ped_predictions.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1u06XueavEOkuZubGh_RspAfQi7O2vApB
"""

from __future__ import print_function

from keras.models import Model
from keras.layers import Input, LSTM, Dense
import numpy as np
import pickle
import matplotlib.pyplot as plt

# Path to the traj data pickle file on disk.
#trajs_filename = 'trajs.p'
#with open(trajs_filename, "rb") as f:
#	dataset = pickle.load(f)

batch_size = 64  # Batch size for training.
latent_dim = 256  # Latent dimensionality of the encoding space.

max_encoder_seq_length = 10  # how many timesteps of history provided to seed the encoder
max_decoder_seq_length = 5  # how many timesteps decoder has to predict

num_encoder_tokens = 8  # how many elements in the input space (xy pos, xy vel per agent)
num_decoder_tokens = 2  # how many elements in the predicted space (xy pos)


# Split the array of trajectories into 3 arrays of traj snippets
def multivariate_data_3d(dataset, first_traj_index, last_traj_index, history_size,
                         target_size, step):
	encoder_input = []
	decoder_input = []
	decoder_target = []

	if last_traj_index is None:
		last_traj_index = len(dataset) + 1

	for traj in dataset[first_traj_index:last_traj_index]:
		target = traj[:, 0:2]
		start_index = 0 + history_size
		end_index = len(traj) - target_size
		for i in range(start_index, end_index):
			indices = range(i - history_size, i, step)
			encoder_input.append(traj[indices])
			decoder_input.append(target[i - 1:i + target_size - 1])
			decoder_target.append(target[i:i + target_size])

	return np.array(encoder_input), np.array(decoder_input), np.array(decoder_target)


#encoder_input_data, decoder_input_data, decoder_target_data = multivariate_data_3d(dataset, 0, None,
#                                                                                   max_encoder_seq_length,
#                                                                                   max_decoder_seq_length, 1)

#print(encoder_input_data.shape)
#print(decoder_input_data.shape)
#print(decoder_target_data.shape)

# print(encoder_input_data[0])
# print(decoder_input_data[0])
# print(decoder_target_data[0])

# Define an input sequence and process it.
encoder_inputs = Input(shape=(None, num_encoder_tokens))
encoder = LSTM(latent_dim, return_state=True)
encoder_outputs, state_h, state_c = encoder(encoder_inputs)
# We discard `encoder_outputs` and only keep the states.
encoder_states = [state_h, state_c]

# Set up the decoder, using `encoder_states` as initial state.
decoder_inputs = Input(shape=(None, num_decoder_tokens))
# We set up our decoder to return full output sequences,
# and to return internal states as well. We don't use the
# return states in the training model, but we will use them in inference.
decoder_lstm = LSTM(latent_dim, return_sequences=True, return_state=True)
decoder_outputs, _, _ = decoder_lstm(decoder_inputs,
                                     initial_state=encoder_states)
decoder_dense = Dense(num_decoder_tokens)
decoder_outputs = decoder_dense(decoder_outputs)

# Define the model that will turn
# `encoder_input_data` & `decoder_input_data` into `decoder_target_data`
model = Model([encoder_inputs, decoder_inputs], decoder_outputs)

EPOCHS = 2

# Run training
model.compile(optimizer='rmsprop', loss='mae')
history = model.fit([encoder_input_data, decoder_input_data], decoder_target_data,
                    batch_size=batch_size,
                    epochs=EPOCHS,
                    validation_split=0.2)
# Save model
model.save('s2s.h5')

# Plot training & validation loss values
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Model loss during training')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Val'], loc='upper left')
plt.show()

# Next: inference mode (sampling).
# Here's the drill:
# 1) encode input and retrieve initial decoder state
# 2) run one step of decoder with this initial state
# and a "start of sequence" token as target.
# Output will be the next target token
# 3) Repeat with the current target token and current states

# Define sampling models
encoder_model = Model(encoder_inputs, encoder_states)

decoder_state_input_h = Input(shape=(latent_dim,))
decoder_state_input_c = Input(shape=(latent_dim,))
decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]
decoder_outputs, state_h, state_c = decoder_lstm(
	decoder_inputs, initial_state=decoder_states_inputs)
decoder_states = [state_h, state_c]
decoder_outputs = decoder_dense(decoder_outputs)
decoder_model = Model(
	[decoder_inputs] + decoder_states_inputs,
	[decoder_outputs] + decoder_states)


def decode_sequence(input_seq, encoder, decoder):
	# Encode the input as state vectors.
	states_value = encoder.predict(input_seq)

	# Populate the first character of target sequence with the start character.
	target_seq = np.zeros((1, 1, num_decoder_tokens))
	target_seq[0, 0, :] = input_seq[0, -1, 0:2]

	# Sampling loop for a batch of sequences
	# (to simplify, here we assume a batch of size 1).
	stop_condition = False
	decoded_sentence = np.zeros((max_decoder_seq_length, num_decoder_tokens))
	step = 0
	while not stop_condition:
		# Exit condition: either hit max length
		# or find stop character.
		if step >= max_decoder_seq_length - 1:
			stop_condition = True

		output_tokens, h, c = decoder.predict(
			[target_seq] + states_value)

		# Sample a token
		sampled_output = output_tokens[0, -1, :]

		decoded_sentence[step, :] = sampled_output

		# Update the target sequence (of length 1).
		target_seq = output_tokens

		# Update states
		states_value = [h, c]

		step += 1

	return decoded_sentence


def multi_step_plot(history, true_future, prediction):
	plt.figure(figsize=(12, 6))

	if history.ndim == 3:
		history = history[0]
	if true_future.ndim == 3:
		true_future = true_future[0]

	plt.plot(history[:, 0], history[:, 1], 'bx--', label='History')
	plt.plot(history[:, 4], history[:, 5], 'b^--', label='Other Agent History')
	plt.plot(true_future[:, 0], true_future[:, 1], 'bo', label='True Future')

	if prediction.any():
		plt.plot(prediction[:, 0], prediction[:, 1], 'ro',
		         label='Predicted Future')

	plt.axis('equal')
	plt.legend(loc='upper left')
	plt.show()


np.random.seed(1)
num_to_plot = 3
inds = np.arange(encoder_input_data.shape[0])
inds_to_plot = np.random.choice(inds, num_to_plot)
for seq_index in inds_to_plot:
	# Take one sequence (part of the training set)
	# for trying out decoding.
	input_seq = encoder_input_data[seq_index: seq_index + 1]
	true_output = decoder_target_data[seq_index:seq_index + 1]
	predicted_output = decode_sequence(input_seq, encoder_model, decoder_model)
	multi_step_plot(input_seq, true_output, predicted_output)

build/lib/seq_to_seq_ped.py
eliatrevisan/IA-MPPI-LBM
import sys
import os

sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pickle as pkl
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import random
import progressbar
import scipy.io as sio
import json
import importlib
from copy import deepcopy
if sys.version_info[0] < 3:
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from Performance import *
	from utils import *
	import Recorder as rec
	from social_force_model import ForceModel
else:
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.Performance import *
	from src.data_utils.utils import *
	from src.data_utils.Recorder import Recorder as rec
	from src.simulators.social_force_model import *
	import src.data_utils.Trajectory as traj


# Model directories
def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default="VGDNN_simple")
	parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=9)
	parser.add_argument('--n_samples', help='Number of samples', type=int, default=10)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
	                    type=str, default="datasets/ewap_dataset/seq_eth")
	parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool,
	                    default=True)
	parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool,
	                    default=True)
	parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float,
	                    default=0)
	parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float,
	                    default=5)
	parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--update_state', help='update_state', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--freeze_other_agents', help='FReeze other agents.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,
	                    default=False)
	args = parser.parse_args()

	return args


test_args = parse_args()

if test_args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

cwd = os.getcwd()

model_path = os.path.normpath(cwd + '/../') + '/trained_models/' + test_args.model_name + "/" + str(test_args.exp_num)

print("Loading data from: '{}'".format(model_path))
file = open(model_path + '/model_parameters.pkl', 'rb')
if sys.version_info[0] < 3:
	model_parameters = pkl.load(file)  # ,encoding='latin1')
else:
	model_parameters = pkl.load(file, encoding='latin1')
file.close()
args = model_parameters["args"]

with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

# change some args because we are doing inference
args.truncated_backprop_length = 1
args.batch_size = 1
args.keep_prob = 1.0

training_scenario = args.scenario
args.scenario = test_args.scenario
args.real_world_data = test_args.real_world_data
args.dataset = '/' + args.scenario + '.pkl'
data_prep = dhlstm.DataHandlerLSTM(args)
# Only used to create a map from png
map_args = {"file_name": 'map.png',
	            "resolution": 0.1,
	            "map_size": np.array([40., 7.]), }

data_prep.processData(**map_args)

data_path =args.data_path + args.scenario
grid = load_map(data_path)
sf_model = ForceModel(prediction_steps=12, occ_grid=grid,ped_num=2)

# Import model
mod = importlib.import_module("src.models." + args.model_name)
print("Importing : " + args.model_name)
model = mod.NetworkModel(args)

# Lists for logging of the input / output data of the model
input_list = []
grid_list = []
goal_list = []
ped_grid_list = []
y_ground_truth_list = []
y_pred_list = []  # uses ground truth as input at every step
other_agents_list = []
all_predictions = []
trajectories = []
batch_y = []
batch_loss = []

config = tf.ConfigProto(
	device_count={'GPU': 0}
)


def clear_plot():
	axarr[0].clear()
	axarr[0].set_title('Training Sample')

if test_args.unit_testing:
	data_handler = dhlstm.DataHandlerLSTM(args)
	data_handler.unit_test_data_(map_args)

with tf.Session(config=config) as sess:
	model.warmstart_model(args, sess)
	try:
			model.warmstart_convnet(args, sess)
	except:
		print("")

	initial_poses = [np.array([[-15.0, 0], [-12.0, 0]]), np.array([[-12.0, 0], [-15.0, 0]]), np.array([[15.0, 0], [-15.0, 0]])]
	initial_velocities = [np.array([[1.0, 0], [1.0, 0]]),np.array([[1.0, 0], [1.0, 0]]),np.array([[-1.0, 0], [1.0, 0]])]
	all_goals = [np.array([[15.0, 0], [15.0, 0]]),np.array([[15.0, 0], [15.0, 0]]),np.array([[-15.0, 0], [15.0, 0]])]

	for exp_idx in range(len(initial_poses)):

		initial_pose = initial_poses[exp_idx]
		initial_velocity = initial_velocities[exp_idx]
		goals = all_goals[exp_idx]
		next_poses = initial_pose
		next_speed = initial_velocity
		trajectory = traj.Trajectory(goal=goals[0])
		pose = np.zeros([1, 3])
		speed = np.zeros([1, 3])
		pose[0, :2] = initial_pose[0]
		speed[0, :2] = initial_velocity[0]
		trajectory.addData(0, pose, speed, goals[0])
		current_pos = next_poses[0]
		current_vel = next_speed[0]

		# Variables for Post-Processing
		predictions = []
		x_input_series = np.concatenate((np.expand_dims(current_pos,axis=0),np.expand_dims(current_vel,axis=0)),axis=1)
		goal_input_series = np.zeros([1, 2])
		y_ground_truth_series = np.zeros([1, args.prediction_horizon * 2])
		other_agents = []
		y_pred_series = np.zeros([1, args.n_mixtures * args.prediction_horizon * args.output_pred_state_dim])
		grid_input_series = np.zeros([1, 1,60,60])
		ped_grid_series = np.zeros([1,1,args.n_other_agents,args.pedestrian_vector_dim*args.prediction_horizon])
		model.reset_test_cells(np.ones((args.batch_size)))
		cell_state_list= []
		cell_ped_list = []
		cell_concat_list = []

		sf_model.set_initial_pose(next_poses, next_speed, goals)

		for step in range(50):
			samples = []
			if step > 0:
				x_input_series = np.append(x_input_series, np.concatenate(
					(np.expand_dims(current_pos, axis=0), np.expand_dims(current_vel, axis=0)), axis=1), axis=0)

			# Assemble feed dict for training

			x = np.reshape(next_speed[0],(1,1,-1))
			agents = np.zeros([1,1,args.n_other_agents,args.pedestrian_vector_dim*args.prediction_horizon])
			# agent zero is the predicting agent
			current_pos = next_poses[0]
			current_vel = next_speed[0]
			for pred_step in range(args.prediction_horizon):
					for ag_id in range(min([next_poses.shape[0]-1, args.n_other_agents])):
						next_pose = np.expand_dims(next_poses[ag_id+1, :2] + args.dt * next_speed[ag_id+1, :2] * (pred_step + 1),axis=0)
						next_pose_rotated = sup.positions_in_local_frame(current_pos, 0, next_pose)
						agents[0, 0, ag_id, args.pedestrian_vector_dim * pred_step:args.pedestrian_vector_dim * pred_step+2] = \
									next_pose_rotated[0,:]#np.tanh(/ (np.abs(dx) + 0.1)
						velocity_rotated = sup.positions_in_local_frame(current_vel, 0, np.expand_dims(next_speed[ag_id+1, :2],axis=0))
						agents[0, 0, ag_id, args.pedestrian_vector_dim * pred_step + 2:args.pedestrian_vector_dim * pred_step + 4] = \
									velocity_rotated[0,:]
					for ag_id in range(min([next_poses.shape[0]-1, args.n_other_agents]), args.n_other_agents):
						agents[0, 0, ag_id, args.pedestrian_vector_dim * pred_step] = -1
						agents[0, 0, ag_id, args.pedestrian_vector_dim * pred_step + 1] = -1
						agents[0, 0, ag_id, args.pedestrian_vector_dim * pred_step + 2] = 0
						agents[0, 0, ag_id, args.pedestrian_vector_dim * pred_step + 3] = 0

			grid = np.expand_dims(np.expand_dims(data_handler.agent_container.occupancy_grid.getSubmapByCoords(current_pos[0],
																																			 current_pos[1],
																																			 args.submap_width,
																																			 args.submap_height),axis=0),axis=0)

			feed_dict_ = model.feed_test_dic(x, grid, agents, 0)

			# Append to logging series
			if step == 0:
				grid_input_series = grid[0]
			else:
				grid_input_series = np.append(grid_input_series, grid[0], axis=0)
			ped_grid_series = np.append(ped_grid_series, agents, axis=0)

			y_model_pred, output_decoder, cell_outputs_series_state, cell_outputs_series_lstm_ped, cell_concat = model.predict(sess, feed_dict_, True)

			pred_poses, pred_velocities = sf_model.predict(0, next_poses[0], next_speed[0], goals[0], args.dt)

			next_poses, next_speed = sf_model.simulation_step(args.dt)

			# does not work yet to propagate the dynamics of the network
			#next_poses[0,:] = x_input_series[step,:2] + args.dt*y_model_pred[0][0,:2]
			#next_speed[0,:] = args.dt*y_model_pred[0][0,:2]

			y_ground_truth_series = np.append(y_ground_truth_series, np.reshape(pred_velocities,newshape=(1,-1)), axis=0)
			other_agents.append(deepcopy(next_poses[1:]))
			pose[0, :2] = next_poses[0]
			speed[0, :2] = next_speed[0]
			cell_state_list.append(cell_outputs_series_state[0][0,:])
			cell_ped_list.append(cell_outputs_series_lstm_ped[0][0, :])
			cell_concat_list.append(cell_concat[0][0, :])
			samples.append(y_model_pred[0])

			# append data sample to trajectory (timestamp in [ns], pose and vel as np 1x3 arrays)
			trajectory.addData((step+1)*args.dt, pose, speed, goals[0])

			"""
			output_grid = output_decoder[0, :, :, 0]
			clear_plot()
			sup.plotGrid(batch_grid[0, step, :, :], axarr[0], color='k', alpha=0.5)
			sup.plotGrid(output_grid, axarr[1], color='r', alpha=0.5)
			fig.canvas.draw()
			pl.show(block=False)
			"""

			predictions.append(samples)

		latent_analysis = False
		if latent_analysis:

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			pl.show(block=False)
			x_state = np.arange(0, args.rnn_state_size, 1)
			x_ped = np.arange(0, args.rnn_state_size_lstm_ped * 2, 1)
			x_concat = np.arange(0, args.rnn_state_size_lstm_concat, 1)

			time = np.arange(0, len(cell_state_list), 1)
			R= np.array(cell_state_list)
			ax.clear()
			X, Y = np.meshgrid(x_state,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			save_img_to_file = model_path + "/figs"
			if not os.path.exists(save_img_to_file):
				os.makedirs(save_img_to_file)
			pl.savefig(save_img_to_file+"/cell_state" + str(exp_idx) + '.png')
			pl.close()

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			time = np.arange(0, len(cell_ped_list), 1)
			R= np.array(cell_ped_list)
			ax.clear()
			X, Y = np.meshgrid(x_ped,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			pl.savefig(save_img_to_file+"/cell_ped" + str(exp_idx) + '.png')
			pl.close()

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			time = np.arange(0, len(cell_concat_list), 1)
			R= np.array(cell_concat_list)
			ax.clear()
			X, Y = np.meshgrid(x_concat,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			pl.savefig(save_img_to_file+"/cell_concat" + str(exp_idx) + '.png')
			pl.close()
		trajectories.append(trajectory)
		all_predictions.append(predictions)
		input_list.append(x_input_series)
		goal_list.append(goal_input_series)
		grid_list.append(grid_input_series)
		ped_grid_list.append(ped_grid_series)
		y_ground_truth_list.append(y_ground_truth_series)
		other_agents_list.append(other_agents)
		# update progress bar

sess.close()

# Save data
if not os.path.exists(args.model_path + '/../results/'):
	os.makedirs(args.model_path + '/../results/')
results= {
	"trajectories": trajectories,
	"predictions": all_predictions
}
scenario = args.scenario.split('/')[-1]
results_file = args.model_path + '/../results/'+scenario + str(args.exp_num) + "_results.mat"
if os.path.exists(results_file):
	sio.savemat(results_file,results)
else:
	print("Directory not found to save results: " + results_file)

if test_args.record:
		recorder = rec(args, data_prep.agent_container.occupancy_grid)
		if ( "real_world" in test_args.scenario) and not test_args.unit_testing:
			print("Real data!!")
			recorder.plot_on_video(input_list, grid_list, ped_grid_list, all_predictions, y_ground_truth_list,
			                       other_agents_list,
			                       trajectories, test_args)
		else:
			recorder.plot_on_image(input_list, grid_list, ped_grid_list, all_predictions, y_ground_truth_list, other_agents_list,
				                       trajectories,test_args)
			recorder.animate_local(input_list, grid_list, ped_grid_list, all_predictions, y_ground_truth_list, other_agents_list,
		                 trajectories,test_args)
		print("Recorder is done!")
else:
	"""Performance tests"""
	print("Average loss: " + str(np.mean(np.asarray(batch_loss))))
	pred_error, pred_error_summary_lstm = compute_trajectory_prediction_mse(args, trajectories, all_predictions)
	pred_fde, pred_error_summary_lstm_fde = compute_trajectory_fde(args, trajectories, all_predictions)
	args.scenario = training_scenario
	write_results_summary(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde), 0, args, test_args)

build/lib/simulate_VGDNN.py
eliatrevisan/IA-MPPI-LBM
# -*- coding: utf-8 -*-

from math import floor, exp, sqrt
import numpy as np
from src.simulators import social_forces_utils as utils
import sys
from os.path import dirname, abspath
import pylab as pl
import random
import csv
import time

def load_map(data_path):
    try:
        file_path = data_path + '/map.npy'
        if sys.version_info[0] < 3:
            grid_map = np.load(file_path,allow_pickle=True)[()]
        else:
            grid_map = np.load(file_path,encoding='latin1',allow_pickle=True)[()]
        print("Map loaded")
        return grid_map
    except IOError:
        print("ERROR: Cannot load map.")


class ForceModel(object):
    """
    This class is used to simulate and predict social force model.
    """

    def __init__(self, prediction_steps, occ_grid,ped_num,dt=0.1):
        self.desired_speed = 1.0
        self.rel_time = 0.5
        self.delta_t = 0.5
        self.tol = 0.3
        self.ped_num = ped_num
        self.prediction_steps = prediction_steps
        self.occ_grid = occ_grid
        self.dt = dt  # Simulation step in seconds
        try:
            self.map_center = self.occ_grid['Size']/2.0
        except:
            self.map_center = self.occ_grid.map_size
        self.corridor = True

    def goal_force(self, pose, current_speed, goal):
        """
        This function calculates the attractive force towards the goal
        """
        direction = utils.calc_direction(current_pose=pose, desired_pose=goal)
        speed_mod = utils.calc_mod(current_speed)

        force = np.zeros((2))
        force = (direction*self.desired_speed - current_speed)/self.rel_time
        return force

    def repulsive_obj_forces(self, pose, speed, goal):
        """
        This function uses the occupancy grid to calculate the repulsive forces
        of obstacles in a given pose.
        It is not done only once and stored so we can avoid the quantization error
        given by the discretization of the map.
        """
        A = 4.3
        B = 1.07
        l = 0.7

        # Rescale position
        position = (np.array(pose) + self.map_center)/self.occ_grid['Resolution']

        # Get indexes for map
        position_idx = np.ndarray.astype(position, 'int')

        closest_obj = np.array((self.occ_grid['Closest X'][tuple(position_idx)],
                                self.occ_grid['Closest Y'][tuple(position_idx)]))

        rel_distance = pose - closest_obj
        step = - np.array(speed) * self.delta_t
        r_s = rel_distance - step

        mod_rd = utils.calc_mod(rel_distance)
        mod_step = utils.calc_mod(step)
        mod_rs = utils.calc_mod(r_s)

        b = sqrt( pow(mod_rd + mod_rs, 2) - pow(mod_step, 2) )/2

        force = np.zeros((2))
        # if mod_rd != 0 and mod_rs != 0 and mod_step != 0 and b != 0:
        g = A * exp(-b/B) * (mod_rd + mod_rs)/(2*b)
        g = g*(rel_distance/mod_rd + r_s/mod_rs)/2

        force = g


        return force

    def repulsive_ped_force(self, pose, speed, goal, other_ped):
        """
        This function calculates the repulsive forces exerted by other pedestrian.
        """
        A = 4.3
        B = 1.07
        l = 0.7
        # other_ped is in the format [[pose_x, pose_y], [vel_x, vel_y], [goal_x, goal_y]]
        r_ij = np.array(pose) - np.array(other_ped[0])
        s_ij = (np.array(other_ped[1]) - np.array(speed)) * self.delta_t
        r_s_ij = r_ij - s_ij

        mod_r_ij = utils.calc_mod(r_ij)
        mod_r_s_ij = utils.calc_mod(r_s_ij)
        mod_s_ij = utils.calc_mod(s_ij)

        b = sqrt(pow(mod_r_ij + mod_r_s_ij, 2) - pow(mod_s_ij, 2))/2
        force = np.zeros((2))
        if (mod_r_ij != 0) and (mod_r_s_ij != 0) and (mod_s_ij != 0) and (b != 0):
            g = A * exp(-b/B) * (mod_r_ij + mod_r_s_ij) / (2*b)
            g = g * (r_ij/mod_r_ij + r_s_ij/mod_r_s_ij)/2

            e_t_i = utils.calc_direction(current_pose=pose,
                                         desired_pose=goal)

            cos_phi = (e_t_i[0]*r_ij[0]/mod_r_ij) + (e_t_i[1]*r_ij[1]/mod_r_ij)
            w = l + (1 - l)*(1 + cos_phi)/2

            force = g * w

        return force

    def predict(self, ped_id, initial_pose, current_speed, goal, dt):
        """
        This function predicts the next positions and speed of the pedestrian
        using the social force model.
        """
        # data must be in the shape [current + predicted_time_step, ped_index, pose_speed, XY_dims]
        current_pose = initial_pose
        pred_velocities = np.zeros((self.prediction_steps, 2))
        pred_poses = np.zeros((self.prediction_steps, 2))

        for i in range(self.prediction_steps):
            force = self.goal_force(current_pose, current_speed, goal)
            force = force + self.repulsive_obj_forces(pose=current_pose,
                                                      speed=current_speed,
                                                      goal=goal)

            for pp in range(self.ped_num):
                if pp != ped_id:
                    ped_data = [self.current_poses[pp], self.current_speed[pp], self.goals[pp]]
                    force = force + self.repulsive_ped_force(pose=initial_pose,
                                                                         speed=current_speed,
                                                                         goal=goal,
                                                                         other_ped=ped_data)

            current_speed = current_speed + dt * force # Integrate speed
            mod_speed = utils.calc_mod(current_speed)
            max_speed = 1.3 * self.desired_speed
            if mod_speed >= max_speed:
                current_speed = current_speed * max_speed/mod_speed
            current_pose = current_pose + dt * current_speed # Integrate pose

            pred_velocities[i] = current_speed
            pred_poses[i] = current_pose
        return pred_poses, pred_velocities

    def set_initial_pose(self,pose,speed,goals):
        self.initial_pose = pose
        self.current_speed = speed
        self.goals = goals
        self.current_poses = pose

    def get_initial_poses(self, ped_num):
        """
        This functions calculates the random initial positions for the
        pedestrians, making sure they do not spawn into obstacles
        """
        index = 0
        poses = np.zeros((ped_num, 2))
        while index < ped_num:
            if (index % 2):
                pose = np.array((random.uniform(-9.9, -9),
                                 random.uniform(-0.5, 0.5)))
            else:
                pose = np.array((random.uniform(9, 9.9),
                                 random.uniform(-0.5, 0.5)))

            position = (pose + self.map_center)/self.occ_grid['Resolution']
            obstacle = self.occ_grid['Map'][int(position[0]), int(position[1])]
            occupied = utils.check_if_occupied(pose, poses[:index+1])
            if obstacle != 1 and not occupied:
                poses[index] = pose
                index = index + 1
        return poses

    def get_goal(self, pose=None, corridor=False):
        done = False

        if corridor:
            if pose[0] >= 0:
                x_range = np.array([-9.9, -9])
                y_range = np.array([-0.5, 0.5])
            elif pose[0] <= 0:
                x_range = np.array([9, 9.9])
                y_range = np.array([-0.5, 0.5])
        else:
            x_range = np.array([-self.occ_grid['Size'][0]/4,
                                self.occ_grid['Size'][0]/4])
            y_range = np.array([-self.occ_grid['Size'][1]/4,
                                self.occ_grid['Size'][1]/4])

        while not done:
            goal = np.array((random.uniform(x_range[0], x_range[1]),
                             random.uniform(y_range[0], y_range[1])))

            # Check that the goal is not into any obstacle
            position = (goal + self.map_center)/self.occ_grid['Resolution']
            obstacle = self.occ_grid['Map'][int(position[0]), int(position[1])]
            position_idx = np.ndarray.astype(position, 'int')
            closest_obj = np.array((self.occ_grid['Closest X'][tuple(position_idx)],
                                    self.occ_grid['Closest Y'][tuple(position_idx)]))
            if obstacle != 1 or utils.calc_mod(closest_obj - goal) >= 0.5:
                done = True
        return goal

    def simulate(self, ped_num, duration, data_path):
        """
        This function is used to simulate SF model in the given map
        for the given time.
        Input:
            ped_num: Number of pedestrians to simulate
            time: [hours, minutes, seconds]
        """
        self.ped_num = ped_num
        self.min_time = duration[0] * 60 + duration[1]
        sec_time = self.min_time * 60 + duration[2]
        #print("Simulation length: {}h {}m {}s".format(duration[0],
        #                                              duration[1],
        #                                              duration[2]))
        dt = 0.1 # Simulation step in seconds
        sim_elapsed_mins = 0

        self.csv_files = []
        self.writers = []
        self.fieldnames = ['Pedestrian', 'Time s', 'Time ns',
                           'Position X', 'Position Y',
                           'Velocity X', 'Velocity Y',
                           'Goal X', 'Goal Y']

        random.seed()
        counts = np.zeros((self.ped_num))
        poses = self.get_initial_poses(self.ped_num)
        speed = np.zeros((self.ped_num, 2))
        goals = np.zeros((self.ped_num, 2))
        forces = np.zeros((self.ped_num, 2))

        for i in range(self.ped_num):
            goals[i] = self.get_goal(poses[i], corridor=self.corridor)
            # Create logger files
            try:
                self.csv_files.append(open('{}/log_ped{}.csv'.format(data_path, i), 'wb'))
            except:
                print("Cannot open log file number {}. Quitting.".format(i))
                return
            self.writers.append(csv.DictWriter(self.csv_files[i], fieldnames=self.fieldnames))
            self.writers[i].writeheader()
        # Create total log file
        try:
            self.csv_files.append(open('{}/total_log.csv'.format(data_path), 'wb'))
        except:
            print("Cannot open total log file. Quitting.".format(i))
            return
        self.total_log_writer = csv.DictWriter(self.csv_files[-1], fieldnames=self.fieldnames)
        self.total_log_writer.writeheader()

        # Now we are ready to start simulate
        t_start = time.time()
        time_steps = int(sec_time/dt)
        for ts in range(time_steps):
            # NB this could be made for all the peds at once but the force functions
            # for now are designed to work with just one ped at a time
            for ped in range(self.ped_num):
                forces[ped] = self.goal_force(pose=poses[ped],
                                              current_speed=speed[ped],
                                              goal=goals[ped])

                forces[ped] = forces[ped] + self.repulsive_obj_forces(pose=poses[ped],
                                                                      speed=speed[ped],
                                                                      goal=goals[ped])
                for pp in range(self.ped_num):
                    if pp != ped:
                        ped_data = [poses[pp], speed[pp], goals[pp]]
                        forces[ped] = forces[ped] + self.repulsive_ped_force(pose=poses[ped],
                                                                             speed=speed[ped],
                                                                             goal=goals[ped],
                                                                             other_ped=ped_data)
                # Add some noise
                forces[ped] = forces[ped]
            # Once we have all the forces, we update velocity for each ped

            speed = speed + dt * (forces + np.random.normal(0, 0.3, (self.ped_num, 2)))

            for ped in range(self.ped_num):

                mod_speed = utils.calc_mod(speed[ped])
                max_speed = 1.3 * self.desired_speed
                if mod_speed >= max_speed:
                    speed[ped] = speed[ped] * max_speed/mod_speed
                new_pose = poses[ped] + dt * speed[ped]

                # Check if we end up into an obstacle
                # Rescale position
                position = (np.array(new_pose) + self.map_center)/self.occ_grid['Resolution']
                # Get indexes for map
                position = np.ndarray.astype(position, 'int')
                # If we would end up into an obst we do not move
                if self.occ_grid['Map'][tuple(position)] == 1:
                    print("Inside obstacle")
                    new_pose = poses[ped]
                    speed[ped] = np.zeros((2))

                # Update poses
                poses[ped] = new_pose
                # Save data
                data_dict = {}
                data_dict['Pedestrian'] = ped
                data_dict['Time s'] = int(ts*dt)
                data_dict['Time ns'] = int((ts*dt * 10**9)%10**9)
                data_dict['Position X'] = new_pose[0]
                data_dict['Position Y'] = new_pose[1]
                data_dict['Velocity X'] = speed[ped][0]
                data_dict['Velocity Y'] = speed[ped][1]
                data_dict['Goal X'] = goals[ped][0]
                data_dict['Goal Y'] = goals[ped][1]
                self.writers[ped].writerow(data_dict)
                self.total_log_writer.writerow(data_dict)
                # Update goal
                if utils.calc_mod(goals[ped] - new_pose) <= self.tol:
                    goals[ped] = self.get_goal(new_pose, corridor=self.corridor)
                    counts[ped] = 0
                # Check if stuck
                if utils.calc_mod(speed[ped]) <= 0.01:
                    counts[ped] = counts[ped] + 1
                    if counts[ped] == 5:
                        goals[ped] = self.get_goal(new_pose, corridor=self.corridor)
                        counts[ped] = 0
                else: # This way we just count when they are consecutive
                    counts[ped] = 0

            # Log passing time
            rt_time = time.time() - t_start
            if (ts*dt/60 - sim_elapsed_mins) >= 1:
                sim_elapsed_mins = int(ts*dt/60)
                fact = ts*dt / rt_time
                rt_time = time.gmtime(rt_time)
                print("Elapsed simulation time since start: {} mins".format(sim_elapsed_mins))
                print("Elapsed real time since start: {}h {}m {}s".format(rt_time.tm_hour,
                                                                          rt_time.tm_min,
                                                                          rt_time.tm_sec))
                print("Sim to RT factor: {}".format(fact))
                print("")

        #
        for i in range(self.ped_num + 1):
            self.csv_files[i].close()

        print("Simulation finished.")
        rt_time = time.gmtime(time.time() - t_start)
        print("Total real time: {}h {}m {}s".format(rt_time.tm_hour,
                                                    rt_time.tm_min,
                                                    rt_time.tm_sec))
    def simulate_ped(self, ped_id,ini_pose,goal,duration,other_agents_pos):
        """
        This function is used to simulate SF model in the given map
        for the given time.
        Input:
            ped_num: Number of pedestrians to simulate
            time: [hours, minutes, seconds]
        """

        self.min_time = duration[0] * 60 + duration[1]
        sec_time = self.min_time * 60 + duration[2]
        print("Simulation length: {}h {}m {}s".format(duration[0],
                                                      duration[1],
                                                      duration[2]))
        dt = 0.3 # Simulation step in seconds
        sim_elapsed_mins = 0

        random.seed()
        counts = np.zeros((self.ped_num))
        poses = np.expand_dims(ini_pose[0:2].copy(),0)
        speed = np.expand_dims(ini_pose[2:].copy(),0)
        goals = np.zeros((self.ped_num, 2))
        forces = np.zeros((self.ped_num, 2))

        for i in range(self.ped_num):
            goals[i] = goal

        # Now we are ready to start simulate
        t_start = time.time()
        time_steps = int(sec_time/dt)
        traj = np.zeros((time_steps,2))
        velocities = np.zeros((time_steps,2))
        velocities[0,:] = speed
        traj[0,:] = ini_pose[0:2]
        for ts in range(time_steps-1):
            # NB this could be made for all the peds at once but the force functions
            # for now are designed to work with just one ped at a time
            ped=0
            forces[ped] = self.goal_force(pose=poses[ped],
                                         current_speed=speed[ped],
                                         goal=goals[ped])

            forces[ped] = forces[ped] + self.repulsive_obj_forces(pose=poses[ped],
                                                                  speed=speed[ped],
                                                                  goal=goals[ped])
            for pp in range(len(other_agents_pos)):
                agent_pos = np.reshape(other_agents_pos[pp],2)
                ped_data = [agent_pos, np.zeros((2)), agent_pos]
                forces[ped] = forces[ped] + self.repulsive_ped_force(pose=poses[ped],
                                                                     speed=speed[ped],
                                                                     goal=goals[ped],
                                                                     other_ped=ped_data)
            # Add some noise
            forces[ped] = forces[ped]
            # Once we have all the forces, we update velocity for each ped

            speed = speed + dt * (forces + np.random.normal(0, 0.01, (self.ped_num, 2)))

            mod_speed = utils.calc_mod(speed[ped])
            max_speed = 1.3 * self.desired_speed
            if mod_speed >= max_speed:
                speed[ped] = speed[ped] * max_speed/mod_speed
            new_pose = poses[ped] + dt * speed[ped]

            # Check if we end up into an obstacle
            # Rescale position
            position = (np.array(new_pose) + self.map_center)/self.occ_grid['Resolution']
            # Get indexes for map
            position = np.ndarray.astype(position, 'int')
            # If we would end up into an obst we do not move
            if self.occ_grid['Map'][tuple(position)] == 1:
                new_pose = poses[ped]
                speed[ped] = np.zeros((2))

            # Update poses
            poses[ped] = new_pose

            # Update goal
            if utils.calc_mod(goals[ped] - new_pose) <= self.tol:
                goals[ped] = self.get_goal(new_pose, corridor=self.corridor)
                counts[ped] = 0
            # Check if stuck
            if utils.calc_mod(speed[ped]) <= 0.01:
                counts[ped] = counts[ped] + 1
                if counts[ped] == 5:
                    goals[ped] = self.get_goal(new_pose, corridor=self.corridor)
                    counts[ped] = 0
            else: # This way we just count when they are consecutive
                counts[ped] = 0

            traj[ts+1,:]=poses[ped]
            velocities[ts+1,:]=speed[ped]
        return traj, velocities
    
    def simulation_step(self,sec_time):
        """
        This function is used to simulate one step of SF model in the given map
        for the given time.
        Input:
            ped_num: Number of pedestrians to simulate
            time: [hours, minutes, seconds]
        """

        forces = np.zeros((self.ped_num, 2))

        # Now we are ready to start simulate
        time_steps = int(sec_time/self.dt)
        for ts in range(time_steps):
            # NB this could be made for all the peds at once but the force functions
            # for now are designed to work with just one ped at a time
            for ped in range(self.ped_num):
                forces[ped] = self.goal_force(pose=self.current_poses[ped],
                                              current_speed=self.current_speed[ped],
                                              goal=self.goals[ped])

                forces[ped] = forces[ped] + self.repulsive_obj_forces(pose=self.current_poses[ped],
                                                                      speed=self.current_speed[ped],
                                                                      goal=self.goals[ped])
                for pp in range(self.ped_num):
                    if pp != ped:
                        ped_data = [self.current_poses[pp], self.current_speed[pp], self.goals[pp]]
                        forces[ped] = forces[ped] + self.repulsive_ped_force(pose=self.current_poses[ped],
                                                                             speed=self.current_speed[ped],
                                                                             goal=self.goals[ped],
                                                                             other_ped=ped_data)
                # Add some noise
                forces[ped] = forces[ped]
            # Once we have all the forces, we update velocity for each ped

            self.current_speed = self.current_speed + self.dt * (forces + np.random.normal(0, 0.1, (self.ped_num, 2)))

            for ped in range(self.ped_num):

                mod_speed = utils.calc_mod(self.current_speed[ped])
                max_speed = 1.3 * self.desired_speed
                if mod_speed >= max_speed:
                    self.current_speed[ped] = self.current_speed[ped] * max_speed/mod_speed
                new_pose = self.current_poses[ped] + self.dt * self.current_speed[ped]

                # Check if we end up into an obstacle
                # Rescale position
                position = (np.array(new_pose) + self.map_center)/self.occ_grid['Resolution']
                # Get indexes for map
                position = np.ndarray.astype(position, 'int')
                # If we would end up into an obst we do not move
                if self.occ_grid['Map'][tuple(position)] == 1:
                    print("Inside obstacle")
                    new_pose = self.current_poses[ped]
                    self.current_speed[ped] = np.zeros((2))

                # Update poses
                self.current_poses[ped] = new_pose

                # Update goal
                if utils.calc_mod(self.goals[ped] - new_pose) <= self.tol:
                    print("Goal reached")
                # Check if stuck
                if utils.calc_mod(self.current_speed[ped]) <= 0.01:
                    print("Stuck")

        return self.current_poses , self.current_speed

# All the stuff here is for test
if __name__ == "__main__":

    data_path ="../../data/multipath"

    grid = load_map(data_path)
    model = ForceModel(predicted=10, occ_grid=grid,ped_num=2)

    model.simulate(2, [5,0,0], data_path)
build/lib/simulators/social_force_model.py
eliatrevisan/IA-MPPI-LBM
# -*- coding: utf-8 -*-

from math import ceil
import numpy as np
import os
import cv2 as cv
import csv
import pylab as pl
from PIL import ImageDraw
from PIL import Image
from numpngw import write_png

# TODO Merge this functions into Support.py
# ------------------------------------------------------------------------------
def create_map_from_obstacles(map_shape=[30, 30],
                              map_resolution=0.1,
                              save_data_path='../data'):
  grid_map = {}
  grid_map['Resolution'] = map_resolution
  grid_map['Size'] = np.array(map_shape)
  grid_map['Map'] = np.zeros((int(grid_map['Size'][0] / grid_map['Resolution']),
                              int(grid_map['Size'][1] / grid_map['Resolution'])))
  map_center = grid_map['Size'] /2 # other datasets /2

  xx, yy = np.mgrid[:map_shape[0] / map_resolution, :map_shape[1] / map_resolution]
  # First read and draw cylinders
  with open("../objects/cylinders.csv", "rb") as f:
    reader = csv.reader(f, delimiter=",")
    for i, line in enumerate(reader):
      if i == 0:
        continue
      center = np.array((float(line[0]) + map_center[0],
                         float(line[1]) + map_center[1])) / map_resolution
      circle = (xx - int(center[0]))**2 + (yy - int(center[1]))**2
      grid_map['Map'][circle <= int(float(line[2]) / map_resolution)**2] = 1

  # Now go with the rectangles
  with open("../objects/rectangles.csv", "rb") as f:
    reader = csv.reader(f, delimiter=",")
    img = Image.fromarray(grid_map['Map'])
    draw = ImageDraw.Draw(img)
    for i, line in enumerate(reader):
      if i == 0:
        continue
      center = np.array((float(line[0]), float(line[1]))) / map_resolution
      shape = np.array((float(line[2]), float(line[3]))) / map_resolution
      theta = (np.pi / 180.0) * float(line[4])  # In radians
      rect = np.array(((-shape[1] / 2., -shape[0] / 2.),
                       (-shape[1] / 2., shape[0] / 2.),
                       (shape[1] / 2., shape[0] / 2.),
                       (shape[1] / 2., -shape[0] / 2.)))

      R = np.array([[np.cos(theta), -np.sin(theta)],
                    [np.sin(theta), np.cos(theta)]])

      rect = np.dot(rect, R) + center[::-1] + map_center / map_resolution
      draw.polygon([tuple(p) for p in rect], fill=1)
    grid_map['Map'] = np.asarray(img)
  print("Occupancy grid done.")

  # Calculate closest obstacles
  grid_map['Closest X'] = np.zeros_like(grid_map['Map']) + 10000
  grid_map['Closest Y'] = np.zeros_like(grid_map['Map']) + 10000

  # Get idx in map of places where there are obstacles
  obst_idx = np.array(np.where(grid_map['Map'] == 1))

  for xx in range(int(grid_map['Size'][0] / grid_map['Resolution'])):
    for yy in range(int(grid_map['Size'][1] / grid_map['Resolution'])):
      distances = np.sqrt(np.sum(np.square(obst_idx - np.expand_dims([xx, yy], axis=1)), axis=0))
      closest_obj = np.argmin(distances)
      # Get index in map of the closest_obj
      grid_map['Closest X'][xx, yy] = obst_idx[0, closest_obj]
      grid_map['Closest Y'][xx, yy] = obst_idx[1, closest_obj]

  grid_map['Closest X'] = grid_map['Closest X'] * grid_map['Resolution'] - map_center[0]
  grid_map['Closest Y'] = grid_map['Closest Y'] * grid_map['Resolution'] - map_center[1]

  print("Closest obstacles positions calculated.")
  print("Saving map...")
  np.save(os.path.join(save_data_path, 'map'), grid_map)
  print("Map saved.")
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def create_map_from_png(load_data_path='../data/datasets/ewap_dataset/seq_eth',
                        save_data_path='../data'):
  # Create grid for SF model
  print("create_map_from_png")
  grid_map = {}
  grid_map['Resolution'] = 0.1
  grid_map['Size'] = np.array([20., 7.])  # map size in [m]
  grid_map['Map'] = np.zeros((int(grid_map['Size'][0] / grid_map['Resolution']),
                              int(grid_map['Size'][1] / grid_map['Resolution'])))

  map_center = grid_map['Size'] / 2. # hack for my dataset
  H = np.genfromtxt(os.path.join(load_data_path, 'H.txt'),
                    delimiter='  ',
                    unpack=True).transpose()

  # Extract static obstacles
  obst_threshold = 200
  static_obst_img = cv.imread(os.path.join(load_data_path, 'map.png'), 0)
  #static_obst_img = cv.imread(os.path.join(load_data_path, 'map.png'), 0)
  obstacles = np.zeros([0, 3])
  #grid = (static_obst_img/255*-1)+1
  #static_obst_img = np.transpose(static_obst_img)
  for xx in range(static_obst_img.shape[0]):
    for yy in range(static_obst_img.shape[1]):
      if static_obst_img[xx, yy] > obst_threshold:
        obstacles = np.append(obstacles,
                              np.dot(H, np.array([[xx], [yy], [1]])).transpose(),
                              axis=0)
  # Compute obstacles in 2D
  obstacles_2d = np.zeros([obstacles.shape[0], 2])
  obstacles_2d[:, 0] = obstacles[:, 0] / obstacles[:, 2]
  obstacles_2d[:, 1] = obstacles[:, 1] / obstacles[:, 2]

  # Get obstacle idx on map
  obst_idx = []
  for obst_ii in range(obstacles_2d.shape[0]):
    obst_idx.append(idx_from_pos(obstacles_2d[obst_ii, 0],
                    obstacles_2d[obst_ii, 1],
                    map_center,grid_map['Resolution']))
    grid_map['Map'][obst_idx[-1]] = 1

  grid_map['Closest X'] = np.zeros_like(grid_map['Map']) + 10000
  grid_map['Closest Y'] = np.zeros_like(grid_map['Map']) + 10000
  # Calculate closest obstacle
  obst_idx = np.array(obst_idx)
  for xx in range(int(grid_map['Size'][0] / grid_map['Resolution'])):
    for yy in range(int(grid_map['Size'][1] / grid_map['Resolution'])):
      delta_idx = obst_idx - np.array([xx, yy])
      distances = np.sqrt(np.sum(np.square(delta_idx), axis=1))
      closest_obj = np.argmin(distances)
      grid_map['Closest X'][xx, yy] = obst_idx[closest_obj][0]
      grid_map['Closest Y'][xx, yy] = obst_idx[closest_obj][1]

  grid_map['Closest X'] = grid_map['Closest X'] * grid_map['Resolution'] - map_center[0]
  grid_map['Closest Y'] = grid_map['Closest Y'] * grid_map['Resolution'] - map_center[1]

  np.save(os.path.join(save_data_path, 'map'), grid_map)
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def rotateGridAroundCenter(grid, angle):
  # Rotate grid into direction of initial heading
  grid = grid.copy()
  rows, cols = grid.shape
  M = cv.getRotationMatrix2D(center=(rows/2, cols/2), angle=angle, scale=1)
  grid = np.array(cv.warpAffine(grid, M, (rows, cols)))

  grid[grid > 0] = 1 # This way we remove artifacts around borders
  grid[grid < 1] = 0 # This way we remove any negative value

  return grid
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def idx_from_pos(x, y, center, res=0.1):
  idx_x = round((x + float(center[0])) / res)
  idx_y = round((y + float(center[1])) / res)

  # Projecting index on map if out of bounds
  idx_x = max(0, min(idx_x, -1 + center[0] * 2. / res))
  idx_y = max(0, min(idx_y, -1 + center[1] * 2. / res))
  return int(idx_x), int(idx_y)
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_mod(vect):
  vect = np.array(vect)
  return np.sqrt(pow(vect[0], 2) + pow(vect[1], 2))
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_angle(x, y):
  angles = np.array(np.arctan2(y, x) * 180 / np.pi)
  angles[angles == 180.] = -180. # This way output is in [-180,180) and not (-180,180]
  return angles
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_rotation_matrix(angle):
  """
  This function calculates the active rotation matrix.
  Thus if you want to rotate the ref system and not the point you shoud use
  its transpose.
  Thus R = [c, -s; s, c]
  """
  theta = np.array(np.radians(angle))
  c, s = np.cos(theta), np.sin(theta)
  R = np.zeros((len(theta), 2, 2))
  R[:, 0, 0] = c
  R[:, 0, 1] = -s
  R[:, 1, 0] = s
  R[:, 1, 1] = c
  return R
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_direction(current_pose, desired_pose):
  """
  This function calculates the normalized vector going from the current pose
  towards the desired pose.
  """
  tmp = np.array(desired_pose) - np.array(current_pose)
  mod_e_t = calc_mod(tmp)

  e_t = np.zeros((2))

  if mod_e_t is not 0:
    e_t = tmp / mod_e_t

  return e_t
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def check_if_occupied(pose, other_poses):
  """
  This function is used in the SF model to check if a position
  is already occupied by another pedestrian.
  """
  ped_radius = 0.2
  for index, ped_pose in enumerate(other_poses):
    distance = calc_mod(pose - ped_pose)
    if distance <= 2 * ped_radius:
      return True

  return False
# ------------------------------------------------------------------------------

def write_map_png(file, grid):
  grid = rotateGridAroundCenter(grid,90)
  #grid = (grid-1)*-255
  img = grid.astype(np.uint8)
  write_png(file, img, bitdepth=1)

# ------------------------------------------------------------------------------
def load_map(data_path):
  try:
    file_path = data_path + '/map.npy'
    grid_map = np.load(file_path)[()]
    print("Map loaded")
    return grid_map
  except IOError:
    print("ERROR: Cannot load map.")
    return False
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def show_map(grid_map):
  i = 1
  for k in grid_map.keys():
    if k in ["Size", "Resolution"]:
      continue
    fig = pl.figure(i)
    pl.imshow(grid_map[k], cmap=pl.cm.ocean)
    pl.colorbar()
    pl.title(k)
    i = i + 1
  pl.show()
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_direction(current_pose, desired_pose):
    """
    This function calculates the normalized vector going from the current pose
    towards the desired pose.
    """
    tmp = np.array(desired_pose) - np.array(current_pose)
    mod_e_t = calc_mod(tmp)

    e_t = np.zeros((2))

    if mod_e_t != 0:
        e_t = tmp/mod_e_t

    return e_t
# ------------------------------------------------------------------------------
if __name__ == "__main__":
  create_map_from_png(
    load_data_path='/media/bdebrito/7697ec91-468a-4763-b1c3-135caa7f5aed/home/code/lstm_pedestrian_prediction/data/multipath/corridor_3peds',
  save_data_path = '/media/bdebrito/7697ec91-468a-4763-b1c3-135caa7f5aed/home/code/lstm_pedestrian_prediction/data/multipath/corridor_3peds')
  grid_map = load_map("/media/bdebrito/7697ec91-468a-4763-b1c3-135caa7f5aed/home/code/lstm_pedestrian_prediction/data/multipath/corridor_3peds")
  show_map(grid_map)
  print("DOne!")
build/lib/simulators/social_forces_utils.py
eliatrevisan/IA-MPPI-LBM
#!/usr/bin/env python3

import os, sys; print("Running with {}".format(sys.version))
sys.path.append('../')
import copy, time, pickle, math, collections

from scipy.spatial.transform import Rotation
import tensorflow as tf
import numpy as np; np.set_printoptions(suppress=True)
from scipy.stats import multivariate_normal

import matplotlib.pyplot as plt

from models.VGDNN import NetworkModel as SocialVRNN
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import rospy, rospkg
import message_filters
from scipy.spatial.transform import Rotation as R
from nav_msgs.msg import OccupancyGrid, Odometry
from geometry_msgs.msg import Point, PoseStamped
from visualization_msgs.msg import Marker, MarkerArray
from social_vrnn.msg import lmpcc_obstacle_array as LMPCC_Obstacle_Array, svrnn_path as SVRNN_Path, svrnn_path_array as SVRNN_Path_Array
from data_utils.OccupancyGrid import OccupancyGrid as OccGrid

#sys.path.remove("/opt/ros/kinetic/lib/python2.7/dist-packages")
import cv2

PACKAGE_NAME = 'social_vrnn'
MAX_QUERY_AGENTS = 3
SWITCH_AXES = False


class SVRNN_Predictor():
    
    def __init__(self, node_name, visual_node_name):
        self.node_name, self.visual_node_name = node_name, visual_node_name

        # Bind node
        rospy.init_node(self.node_name)
        rospy.on_shutdown(self.shutdown_callback)
        rospy.loginfo('{} has started.'.format(self.visual_node_name))

        # Set up class variables
        self.get_submap = None # function to request submap for some position
        self.agents_pos, self.agents_vel = {}, {} # dictionary of world state

        # Generic marker template
        mrk = Marker()
        mrk.header.frame_id = 'odom'
        mrk.action = Marker.ADD
        mrk.lifetime.secs = int(1.0)
        mrk.scale.x, mrk.scale.y, mrk.scale.z = 1.0, 1.0, 1.0
        mrk.color.r, mrk.color.g, mrk.color.b, mrk.color.a = 1.0, 1.0, 1.0, 1.0
        mrk.pose.orientation.w = 1.0
        self.marker_template = mrk

        # Read parameters
        self.n_agents = rospy.get_param("/n_agents")

        # Initialize map class
        self.occgrid = OccGrid()

        # Load the model
        self.model_args = self.get_model_args('VGDNN', '33')
        self.model, self.tf_session = self.load_model(SocialVRNN, self.model_args)

        # Set up subscribers
        rospy.Subscriber('/roboat_cloud/obstacle/map', OccupancyGrid, self.store_occupancy_grid)

        # Set up publishers
        self.pred_path_publisher = rospy.Publisher('/{}/predictions'.format(self.node_name), SVRNN_Path_Array, latch=True, queue_size=10)
        self.pos_mark_publisher = rospy.Publisher('/{}/markers/positions'.format(self.node_name), MarkerArray, latch=True, queue_size=10)
        self.pred_mark_publisher = rospy.Publisher('/{}/markers/predictions'.format(self.node_name), MarkerArray, latch=True, queue_size=10)


    def infer(self):
        if self.get_submap is None: return None, None
        #print("Agent pos: ", self.agents_pos)
        if len(self.agents_pos) < 2:
            rospy.logwarn('No other agents present, skipping path prediction')
            return None, None

        # Build NumPy matrices from world state
        active_agents = np.full((MAX_QUERY_AGENTS, ), False, dtype=bool)
        for id in self.agents_pos.keys(): active_agents[id] = True

        agents_pos_np = np.zeros((MAX_QUERY_AGENTS, 2), dtype=float)
        for id, pos in self.agents_pos.items(): agents_pos_np[id] = pos

        agents_vel_np = np.zeros((MAX_QUERY_AGENTS, 2), dtype=float)
        for id, vel in self.agents_vel.items(): agents_vel_np[id] = vel
        
        # define multiplier based on dt. (dt = 0.4 becomes 4 since incoming time steps are recorded 0.1 seconds from each other)
        dt = int(self.model_args.dt * 10)

        # Maintain past agents velocities
        if not hasattr(self, '_agents_vel_ls_all'):
            self._agents_vel_ls_all = np.zeros((MAX_QUERY_AGENTS, 2 * (self.model_args.prev_horizon + 1) * dt), dtype=float)
        self._agents_vel_ls_all = np.roll(self._agents_vel_ls_all, 2, axis=1) # Brings oldest velocities to front of array
        self._agents_vel_ls_all[:, :2] = agents_vel_np # Replaces oldest velocities with newest

        # Only use time steps at dt interval
        if not hasattr(self, '_agents_vel_ls'):
            self._agents_vel_ls = np.zeros((MAX_QUERY_AGENTS, 2 * (self.model_args.prev_horizon + 1)), dtype=float)
        self._agents_vel_ls[:, ::2] = self._agents_vel_ls_all[:, ::(dt*2)]
        self._agents_vel_ls[:, 1::2] = self._agents_vel_ls_all[:, 1::(dt*2)]

        # Calculate the relative odometry
        relative_odometry = np.zeros((MAX_QUERY_AGENTS, 4 * self.model_args.n_other_agents), dtype=float)
        for sub_ind in np.arange(MAX_QUERY_AGENTS)[active_agents]:

            # Get n-nearest neighbours
            nn_inds = np.argsort(np.linalg.norm(agents_pos_np[active_agents] - agents_pos_np[sub_ind], axis=1))[1:]
            #if len(nn_inds) < self.model_args.n_other_agents: nn_inds = np.concatenate((nn_inds, np.repeat(nn_inds[-1], self.model_args.n_other_agents - len(nn_inds))))
            #if len(nn_inds) > self.model_args.n_other_agents: nn_inds = nn_inds[:self.model_args.n_other_agents]

            # Calculate relative positions
            nn_pos = agents_pos_np[active_agents][nn_inds].copy()
            #multivars = []
            #for i in range(len(nn_inds)):
            #    multivars.append( multivariate_normal.pdf(  np.linalg.norm(nn_pos - agents_pos_np[sub_ind], axis=1)[i] , mean=0.0, cov=5.0) )
            nn_pos -= agents_pos_np[sub_ind] 
            
            #for i in range(len(nn_inds)):
            #    nn_pos[i] = nn_pos[i] * multivars[i]

            # Calculate relative velocities
            nn_vel = agents_vel_np[active_agents][nn_inds].copy()
            nn_vel -= agents_vel_np[sub_ind]

            #relative_odometry[sub_ind] = np.concatenate((nn_pos, nn_vel,np.linalg.norm(nn_pos),np.arctan2(nn_pos[1], nn_pos[0])), axis=1).flatten()
            #relative_odometry[sub_ind, sub_ind*4:sub_ind*4+((self.n_agents-1)*4)] = np.concatenate([nn_pos, nn_vel], axis=1).flatten()
            relative_odometry[sub_ind, 0:((self.n_agents-1)*4)] = np.concatenate([nn_pos, nn_vel], axis=1).flatten()

        # Get the submaps
        submaps = np.zeros((MAX_QUERY_AGENTS, int(self.model_args.submap_width / self.model_args.submap_resolution), int(self.model_args.submap_height / self.model_args.submap_resolution))) #, dtype=float
        for id in np.arange(MAX_QUERY_AGENTS)[active_agents]:
            if self.get_submap is None: return
            submaps[id] = self.get_submap(agents_pos_np[id], agents_vel_np[id])

        
        #print("Batch_vel : ", self._agents_vel_ls)
        """
        print("batch vel shape: ", self._agents_vel_ls.shape[0], self._agents_vel_ls.shape[1]) 

        print("Relative vector: ", relative_odometry)
        print("batch rel shape: ", relative_odometry.shape[0], relative_odometry.shape[1]) 
        """

        # Predict the future velocities
        return copy.deepcopy(self.agents_pos), self.model.predict(
            self.tf_session,
            self.model.feed_pred_dic(
                batch_vel = self._agents_vel_ls,
                batch_ped_grid = relative_odometry,
                batch_grid = submaps,
                step = 0
            ),
            True
        )[0]

    def rotate_to_world_frame(self, msg):

        # Read current position from message in world frame
        curr_x, curr_y = msg.pose.pose.position.x, msg.pose.pose.position.y
        
        # Read current velocity from message in body frame of roboat
        curr_vx, curr_vy = msg.twist.twist.linear.x, msg.twist.twist.linear.y

        # Transfrom velocity to world frame
        q0, q1, q2, q3 = (msg.pose.pose.orientation.x,
                            msg.pose.pose.orientation.y,
                            msg.pose.pose.orientation.z,
                            msg.pose.pose.orientation.w)
        
        rot_mat = R.from_quat([q0, q1, q2, q3])
        curr_v = np.array([curr_vx, curr_vy, 0])

        # Transformation
        v = rot_mat.apply(curr_v)

        return curr_x, curr_y, v[0], v[1]

    def publish(self, agents_pos, pred_velocities):
        if pred_velocities is None: return

        # Publish predicted positions
        pred_velocities_path = SVRNN_Path_Array()
        pred_velocities_mrk = MarkerArray()
        for id in agents_pos.keys():
            pred = pred_velocities[id]
            path_msg = SVRNN_Path()
            path_msg.id = float(id)
            path_msg.dt = self.model_args.dt
            for mx_id in range(self.model_args.n_mixtures):
                main_mixture = mx_id == 0
                path_mrk = copy.deepcopy(self.marker_template)
                path_mrk.id = self.model_args.n_mixtures * id + mx_id
                path_mrk.type = Marker.LINE_STRIP
                path_mrk.scale.x = 0.1
                path_mrk.pose.position.x = agents_pos[id][1] if SWITCH_AXES else agents_pos[id][0]
                path_mrk.pose.position.y = agents_pos[id][0] if SWITCH_AXES else agents_pos[id][1]
                path_mrk.color.r = float(mx_id == 0)
                path_mrk.color.g = float(mx_id == 1)
                path_mrk.color.b = float(mx_id == 2)
                prev_x, prev_y = 0.0, 0.0
                for ts_id in range(self.model_args.prediction_horizon):
                    pose = PoseStamped()
                    pt = Point()
                    idx = ts_id * self.model_args.output_pred_state_dim * self.model_args.n_mixtures + mx_id
                    pt.x = prev_x + self.model_args.dt * pred[0][idx + (self.model_args.n_mixtures if SWITCH_AXES else 0)]
                    pt.y = prev_y + self.model_args.dt * pred[0][idx + (0 if SWITCH_AXES else self.model_args.n_mixtures)]
                    pose.pose.position = pt
                    if main_mixture: path_msg.path.poses.append(pose)
                    pt.z = 0.2
                    path_mrk.points.append(pt)
                    prev_x, prev_y = pt.x, pt.y
                pred_velocities_mrk.markers.append(path_mrk)
            pred_velocities_path.paths.append(path_msg)
        #print("Pred positions", pred_velocities_path)
        self.pred_path_publisher.publish(pred_velocities_path)
        self.pred_mark_publisher.publish(pred_velocities_mrk)


    def store_occupancy_grid(self, occupancy_grid_msg):
        # Reset class state
        def safedelattr(attr):
            if hasattr(self, attr): delattr(self, attr)
        safedelattr('_roboat_pos'); safedelattr('_roboat_last_update')
        self.agents_pos, self.agents_vel = {}, {}; safedelattr('_agents_last_update')
        self.get_submap = None; safedelattr('_agents_vel_ls')
        rospy.loginfo('Cleared history context')

        # Makes 1 pixel equal to 1 submap pixel
        scale_factor = occupancy_grid_msg.info.resolution / self.model_args.submap_resolution

        # Transform occupancy grid data into Social-VRNN format
        grid_info = occupancy_grid_msg.info
        occupancy_grid = np.asarray(occupancy_grid_msg.data, dtype=float).reshape((grid_info.height, grid_info.width))
        occupancy_grid[occupancy_grid > 0.0] = 1.0
        occupancy_grid[occupancy_grid < 1.0] = 0.0
        occupancy_grid = cv2.flip(occupancy_grid, 0)
        occupancy_grid = cv2.resize(
            occupancy_grid,
            (int(scale_factor * grid_info.width), int(scale_factor * grid_info.height)),
            fx=0, fy=0,
            interpolation=cv2.INTER_NEAREST
        )
        self.occgrid.resolution = self.model_args.submap_resolution
        self.occgrid.gridmap = occupancy_grid
        self.occgrid.map_size = np.array([grid_info.height, grid_info.width])

        # Create custom function for requesting submap
        def get_submap(position, orientation):
            """
            position: (2,) numpy matrix
                The offset of the position in meters from the origin
                Index 0 is x, index 1 is y
            orientation: (2,) numpy matrix
                The orientation of the position (does not have to be normalised)
                Index 0 is x, index 1 is y
            """
            # Translate to pixel coordinates
            origin_offset = position.copy().astype(float)
            origin_offset /= self.model_args.submap_resolution
            origin_offset[0] = origin_offset[0] + (78 / self.occgrid.resolution) 
            origin_offset[1] = self.occgrid.gridmap.shape[0] - origin_offset[1] - (40 / self.occgrid.resolution) 
            origin_offset = origin_offset.astype(int)
                    
            # Do bounds-check, returns submaps with zeros if near border or out of bounds
            if (origin_offset[0] < self.model_args.submap_width/ self.model_args.submap_resolution/2 or
                origin_offset[0] > occupancy_grid.shape[1] - self.model_args.submap_width/ self.model_args.submap_resolution/2 or
                origin_offset[1] < self.model_args.submap_height/ self.model_args.submap_resolution/2 or
                origin_offset[1] > occupancy_grid.shape[0] - self.model_args.submap_height/ self.model_args.submap_resolution/2):
                rospy.logerr('Out-of-bounds submap requested!')
                return np.zeros((int(self.model_args.submap_height / self.model_args.submap_resolution), int(self.model_args.submap_width / self.model_args.submap_resolution)))
            
            return occupancy_grid[
                int(origin_offset[1]-self.model_args.submap_height/ self.model_args.submap_resolution/2) : int(origin_offset[1]+self.model_args.submap_height/ self.model_args.submap_resolution/2),
                int(origin_offset[0]-self.model_args.submap_width/ self.model_args.submap_resolution/2) : int(origin_offset[0]+self.model_args.submap_width/ self.model_args.submap_resolution/2)
            ]

        self.get_submap = get_submap

    def callback(self, *roboat_odometry):

        # Check if attribute exists
        if not hasattr(self, '_roboat_pos'):
            self._roboat_pos = np.zeros((2, ), dtype=float)
            self._roboat_vel = np.zeros((2, ), dtype=float)

        curr_pos, curr_vel = {}, {}

        # Read data from msgs and transform velocity to world frame
        for id in range(self.n_agents):
            
            curr_x, curr_y, curr_vx, curr_vy = self.rotate_to_world_frame(roboat_odometry[id])

            curr_pos[id] = np.array([curr_x, curr_y])
            curr_vel[id] = np.array([curr_vx, curr_vy])

        # Update agents' states
        self.agents_pos, self.agents_vel = curr_pos, curr_vel

        # Make prediction
        agents_pos, pred_velocities = self.infer()

        # Publish predictions
        self.publish(agents_pos, pred_velocities)
        
    
    def listener(self):

        # Set up occupancy grid subscriber
        rospy.Subscriber('/roboat_cloud/obstacle/map', OccupancyGrid, self.store_occupancy_grid)

        # Set up odometry subscribers
        roboat_subscribers = []
        for i in range(self.n_agents):
            roboat_subscribers.append( message_filters.Subscriber("/roboat_" + str(i) + "/odometry_enu", Odometry) )
        agent_ids = range(self.n_agents)
        # Set up message filters and callback
        ts = message_filters.ApproximateTimeSynchronizer(roboat_subscribers, 1, 0.1)
        ts.registerCallback(self.callback)

        # Let it spin
        rospy.spin()


    def get_model_args(self, model_name, train_run):
        trained_dir = os.path.join(rospkg.RosPack().get_path(PACKAGE_NAME), 'trained_models')
        model_dir = os.path.join(trained_dir, model_name, train_run)
        convnet_dir = os.path.join(trained_dir, 'autoencoder_with_ped')
        with open(os.path.join(model_dir, 'model_parameters.pkl'), 'rb') as f:
            model_args = pickle.load(f)["args"]
        model_args.model_path = model_dir
        model_args.pretrained_convnet_path = convnet_dir
        model_args.batch_size = MAX_QUERY_AGENTS
        model_args.truncated_backprop_length = 1
        model_args.keep_prob = 1.0
        return model_args


    def load_model(self, model_class, model_args):
        model = model_class(model_args)
        tf_session = tf.Session()
        model.warmstart_model(model_args, tf_session)
        try: model.warmstart_convnet(model_args, tf_session)
        except: rospy.logwarn('Could not warm-start ConvNet')
        rospy.loginfo('Model loaded!')
        return model, tf_session

    def shutdown_callback(self):
      rospy.loginfo('{} was terminated.'.format(self.visual_node_name))


if __name__ == '__main__':

    # Get CLI arguments
    args = rospy.myargv(sys.argv)

    # Start main logic
    node = SVRNN_Predictor('socialvrnn_node', 'Social-VRNN roboat node')

    print("Number of agents: ", node.n_agents)

    # Infer predictions every 'dt' seconds
    rospy.loginfo('Inferring predictions every {} seconds'.format(node.model_args.dt))

    node.listener()



build/lib/socialvrnn_node.py
eliatrevisan/IA-MPPI-LBM
import sys
import os

sys.path.append('../')

# With debugger:
#os.chdir(os.getcwd() + "/src/")

import numpy as np
import argparse
import pylab as pl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pickle as pkl
import importlib
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import random
import progressbar
import scipy.io as sio
import json
from copy import deepcopy
import colorama
from colorama import Fore, Style

if sys.version_info[0] < 3:
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from Performance import *
	from utils import *
	import Recorder as rec
else:
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.Performance import *
	from src.data_utils.utils import *
	from src.data_utils.Recorder import Recorder as rec
	from src.models.CV_model import ConstantVelocity

"""
exp_num = 33
model_name = "VGDNN" 
num_test_sequences = 100 
scenario = "simulation/roboat" 
record = False 
n_samples = 1 
unit_testing = False 
freeze_other_agents = False
"""

from random import seed
seed(1)

# Model directories
def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
											help='Path to directory that comprises the model (default="model_name").',
											type=str, default="VGDNN")
	parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=33)
	parser.add_argument('--n_samples', help='Number of samples', type=int, default=10)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
											type=str, default="datasets/ewap_dataset/seq_eth")
	parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool,
											default=True)
	parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool,
											default=True)
	parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float,
											default=0)
	parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float,
											default=5)
	parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool,
											default=False)
	parser.add_argument('--update_state', help='update_state', type=sup.str2bool,
											default=False)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
											default=False)
	parser.add_argument('--freeze_other_agents', help='FReeze other agents.', type=sup.str2bool,
											default=False)
	parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,
											default=False)
	parser.add_argument('--constant_velocity', help='Run CV comparison', type=sup.str2bool,
											default=False)
	args = parser.parse_args()

	return args


test_args = parse_args()

if test_args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

cwd = os.getcwd()

model_path = os.path.normpath(cwd + '/../') + '/trained_models/' + test_args.model_name + "/" + str(test_args.exp_num)

print("Loading data from: '{}'".format(model_path))
file = open(model_path + '/model_parameters.pkl', 'rb')
if sys.version_info[0] < 3:
	model_parameters = pkl.load(file)  # ,encoding='latin1')
else:
	model_parameters = pkl.load(file, encoding='latin1')
file.close()
args = model_parameters["args"]

with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

print("Training steps: ", args.total_training_steps)

# change some args because we are doing inference
truncated_backprop_length = args.truncated_backprop_length
args.truncated_backprop_length = 1
args.batch_size = 1
args.keep_prob = 1.0


training_scenario = args.scenario
args.scenario = test_args.scenario
args.real_world_data = test_args.real_world_data
args.dataset = '/' + args.scenario + '.pkl'
data_prep = dhlstm.DataHandlerLSTM(args)

# Only used to create a map from png
# Load Map Parameters
map_params = os.path.join(args.data_path+args.scenario, 'map.json')
with open(map_params) as json_file:
	data = json.load(json_file)
map_args = {"file_name": data["file_name"],
						"resolution": data["resolution"],
						"map_size": np.array(data["map_size"]),
						"map_center": np.array(data["map_center"])}

data_prep.processData(**map_args)
if args.normalize_data:
	data_prep.compute_min_max_values()

# Import model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

model = NetworkModel(args)

# Lists for logging of the input / output data of the model
input_list = []
grid_list = []
goal_list = []
ped_grid_list = []
y_ground_truth_list = []
y_pred_list = []  # uses ground truth as input at every step
other_agents_list = []
all_predictions = []
all_traj_likelihood = []
trajectories = []
batch_y = []
batch_loss = []

config = tf.ConfigProto(
	device_count={'GPU': 0}
)


def clear_plot():
	axarr[0].clear()
	axarr[0].set_title('Training Sample')

if test_args.unit_testing:
	data_handler = dhlstm.DataHandlerLSTM(args)
	data_handler.unit_test_data_(map_args)

with tf.Session(config=config) as sess:
	model.warmstart_model(args, sess)
	try:
			model.warmstart_convnet(args, sess)
	except:
		print("")

	for exp_id in range(np.minimum(test_args.num_test_sequences,len(data_prep.trajectory_set)-1)):
		predictions = []
		traj_likelihood = []

		# sample a trajectory id for testing
		traj_id = random.randint(0, len(data_prep.trajectory_set) - 1)
		

		batch_x, batch_vel, batch_pos,batch_goal, batch_grid, other_agents_info, batch_target,batch_end_pos, other_agents_pos, traj = data_prep.getTrajectoryAsBatch(
			traj_id,freeze = test_args.freeze_other_agents)  
		
		trajectories.append(traj)
		x_input_series = np.zeros([0, (args.prev_horizon + 1) * args.input_dim])
		goal_input_series = np.zeros([0, 2])
		grid_input_series = np.zeros(
			[0, int(args.submap_size), int(args.submap_size)])
		if args.others_info == "relative":
			ped_grid_series = np.zeros([0, args.n_other_agents,args.pedestrian_vector_dim])
		elif "sequence" in args.others_info:
			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim*args.prediction_horizon])
		elif args.others_info == "prev_sequence":
			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim *(args.prev_horizon+1)])
		elif args.others_info == "sequence2":
			ped_grid_series = np.zeros([0, args.n_other_agents, args.prediction_horizon,args.pedestrian_vector_dim])
		else:
			ped_grid_series = np.zeros([0, args.pedestrian_vector_dim])
		y_ground_truth_series = np.zeros([0, args.prediction_horizon * 2])
		y_pred_series = np.zeros([0, args.n_mixtures * args.prediction_horizon * args.output_pred_state_dim])

		batch_y.append(batch_target)
		model.reset_test_cells(np.ones((args.batch_size)))
		cell_state_list= []
		cell_ped_list = []
		cell_concat_list = []
		if "grid" in args.model_name:
			batch_ped_grid_backup = np.zeros_like(batch_grid)
			data_prep.add_other_agents_to_grid(batch_ped_grid_backup, batch_x, [other_agents_pos])

		for step in range(batch_x.shape[1]):
			samples = []
			# Assemble feed dict for training
			if "future" in args.others_info:
				if step == 0:
					batch_y_pred = deepcopy(batch_vel)
					for pred_step in range(1, args.prediction_horizon):
						batch_y_pred[:, :, pred_step * 2:pred_step * 2 + 2] = batch_vel[:, :, :2]
				else:
					batch_y_pred = deepcopy(batch_vel)
					batch_y_pred[:, :, 2:] = y_model_pred[:, :, 2:]
			"""
			dict = {"batch_x": batch_x,
							"batch_vel": batch_vel,
							"batch_pos": batch_pos,
							"batch_grid": batch_grid,
							"batch_ped_grid": other_agents_info,
							"step": step,
							"batch_goal": batch_goal,
							"state_noise": 0.0,
							"grid_noise": 0.0,
							"other_agents_pos": [other_agents_pos]
			}
		
			feed_dict_ = model.feed_test_dic(**dict)
			"""

			dict = {
							"batch_vel": batch_vel,
							"batch_grid": batch_grid,
							"batch_ped_grid": other_agents_info,
							"step": step,
							"other_agents_pos": [other_agents_pos]
			}

			feed_dict_ = model.feed_test_dic(
                **dict
            )

			# Append to logging series
			x_input_series = np.append(x_input_series, batch_x[:, step, :], axis=0)
			grid_input_series = np.append(grid_input_series, batch_grid[:, step, :, :], axis=0)
			goal_input_series = np.append(goal_input_series, batch_goal[:, step, :], axis=0)

			y_ground_truth_series = np.append(y_ground_truth_series, batch_target[:, step, :], axis=0)

			y_model_pred, likelihood = model.predict(sess, feed_dict_, True)

			# Backup cell states for later analysis
			cell_state_list.append(model.test_cell_state_current[0,:])
			#cell_ped_list.append(model.test_cell_state_current_lstm_ped[0, :])
			cell_concat_list.append(model.test_cell_state_current_lstm_concat[0, :])

			# Rotate predictions to global frame
			if args.rotated_grid:
				heading = math.atan2(batch_vel[0, step, 1], batch_vel[0, step, 0])
				rot_mat = np.array([[np.cos(heading), -np.sin(heading)], [np.sin(heading), np.cos(heading)]])
				for pred_step in range(args.prediction_horizon):
					y_model_pred[0][0, 2 * pred_step:2 * pred_step + 2] = np.dot(rot_mat, y_model_pred[0][0,2 * pred_step:2 * pred_step + 2])
				samples.append(y_model_pred[0])
			else:
				samples.append(y_model_pred[:,0,:])

			# If sample more than one trajectory from the model
			for sample_id in range(test_args.n_samples - 1):
				dict = {"batch_x": batch_x,
								"batch_vel": batch_vel,
								"batch_pos": batch_pos,
								"batch_grid": batch_grid,
								"batch_ped_grid": other_agents_info,
								"step": step,
								"batch_goal": batch_goal,
								"state_noise": 0.0,
								"grid_noise": 0.0
								}
				feed_dict_ = model.feed_test_dic(**dict)
				y_model_pred = model.predict(sess, feed_dict_, test_args.update_state)
				samples.append(y_model_pred[:,0,:])

			traj_likelihood.append(likelihood)
			predictions.append(samples)

		latent_analysis = False

		# Plot hidden states evolution over a trajectory
		if latent_analysis:

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			pl.show(block=False)
			x_state = np.arange(0, args.rnn_state_size, 1)
			x_ped = np.arange(0, args.rnn_state_size_lstm_ped * 2, 1)
			x_concat = np.arange(0, args.rnn_state_size_lstm_concat, 1)

			time = np.arange(0, len(cell_state_list), 1)
			R= np.array(cell_state_list)
			print(R)
			ax.clear()
			X, Y = np.meshgrid(x_state,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			save_img_to_file = model_path + "/figs"
			if not os.path.exists(save_img_to_file):
				os.makedirs(save_img_to_file)
			pl.savefig(save_img_to_file+"/cell_state" + str(exp_id) + '.png')
			pl.show()
			pl.close()

			# fig = pl.figure()
			# ax = pl.axes(projection="3d")
			# time = np.arange(0, len(cell_ped_list), 1)
			# R= np.array(cell_ped_list)
			# print(R)
			# ax.clear()
			# X, Y = np.meshgrid(x_ped,time)
			# surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			# ax.view_init(elev=90,azim=90)
			# ax.set_xlabel("Time [s]")
			# ax.set_xlabel("Features")
			# fig.colorbar(surf)
			# pl.savefig(save_img_to_file+"/cell_ped" + str(exp_id) + '.png')
			# pl.close()

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			time = np.arange(0, len(cell_concat_list), 1)
			R= np.array(cell_concat_list)
			ax.clear()
			X, Y = np.meshgrid(x_concat,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			pl.savefig(save_img_to_file+"/cell_concat" + str(exp_id) + '.png')
			pl.show()
			pl.close()

		all_predictions.append(predictions)
		all_traj_likelihood.append(traj_likelihood)
		input_list.append(x_input_series)
		goal_list.append(goal_input_series)
		grid_list.append(grid_input_series)
		y_ground_truth_list.append(y_ground_truth_series)
		other_agents_list.append(other_agents_pos)
		# update progress bar

sess.close()

# Save data in Matlab format for post-processing
if not os.path.exists(args.model_path + '/../results/'):
	os.makedirs(args.model_path + '/../results/')
results= {
	"trajectories": trajectories,
	"predictions": all_predictions
}
scenario = args.scenario.split('/')[-1]
results_file = args.model_path + '/'+ scenario + "_results.mat"
sio.savemat(results_file,results)

if test_args.constant_velocity:
	# Make constant velocity predictions
	CVModel = ConstantVelocity(args)
	cv_predictions = CVModel.predict(input_list)
else:
	cv_predictions = None


if test_args.record:
		recorder = rec(args, data_prep.agent_container.occupancy_grid)
		if ( "real_world" in test_args.scenario) and not test_args.unit_testing:
			print("Real data!!")
			recorder.plot_on_video(input_list, grid_list, all_predictions, y_ground_truth_list,
														 other_agents_list,
														 trajectories,all_traj_likelihood, test_args)
		else:
			#recorder.plot_on_image(input_list, grid_list, all_predictions, y_ground_truth_list, other_agents_list,
			#	                       trajectories,test_args)
			recorder.animate_local(input_list, grid_list, ped_grid_list, all_predictions, y_ground_truth_list, other_agents_list,
										 trajectories,test_args)

			recorder.animate_global(input_list, grid_list, all_predictions, y_ground_truth_list,
														 other_agents_list,
														 trajectories, all_traj_likelihood, cv_predictions, test_args)
			#cv_fde_error = compute_fde(args, trajectories, cv_predictions)
			#cv_ade_error = compute_ade(args, trajectories, cv_predictions)

			#print(" ADE: ", cv_ade_error, " FDE: ", cv_fde_error)

			# cv_fde = compute_rolling_fde(args, trajectories, cv_predictions)
	
			# fde_roll = []
			# for horizon in range(0,args.prediction_horizon):
			# 	pred_fde, pred_error_summary_lstm_fde = compute_rolling_trajectory_fde(args, horizon, trajectories, all_predictions)
			# 	fde_roll.append(pred_fde)
			# 	print("Step: ", horizon, " FDE: ", pred_fde)

			# print("CV: ", cv_fde)
			# print("SVRNN: ", fde_roll)

			print("Recorder is done!")
else:
	print("Performance tests")
	

	cv_fde, cv_std = compute_rolling_fde(args, trajectories, cv_predictions)
	
	fde_roll = []
	std_roll = []
	for horizon in range(0,args.prediction_horizon):
		pred_fde, pred_error_summary_lstm_fde, std = compute_rolling_trajectory_fde(args, horizon, trajectories, all_predictions)
		fde_roll.append(np.mean(pred_error_summary_lstm_fde))
		std_roll.append(std)

	print("CV: ", cv_fde)
	print("SVRNN: ", fde_roll)

	# 30% of STD, just like in bruno's paper, but why?
	std_roll = np.array(std_roll)*0.3
	cv_std = cv_std * 0.3

	fig, ax = plt.subplots(figsize=(4,3))
	x = [1,2,3,4,5,6,7,8,9,10,11,12]
	ax.set_title("Herengracht")
	ax.plot(x, cv_fde, label="CV", color="b")
	ax.plot(x, fde_roll, label="SVRNN", color="r")
	ax.fill_between(x, cv_fde+cv_std, cv_fde-cv_std, color="b", alpha=0.2)
	ax.fill_between(x, fde_roll+std_roll, fde_roll-std_roll, color="r", alpha=0.2)
	ax.set_xlabel("Timestep")
	ax.set_ylabel("Displacement Error [m]")
	plt.tight_layout()
	plt.legend()
	plt.show()
	
	pred_error, pred_error_summary_lstm = compute_trajectory_prediction_mse(args, trajectories, all_predictions)
	pred_fde, pred_error_summary_lstm_fde = compute_trajectory_fde(args, trajectories, all_predictions)
	if test_args.constant_velocity: 
		#cv_pred_error = compute_ade_cv(args, trajectories, cv_predictions)
		#cv_fde_error = compute_fde_cv(args, trajectories, cv_predictions)
		cv_fde_error = compute_fde(args, trajectories, cv_predictions)
		cv_pred_error = compute_ade(args, trajectories, cv_predictions)
		print(
		Fore.LIGHTBLUE_EX + "\nConstant Velocity: MSE: {:01.2f}, FDE: {:01.2f}".format(np.mean(cv_pred_error), np.mean(cv_fde_error))+Style.RESET_ALL)
	diversity, diversity_summary = compute_2_wasserstein(args, all_predictions)
	args.scenario = training_scenario
	args.truncated_backprop_length = truncated_backprop_length
	write_results_summary(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde), np.mean(diversity_summary), args, test_args, cv_pred_error, cv_fde_error)
	print(
		Fore.LIGHTBLUE_EX + "\nMSE: {:01.2f}, FDE: {:01.2f}, DIVERSITY: {:01.2f}".format(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde),np.mean(diversity_summary))+Style.RESET_ALL)
	

build/lib/test_VGDNN.py
eliatrevisan/IA-MPPI-LBM
import sys
import os

sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pickle as pkl
import importlib
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import random
import progressbar
import scipy.io as sio
import json
from copy import deepcopy
if sys.version_info[0] < 3:
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from Performance import *
	from utils import *
	import Recorder as rec
else:
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.Performance import *
	from src.data_utils.utils import *
	from src.data_utils.Recorder import Recorder as rec


# Model directories
def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default="VGDNN_simple")
	parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=9)
	parser.add_argument('--n_samples', help='Number of samples', type=int, default=10)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
	                    type=str, default="datasets/ewap_dataset/seq_eth")
	parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool,
	                    default=True)
	parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool,
	                    default=True)
	parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float,
	                    default=0)
	parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float,
	                    default=5)
	parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--update_state', help='update_state', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--freeze_other_agents', help='FReeze other agents.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,
	                    default=False)
	args = parser.parse_args()

	return args


test_args = parse_args()

if test_args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

cwd = os.getcwd()

model_path = os.path.normpath(cwd + '/../') + '/trained_models/' + test_args.model_name + "/" + str(test_args.exp_num)

print("Loading data from: '{}'".format(model_path))
file = open(model_path + '/model_parameters.pkl', 'rb')
if sys.version_info[0] < 3:
	model_parameters = pkl.load(file)  # ,encoding='latin1')
else:
	model_parameters = pkl.load(file, encoding='latin1')
file.close()
args = model_parameters["args"]

with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

# change some args because we are doing inference
args.truncated_backprop_length = 1
args.batch_size = 1
args.keep_prob = 1.0
# args.data_path = "../data/2_agents_swap/trajs/"
# args.scenario = "GA3C-CADRL-10-py27"
# args.dataset = args.scenario+'.pkl'
# args.data_path = '../data/'

training_scenario = args.scenario
args.scenario = test_args.scenario
args.real_world_data = test_args.real_world_data
args.dataset = '/' + args.scenario + '.pkl'
data_prep = dhlstm.DataHandlerLSTM(args)
# Only used to create a map from png
map_args = {"file_name": 'map.png',
	            "resolution": 0.1,
	            "map_size": np.array([30., 30.]), }

data_prep.processData(**map_args)
if args.normalize_data:
	data_prep.compute_min_max_values()

# Import model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

model = NetworkModel(args)

# Lists for logging of the input / output data of the model
input_list = []
grid_list = []
goal_list = []
ped_grid_list = []
y_ground_truth_list = []
y_pred_list = []  # uses ground truth as input at every step
other_agents_list = []
all_predictions = []
trajectories = []
batch_y = []
batch_loss = []

config = tf.ConfigProto(
	device_count={'GPU': 0}
)


def clear_plot():
	axarr[0].clear()
	axarr[0].set_title('Training Sample')

if test_args.unit_testing:
	data_handler = dhlstm.DataHandlerLSTM(args)
	data_handler.unit_test_data_(map_args)

with tf.Session(config=config) as sess:
	model.warmstart_model(args, sess)
	try:
			model.warmstart_convnet(args, sess)
	except:
		print("")

	for exp_idx in range(0,np.minimum(test_args.num_test_sequences,len(data_prep.trajectory_set)-1),args.batch_size):
		predictions = []

		batch_x_list, batch_vel_list, batch_pos_list, batch_goal_list, batch_grid_list, \
		pedestrian_grid_list, batch_y_list, batch_pos_final_list, other_agents_pos_list, \
		traj_list = data_prep.getGroupOfTrajectoriesAsBatch(exp_idx,n_trajs=6)  # trajectory_set random.randint(0, len(data_prep.dataset) - 1)

		trajectories.append(traj_list)

		x_input_series = []
		goal_input_series = []
		grid_input_series = []
		y_ground_truth_series = []

		for traj_id in range(len(traj_list)):
			batch_x = batch_x_list[traj_id]
			batch_vel = batch_vel_list[traj_id]
			batch_pos = batch_pos_list[traj_id]
			batch_goal = batch_goal_list[traj_id]
			batch_grid = batch_grid_list[traj_id]
			other_agents_info = pedestrian_grid_list[traj_id]
			batch_pos_final = batch_pos_final_list[traj_id]
			other_agents_pos = other_agents_pos_list[traj_id]
			batch_grid = batch_grid_list[traj_id]
			traj = traj_list[traj_id]

			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim * args.prediction_horizon])
			model.reset_test_cells(np.ones((args.batch_size)))

			samples = []
			for step in range(batch_x.shape[1]):
				# Assemble feed dict for training
				if "future" in args.others_info:
					if step == 0:
						batch_y_pred = deepcopy(batch_vel)
						for pred_step in range(1, args.prediction_horizon):
							batch_y_pred[:, :, pred_step * 2:pred_step * 2 + 2] = batch_vel[:, :, :2]
					else:
						batch_y_pred = deepcopy(batch_vel)
						batch_y_pred[:, :, 2:] = y_model_pred[:, :, 2:]
				dict = {"batch_x": batch_vel,
				        "batch_grid": batch_grid,
				        "batch_ped_grid": other_agents_info,
				        "step": step,
				}
				feed_dict_ = model.feed_test_dic(**dict)

				y_model_pred, output_decoder, cell_outputs_series_state, cell_outputs_series_lstm_ped, cell_concat = model.predict(sess, feed_dict_, True)

				samples.append(y_model_pred[:,0,:])

			predictions.append(samples)

		all_predictions.append(predictions)
		# update progress bar

sess.close()

# Save data
if not os.path.exists(args.model_path + '/../results/'):
	os.makedirs(args.model_path + '/../results/')
results= {
	"trajectories": trajectories,
	"predictions": all_predictions
}
scenario = args.scenario.split('/')[-1]
results_file = args.model_path + '/'+ scenario + "_results.mat"

sio.savemat(results_file,results)

if test_args.record:
		recorder = rec(args, data_prep.agent_container.occupancy_grid)
		if ( "real_world" in test_args.scenario) and not test_args.unit_testing:
			print("Real data!!")
			recorder.plot_on_video(input_list, grid_list, all_predictions, y_ground_truth_list,
			                       other_agents_list,
			                       trajectories, test_args)
		else:
			#recorder.plot_on_image(input_list, grid_list, all_predictions, y_ground_truth_list, other_agents_list,
			#	                       trajectories,test_args)
			#recorder.animate_local(input_list, grid_list, ped_grid_list, all_predictions, y_ground_truth_list, other_agents_list,
		  #               trajectories,test_args)
			recorder.animate_group_of_agents(trajectories, all_predictions,test_args)

		print("Recorder is done!")
else:
	"""Performance tests"""
	print("Average loss: " + str(np.mean(np.asarray(batch_loss))))
	pred_error, pred_error_summary_lstm = compute_trajectory_prediction_mse(args, trajectories, all_predictions)
	pred_fde, pred_error_summary_lstm_fde = compute_trajectory_fde(args, trajectories, all_predictions)
	args.scenario = training_scenario
	write_results_summary(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde), 0, args, test_args)

build/lib/test_VGDNN2.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import tensorflow as tf
import pickle as pkl
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
if sys.version_info[0] < 3:
  sys.path.append('../src/data_utils')
  sys.path.append('../src/models')
  import DataHandler as dh
  from seqModel import *
  from plot_utils import *
else:
  from src.data_utils import DataHandler as dh
  from src.models.seqBiModel import *
  from src.data_utils.plot_utils import *

tfk = tf.keras

model_name = "seq_to_seq"
data_path = '../data/2_agents/trajs'
input_dim = 8
latent_dim = 256
output_dim = 2
cmd_vector_dim = 2
batch_size = 16
learning_rate_init = 0.01
grads_clip = 10.0
n_epochs = 100
stateful = False
prediction_horizon = 10
prev_horizon = 10
truncated_backprop_length = 10
regularization_weight = 0.0001
dropout = 0.2

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default=model_name)
	parser.add_argument('--data_path', help='Data directory', type=str,
	                    default=data_path)
	parser.add_argument('--save_path', help='Path to directory that saves pickle data (default=" ").', type=str,
	                    default=data_path+'/GA3C-CADRL-10-py27.pkl')
	parser.add_argument('--batch_size', help='Batch size.', type=int,
	                    default=batch_size)
	parser.add_argument('--learning_rate_init', help='Learning rate.', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--grads_clip', help='Grad clip.', type=float,
	                    default=grads_clip)
	parser.add_argument('--n_epochs', help='Number of epochs.', type=float,
	                    default=n_epochs)
	parser.add_argument('--stateful', help='Stateful.', type=bool,
	                    default=stateful)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--prediction_horizon', help='Prediction Horizon.', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--dropout', help='Dropout probability (default=0.005).', type=float,
	                    default=dropout)
	parser.add_argument('--output_dim', help='Output dimension.', type=int,
	                    default=output_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
	                    default=cmd_vector_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=int,
	                    default=input_dim)
	parser.add_argument('--latent_dim', help='Latent dimension.', type=int,
	                    default=latent_dim)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--distance_threshold', help='Distance threshold to start new trajectory', type=float, default=1)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
	                    default='../config/topics.yaml')
	parser.add_argument('--grid_width', help='width of occupancy grid', type=int, default=120)
	parser.add_argument('--grid_height', help='height of occupancy grid', type=int, default=120)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--file_name', help='PNG map file name', type=str,
	                    default='simple_multipath_rotated.png')
	parser.add_argument('--resolution', help='Distance threshold to start new trajectory', type=float,
	                    default=0.1)
	parser.add_argument('--map_size_x', help='Distance threshold to start new trajectory', type=float,
	                    default=20.0)
	parser.add_argument('--map_size_y', help='Distance threshold to start new trajectory', type=float,
	                    default=7.0)

	args = parser.parse_args()

	return args


args = parse_args()

print("Loading data from: '{}'".format(args.save_path))
file = open(args.save_path, 'rb')
if sys.version_info[0] < 3:
	tmp_self = pkl.load(file)
else:
	tmp_self = pkl.load(file, encoding='latin1')
	dataset = tmp_self

data_prep = dh.DataHandler(args)
# Collect data online
if os.path.isfile(args.save_path):
	data_prep.load_data()
else:
	data_prep.start_node()

batch_enc, batch_dec, batch_dec_target, new_epoch = data_prep.getBatch()

model = SeqToSeqModel2(args)

model.model.compile(optimizer=tfk.optimizers.Adam(lr=args.learning_rate_init,clipnorm=args.grads_clip),
                             loss='mean_squared_error')


history = model.model.fit([batch_enc,batch_dec],batch_dec_target,
                epochs=n_epochs,
                batch_size=16,
								#validation_split=0.2,
                shuffle=True)

# Plot training & validation loss values
plt.plot(history.history['loss'])
plt.title('Modtrain.pyel loss during training')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train'], loc='upper left')
plt.show()

batch_enc_seq, batch_dec_seq, traj = data_prep.getTrajectoryAsBatch()

predictions = model.decode_sequence([batch_enc_seq,batch_dec_seq])

#plot_scenario(batch_enc_seq,batch_dec_seq,traj,predictions,data_prep.occupancy_grid,args)
plot_scenario_vel(batch_enc_seq,batch_dec_seq,traj,predictions,args)
print("test")


build/lib/train.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
import json
import importlib
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
from datetime import datetime
import pickle as pkl
import time
from copy import deepcopy
from multiprocessing.pool import ThreadPool
import colorama
from colorama import Fore, Style
import wandb

#os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

pretrained_convnet_path = "../trained_models/autoencoder_with_ped"

#Defaults
exp_num = 6
data_path = '../data/roboat/'
scenario = 'simulation/'

# Hyperparameters
n_epochs = 0

batch_size = 16 #128
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 8
prediction_horizon = 8
prev_horizon = 0

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 128
rnn_state_ped_size = 16
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 256
x_dim = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 18
tensorboard_logging = True
wandb_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 5

print_freq = 200
save_freq = 500
total_training_steps = 50000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = False
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
#submap_resolution = 0.081
#submap_width = 4.86
#submap_height = 4.86
submap_size = 60 # In pixels
submap_span_real = 10 # Coordinates
relative_covariance = 10.0
diversity_update = False
predict_positions = False
warm_start_convnet = True
debug_plotting = False

# Dataset division
train_set = 0.8

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default= "RNN")
	parser.add_argument('--model_path',
	                    help='Path to directory to save the model (default=""../trained_models/"+model_name").',
	                    type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
	                    help='Path to directory that comprises the pre-trained convnet model (default=" ").',
	                    type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
	                    help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
	                    type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
	                    type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
	                    default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
	                    default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
	                    default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
	                    default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
	                    default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
	                    default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
	                    default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--total_training_steps', help='Number of training steps (default=20000).', type=int, default=total_training_steps)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
	                    default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
	                    default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
	                    default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
	                    default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
	                    default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
	                    help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
	                    type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
	                    type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
	                    type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
	                    help='Number of fully connected layer units after LSTM layer (default=64).',
	                    type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
	                    default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
	                    default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
	                    default=output_dim)
	parser.add_argument('--goal_size', help='Goal dimension (default=).', type=int,
	                    default=2)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
	                    default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
	                    default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
	                    default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
	                    default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
	                    default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
	                    default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
	                    default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the network.', type=int,
	                    default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=sup.str2bool, default=debug_plotting)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
	                    default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
	                    type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
	                    default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
	                    default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
	                    default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
	                    default=0)
	#parser.add_argument('--submap_width', help='width of occupancy grid (For initialization of arrays)', type=int, default=submap_width)
	#parser.add_argument('--submap_height', help='height of occupancy grid (for initiatization of arrays)', type=int, default=submap_height)
	parser.add_argument('--submap_size', help='size of submap in pixels (60)', type=int, default=submap_size)
	parser.add_argument('--submap_span_real', help='span of submap occupancy grid', type=int, default=submap_span_real)
	#parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	#                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
	                    default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
	                    default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
	                    default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
	                    default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
	                    default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	parser.add_argument('--kl_weight', help='Experimenting with additional weight for KL loss', type=float, default=1)
	parser.add_argument('--relative_covariance', help='Changing covariance of relative positions', type=float, default=relative_covariance)
	args = parser.parse_args()

	return args

args = parse_args()


# Cyclical KL Annealing weight
def frange_cycle_linear(n_iter, start=0.0, stop=1.0,  n_cycle=4, ratio=0.5):
	L = np.ones(n_iter) * stop
	period = n_iter/n_cycle
	step = (stop-start)/(period*ratio) # linear schedule

	for c in range(n_cycle):
		v, i = start, 0
		while v <= stop and (int(i+c*period) < n_iter):
			L[int(i+c*period)] = v
			v += step
			i += 1
	return L

# Enable / Disable GPU
if args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf
"""
# Initialize wandb logging
wandb.init(project="SocialVRNN", entity="walterj", config=tf.flags.FLAGS, sync_tensorboard=True)
wandb.config = {
  "learning_rate": learning_rate_init,
  "epochs": n_epochs,
  "experiment number": exp_num,
  "batch_size": batch_size,
  "training steps": total_training_steps,
  "truncated backpropagation steps": truncated_backprop_length,
  "previous horizon": prev_horizon,
  "prediction horizon": prediction_horizon,
  "gradient clipping": grads_clip,
  "dt": dt,
  "regularization weight": regularization_weight,
  "resolution": submap_resolution,
  "submap_width": submap_width,
  "submap_height": submap_height,
  "submap_width_real": submap_span_real,
  "submap_height_real": submap_span_real
}
"""

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

# Create Log and Model Directory to save training model
args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)
args.log_dir = args.model_path + '/log'
if not os.path.exists(args.log_dir):
	os.makedirs(args.log_dir)
args.dataset = '/' + args.scenario + '.pkl'
model_parameters = {"args": args}

# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)

# Save Model Parameters
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()
with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)



# Create Datahandler class for the specified scenario
if args.scenario == "herengracht":
	data_prep = dhlstm.DataHandlerLSTM(args)
	map_args = {"file_name": 'map.png',
	            "resolution": 0.081,
	            "map_size": np.array([1715,881]),
				"map_center": np.array([-78,-40])} # THIS IS HARDCODED FOR SOME REASON AND NOT EXTRACTED FROM MAP.JSON ?
	data_prep.processData(**map_args)

elif args.scenario == "bloemgracht":
	data_prep = dhlstm.DataHandlerLSTM(args)
	map_args = {"file_name": 'map.png',
				"resolution": 0.22,
				"map_size": np.array([1469,928]),
				"map_center": np.array([-161.59,-102.08])} 
	data_prep.processData(**map_args)

elif args.scenario == "open_crossing":
	data_prep = dhlstm.DataHandlerLSTM(args)
	map_args = {"file_name": 'map.png',
				"resolution": 0.6,
				"map_size": np.array([100,100]),
				"map_center": np.array([-30,-30])} 
	data_prep.processData(**map_args)

elif args.scenario == "prinsengracht":
	data_prep = dhlstm.DataHandlerLSTM(args)
	map_args = {"file_name": 'map.png',
				"resolution": 0.091,
				"map_size": np.array([816,816]),
				"map_center": np.array([-37.1, -37.1])} 
	data_prep.processData(**map_args)

# Import Deep Learning model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

# Create Model Graph
model = NetworkModel(args)

config = tf.ConfigProto()

beta_list = frange_cycle_linear(args.total_training_steps)

# Start Training Session
with tf.Session(config=config) as sess:

	# Load a pre-trained model
	if args.warmstart_model:
		model.warmstart_model(args, sess)
	else:
		# Initialize all TF variables
		sess.run(tf.global_variables_initializer())

	# Load Convnet Model
	try:
		if args.warm_start_convnet:
			model.warmstart_convnet(args, sess)
	except:
		print("Failed to initialized Convnet or Convnet does not exist")

	# if the training was interrupted load last training step index
	try:
		initial_step = int(open(args.model_path + "/tf_log", 'r').read().split('\n')[-2]) + 1
	except:
		initial_step = 1

	epoch = 0
	training_loss = []
	diversity_loss = []
	training_loss.append(0)

	# Set up multithreadign for data handler
	pool = ThreadPool(1)
	res = None
	_model_prediction = []
	start_time = time.time()
	best_loss = float('inf')
	avg_training_loss = np.ones(100)

	for step in range(initial_step,args.total_training_steps):
		start_time_loop = time.time()

		# Get Next Batch of Data
		if res == None:
			batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch = data_prep.getBatch()
		else:
			batch = res.get(timeout=5)

		# Create dictionary to feed into the model
		dict = {"batch_x": batch_x,
		        "batch_vel": batch_vel,
		        "batch_pos": batch_pos,
		        "batch_goal": batch_goal,
			      "batch_grid": batch_grid,
			      "batch_ped_grid": batch_ped_grid,
			      "step": step,
				  "beta": beta_list[step],
			      "batch_y": batch_y,
			      "batch_pos_target": batch_pos_target,
		        "batch_div": batch_y,
		        "other_agents_pos": other_agents_pos
			      }

		feed_dict_train = model.feed_dic(**dict)

		#res = pool.apply_async(data_prep.getBatch)

		epoch += new_epoch

		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		model.reset_cells(data_prep.sequence_reset)

		start_time_training = time.time()

		model_output = model.train_step(sess, feed_dict_train,step)
		

		avg_training_time = time.time() - start_time_training
		avg_loop_time = time.time() - start_time_loop

		training_loss.append(model_output["batch_loss"])

		#if step == 1 :
			#avg_training_loss *= model_output["batch_loss"]
		#else:
			#avg_training_loss = np.roll(avg_training_loss,shift=1)
			#avg_training_loss[0] = model_output["batch_loss"]

		# Print training info
		if step % print_freq == 0:

			# Get batch to compute validation loss
			validation_dict = data_prep.getTestBatch()

			model.reset_test_cells(data_prep.val_sequence_reset)

			feed_dict_validation = model.feed_val_dic(**validation_dict)

			#validation_loss, validation_summary, validation_predictions = model.validation_step(sess, feed_dict_validation)
			validation_loss, validation_summary, validation_predictions = model.validation_step(sess, feed_dict_train)
			ellapsed_time = time.time() - start_time

			print(Fore.BLUE + "\n\nEpoch {:d}, Steps: {:d}, Train loss: {:01.2f}, Validation loss: {:01.2f}, Epoch time: {:01.2f} sec"
			      .format(epoch + 1, step, np.mean(avg_training_loss), validation_loss, ellapsed_time)+Style.RESET_ALL)
			#print(Fore.BLUE + "\n\nEpoch {:d}, Steps: {:d}, Epoch time: {:01.2f} sec"
			#      .format(epoch + 1, step, ellapsed_time)+Style.RESET_ALL)

			if tensorboard_logging:
				model.summary_writer.add_summary(model_output["summary"], step)
				model.summary_writer.flush()
				model.validation_summary_writer.add_summary(validation_summary, step)
				model.validation_summary_writer.flush()

			# Plot Global and Local Scenarios to validate datasets
			if args.debug_plotting:
				for seq_index in range(args.batch_size):
					for t in range(args.truncated_backprop_length):
						data_prep.plot_global_scenario(batch_grid, batch_x, batch_y, batch_goal, other_agents_pos,
						                     model_output["model_predictions"], t, seq_index)
						data_prep.plot_local_scenario(batch_grid, batch_x, batch_y, batch_goal, other_agents_pos,
						                    model_output["model_predictions"], t, seq_index)

			with open(args.model_path + "/tf_log", 'a') as f:
				f.write(str(step) + '\n')
			curr_loss = (validation_loss+np.mean(avg_training_loss))/2.0
			if curr_loss < best_loss:
				save_path = args.model_path + '/model_ckpt'
				model.full_saver.save(sess, save_path, global_step=step)
				best_loss = curr_loss
				print(Fore.LIGHTCYAN_EX+'Step {}: Saving model under {}'.format(step, save_path))

		#step = step + 1

	write_summary(training_loss[-1], args)
	full_path = args.model_path + '/final-model.ckpt'
	model.full_saver.save(sess, full_path)
	print('Saved final model under "{}"'.format(full_path))

	#if wandb_logging:
	#	wandb.tensorflow.log(tf.summary.merge_all())

	if tensorboard_logging:
		model.summary_writer.close()

sess.close()



build/lib/train_VGDNN.py
eliatrevisan/IA-MPPI-LBM
import sys
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl

if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
	from autoencoder_grid import ae_model as ae
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
	from src.autoencoder_grid import ae_model as ae

import tensorflow as tf
import pickle as pkl
import time
from copy import deepcopy

tfk = tf.keras

model_name = "autoencoder_with_ped2"
model_path = '../trained_models/' + model_name
pretrained_convnet_path = "../trained_models/autoencoder_with_ped"
log_dir = model_path + '/log'
data_path = '../data/2_agents_swap/trajs/'
scenario = 'GA3C-CADRL-10-py27'
data_path = '../data/cyberzoo_experiments/'
scenario = 'all_trajectories'
exp_num = 6
data_path = '../data/'
scenario = '20_ped_with_obstacles/short_few_obstacles'

# Hyperparameters
n_epochs = 1200

batch_size = 16
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 10
prediction_horizon = 10
prev_horizon = 0

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 128
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 0.8
dropout = False
reg = 1e-4
n_mixtures = 3  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
dropout = False
tensorboard_logging = False

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 5  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 10

print_freq = 2000
save_freq = 500
dt = 0.4

warmstart_model = False
pretrained_convnet = True
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = True
rotated_grid = False
centered_grid = True
noise = False
normalize_data = True
real_world_data = False

# Map parameters
submap_resolution = 0.1
submap_width = 6
submap_height = 6


def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default=model_name)
	parser.add_argument('--model_path',
	                    help='Path to directory to save the model (default=""../trained_models/"+model_name").',
	                    type=str, default=model_path)
	parser.add_argument('--pretrained_convnet_path',
	                    help='Path to directory that comprises the pre-trained convnet model (default=" ").',
	                    type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
	                    help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
	                    type=str, default=log_dir)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
	                    type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
	                    default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
	                    default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
	                    default='/' + scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
	                    default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
	                    default=warmstart_model)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
	                    default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
	                    default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
	                    default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
	                    default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
	                    default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
	                    default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_size_lstm_grid',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
	                    help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
	                    type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
	                    type=int, default=latent_space_size)
	parser.add_argument('--fc_hidden_unit_size',
	                    help='Number of fully connected layer units after LSTM layer (default=64).',
	                    type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=float,
	                    default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
	                    default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
	                    default=output_dim)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
	                    default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
	                    default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
	                    default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
	                    default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
	                    default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
	                    default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
	                    default=pedestrian_radius)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=int, default=0)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
	                    default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
	                    type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
	                    default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
	                    default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
	                    default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
	                    default=0)
	parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=submap_width)
	parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=submap_height)
	parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
	                    default='../config/topics.yaml')
	args = parser.parse_args()

	return args


args = parse_args()

args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)

model_parameters = {"args": args}
print(args)
# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()

if "GA3C" in args.scenario:
	data_prep = dh.DataHandler(args)
	# Collect data online
	if os.path.isfile(args.data_path + args.dataset):
		data_prep.load_data()
	else:
		data_prep.start_node()
else:
	data_prep = dhlstm.DataHandlerLSTM(args)
	# Only used to create a map from png
	map_args = {"file_name": 'map.png',
	            "resolution": 0.1,
	            "map_size": np.array([30., 30.]), }
	data_pickle = args.data_path + args.scenario + "/data" + str(prediction_horizon) + ".pickle"
	data_prep.processData(data_pickle, **map_args)

model = ae.ConvAEModel(learning_rate_init=learning_rate_init, grid_width=submap_width*10, grid_height=submap_height*10,
                         log_dir=log_dir, is_training=True, latent_space_dim=64)

config = tf.ConfigProto(
	device_count={'GPU': 0}
)

fig, axarr = pl.subplots(1, 2)
pl.show(block=False)


def clear_plot():
	axarr[0].clear()
	axarr[1].clear()
	axarr[0].set_title('Training Sample')
	axarr[1].set_title('Training Output')


with tf.Session(config=config) as sess:
	# Timing
	avg_time_comp = 0.0
	avg_time_total = 0.0

	if args.warmstart_model:
		model.warmstart_model(args, sess)
	else:
		# Initialize all TF variables
		sess.run(tf.global_variables_initializer())
	model.warmstart_convnet(args, sess)

	# if the trainning was interrupted
	try:
		step = int(open(args.model_path + "/tf_log", 'r').read().split('\n')[-2]) + 1
	except:
		step = 1
	epoch = 0
	training_loss = []
	diversity_loss = []
	while epoch <= n_epochs:
		start_time_loop = time.time()
		# if res == None:
		batch_x, batch_vel, batch_grid, batch_ped_grid, batch_y, other_agents_pos, new_epoch = data_prep.getBatch()

		if new_epoch:
			epoch += 1

		data_prep.add_other_agents_to_grid(batch_grid, batch_x, other_agents_pos, args)

		feed_dict_train = {model.input_placeholder: batch_grid[:, 0, :, :],
		                   model.output_placeholder: batch_grid[:, 0, :, :]}
		time_start = time.time()
		_, total_loss, deconv1, summary_str = sess.run([model.update, model.total_loss, model.deconv1, model.summary],
		                                               feed_dict=feed_dict_train)

		model.summary_writer.add_summary(summary_str,   step)
		time_avg += time.time() - time_start
		if step % print_freq == 0:
			clear_plot()
			sup.plotGrid(batch_grid[0, 0, :, :], axarr[0, 0], color='k', alpha=0.5)
			sup.plotGrid(deconv1[0, :, :], axarr[0, 1], color='k', alpha=0.5)
			fig.canvas.draw()
			time_avg = 0.0

		# Save model every now and then
		if step % save_freq == 0:
			save_path = model_path + '/model_ckpt'
			model.saver.save(sess, save_path, global_step=  step)
			print('Epoch {} Step {}: Saving model under {}'.format(epoch,   step, save_path))
			full_path = model_path + '/final-model.ckpt'
			model.saver.save(sess, full_path)
			print('Saved final model under "{}"'.format(full_path))

		step +=1

sess.close()

print("test")

build/lib/train_autoencoder.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
import json
import importlib
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
from datetime import datetime
import pickle as pkl
import time
from copy import deepcopy
from multiprocessing.pool import ThreadPool
import colorama
from colorama import Fore, Style

#os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

pretrained_convnet_path = "../trained_models/autoencoder_with_ped"
exp_num = 6
data_path = '../data/roboat/'
scenario = 'herengracht'

# Hyperparameters
n_epochs = 0

batch_size = 16 #128
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 8
prediction_horizon = 8
prev_horizon = 0

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 128
rnn_state_ped_size = 16
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 256
x_dim = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 18
tensorboard_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 5

print_freq = 200
save_freq = 500
total_training_steps = 50000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = False
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
#submap_resolution = 0.081
#submap_width = 4.86
#submap_height = 4.86
submap_size = 60 # In pixels
submap_span_real = 10 # In Map coordinates
relative_covariance = 10.0
diversity_update = False
predict_positions = False
warm_start_convnet = True
debug_plotting = False

# Dataset division
train_set = 0.8

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
						help='Path to directory that comprises the model (default="model_name").',
						type=str, default= "RNN")
	parser.add_argument('--model_path',
						help='Path to directory to save the model (default=""../trained_models/"+model_name").',
						type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
						help='Path to directory that comprises the pre-trained convnet model (default=" ").',
						type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
						help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
						type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
						type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
						default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
						default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
						default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
						default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
						default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
						default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
						default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--total_training_steps', help='Number of training steps (default=20000).', type=int, default=total_training_steps)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
						default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
						default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
						default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
						default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
						default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
						default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
						type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
						default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
						default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
						default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
						help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
						type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
						type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
						type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
						help='Number of fully connected layer units after LSTM layer (default=64).',
						type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
						default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
						default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
						default=output_dim)
	parser.add_argument('--goal_size', help='Goal dimension (default=).', type=int,
						default=2)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
						default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
						default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
						default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
						default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
						default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
						default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
						default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the network.', type=int,
						default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=sup.str2bool, default=debug_plotting)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
						default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
						type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
						default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
						default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
						default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
						default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
						default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
						default=0)
	#parser.add_argument('--submap_width', help='width of occupancy grid (For initialization of arrays)', type=int, default=submap_width)
	#parser.add_argument('--submap_height', help='height of occupancy grid (for initiatization of arrays)', type=int, default=submap_height)
	parser.add_argument('--submap_size', help='size of submap in pixels (60)', type=int, default=submap_size)
	parser.add_argument('--submap_span_real', help='span of submap occupancy grid', type=int, default=submap_span_real)
	#parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	#                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
						default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
						default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
						default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
						default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
						default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
						default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
						default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
						default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	parser.add_argument('--kl_weight', help='Experimenting with additional weight for KL loss', type=float, default=1)
	parser.add_argument('--relative_covariance', help='Changing covariance of relative positions', type=float, default=relative_covariance)
	args = parser.parse_args()

	return args

args = parse_args()


# Cyclical KL Annealing weight
def frange_cycle_linear(n_iter, start=0.0, stop=1.0,  n_cycle=4, ratio=0.5):
	L = np.ones(n_iter) * stop
	period = n_iter/n_cycle
	step = (stop-start)/(period*ratio) # linear schedule

	for c in range(n_cycle):
		v, i = start, 0
		while v <= stop and (int(i+c*period) < n_iter):
			L[int(i+c*period)] = v
			v += step
			i += 1
	return L

# Enable / Disable GPU
if args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

# Create Log and Model Directory to save training model
args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)
args.log_dir = args.model_path + '/log'
if not os.path.exists(args.log_dir):
	os.makedirs(args.log_dir)
args.dataset = '/' + args.scenario + '.pkl'
model_parameters = {"args": args}

# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)

# Save Model Parameters
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()
with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

# Specify the number of datasets / scenarios 
n_scenarios = 3

# Create Datahandler class for Herengracht experiment
original_scenario = args.scenario
data_prep = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.081,
				"map_size": np.array([1715,881]),
				"map_center": np.array([-78,-40])}
data_prep.processData(**map_args)

# Create Datahandler class for Prinsengracht scenario
args.scenario = "prinsengracht"
data_prep_2 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.091,
				"map_size": np.array([816,816]),
				"map_center": np.array([-37.1, -37.1])} 
data_prep_2.processData(**map_args)

args.scenario = "bloemgracht"
data_prep_3 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.22,
				"map_size": np.array([1469,928]),
				"map_center": np.array([-161.59,-102.08])} 
data_prep_3.processData(**map_args)

args.scenario = "open_crossing"
data_prep_4 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.6,
				"map_size": np.array([100,100]),
				"map_center": np.array([-30,-30])} 
#data_prep_4.processData(**map_args)

args.scenario = original_scenario

# Import Deep Learning model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

# Create Model Graph
model = NetworkModel(args)

config = tf.ConfigProto()

beta_list = frange_cycle_linear(args.total_training_steps)

# Start Training Session
with tf.Session(config=config) as sess:

	# Load a pre-trained model
	if args.warmstart_model:
		model.warmstart_model(args, sess)
	else:
		# Initialize all TF variables
		sess.run(tf.global_variables_initializer())

	# Load Convnet Model
	try:
		if args.warm_start_convnet:
			model.warmstart_convnet(args, sess)
	except:
		print("Failed to initialized Convnet or Convnet does not exist")

	# if the training was interrupted load last training step index
	try:
		initial_step = int(open(args.model_path + "/tf_log", 'r').read().split('\n')[-2]) + 1
	except:
		initial_step = 1

	epoch = 0
	training_loss = []
	diversity_loss = []
	training_loss.append(0)

	# Set up multithreadign for data handler
	pool = ThreadPool(1)
	res = None
	_model_prediction = []
	start_time = time.time()
	best_loss = float('inf')
	avg_training_loss = np.ones(100)

	dataset = random.randint(0,n_scenarios-1)
	print("Dataset: ", dataset)

	for step in range(initial_step,args.total_training_steps):
		start_time_loop = time.time()

		# Get Next Batch of Data
		if res == None:
			if dataset == 0:
				batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch = data_prep.getBatch()
			if dataset == 1:
				batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch = data_prep_2.getBatch()
			if dataset == 2:
				batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch = data_prep_3.getBatch()
			if dataset == 3:
				batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch = data_prep_4.getBatch()
		else:
			batch = res.get(timeout=5)


		# Create dictionary to feed into the model
		dict = {"batch_x": batch_x,
				"batch_vel": batch_vel,
				"batch_pos": batch_pos,
				"batch_goal": batch_goal,
				  "batch_grid": batch_grid,
				  "batch_ped_grid": batch_ped_grid,
				  "step": step,
				  "beta": beta_list[step],
				  "batch_y": batch_y,
				  "batch_pos_target": batch_pos_target,
				"batch_div": batch_y,
				"other_agents_pos": other_agents_pos
				  }

		feed_dict_train = model.feed_dic(**dict)

		# Here was new epoch before....

		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if dataset == 0:
			model.reset_cells(data_prep.sequence_reset)
		if dataset == 1:
			model.reset_cells(data_prep_2.sequence_reset)
		if dataset == 2:
			model.reset_cells(data_prep_3.sequence_reset)
		if dataset == 3:
			model.reset_cells(data_prep_4.sequence_reset)

		start_time_training = time.time()

		model_output = model.train_step(sess, feed_dict_train, step)

		avg_training_time = time.time() - start_time_training
		avg_loop_time = time.time() - start_time_loop

		training_loss.append(model_output["batch_loss"])

		# Print training info
		if step % print_freq == 0:
			
			if dataset == 0:
				validation_dict = data_prep.getTestBatch()
				model.reset_test_cells(data_prep.val_sequence_reset)
			if dataset == 1:
				validation_dict = data_prep_2.getTestBatch()
				model.reset_test_cells(data_prep_2.val_sequence_reset)
			if dataset == 2:
				validation_dict = data_prep_3.getTestBatch()
				model.reset_test_cells(data_prep_3.val_sequence_reset)
			if dataset == 3:
				validation_dict = data_prep_4.getTestBatch()
				model.reset_test_cells(data_prep_4.val_sequence_reset)

			feed_dict_validation = model.feed_val_dic(**validation_dict)

			#validation_loss, validation_summary, validation_predictions = model.validation_step(sess, feed_dict_train)
			validation_loss, validation_summary, validation_predictions = model.validation_step(sess, feed_dict_validation)

			ellapsed_time = time.time() - start_time

			print(Fore.BLUE + "\n\nEpoch {:d}, Steps: {:d}, Train loss: {:01.2f}, Validation loss: {:01.2f}, Epoch time: {:01.2f} sec"
				  .format(epoch + 1, step, np.mean(training_loss), validation_loss, ellapsed_time)+Style.RESET_ALL)

			if tensorboard_logging:
				model.summary_writer.add_summary(model_output["summary"], step)
				model.summary_writer.flush()
				model.validation_summary_writer.add_summary(validation_summary, step)
				model.validation_summary_writer.flush()

			with open(args.model_path + "/tf_log", 'a') as f:
				f.write(str(step) + '\n')
			curr_loss = (validation_loss+np.mean(avg_training_loss))/2.0
			if curr_loss < best_loss:
				save_path = args.model_path + '/model_ckpt'
				model.full_saver.save(sess, save_path, global_step=step)
				best_loss = curr_loss
				print(Fore.LIGHTCYAN_EX+'Step {}: Saving model under {}'.format(step, save_path))
		
		# Change this back to the top to test what happens...
		if new_epoch:
			dataset = random.randint(0,n_scenarios-1)
			epoch += 1
			print("dataset: ", dataset)
		
		# Makes no sense in a for loop...probably a mistake from bruno because he used a while loop before?
		#step = step + 1

	write_summary(training_loss[-1], args)
	full_path = args.model_path + '/final-model.ckpt'
	model.full_saver.save(sess, full_path)
	print('Saved final model under "{}"'.format(full_path))

	if tensorboard_logging:
		model.summary_writer.close()

sess.close()



build/lib/train_crossdata.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0])) 
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *


import pickle as pkl
import time
from copy import deepcopy

model_name = "VDGNN_multihead"
model_path = '../trained_models/' + model_name
pretrained_convnet_path = "../trained_models/autoencoder_with_ped"
log_dir = model_path + '/log'
data_path = '../data/2_agents_swap/trajs/'
scenario = 'GA3C-CADRL-10-py27'
data_path = '../data/cyberzoo_experiments/'
scenario = 'all_trajectories'
exp_num = 6
data_path = '../data/'
scenario = '20_ped_with_obstacles/short_few_obstacles'

# Hyperparameters
n_epochs = 2

batch_size = 16
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 10
prediction_horizon = 10
prev_horizon = 0

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 128
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 256
x_dim = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 0.8
dropout = False
reg = 1e-4
n_mixtures = 3  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
dropout = False
tensorboard_logging = False

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 5  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 10

print_freq = 2000
save_freq = 500
dt = 0.4

warmstart_model = False
pretrained_convnet = True
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = True
rotated_grid = False
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
submap_resolution = 0.1
submap_width = 6
submap_height = 6
diversity_update = False
predict_positions = False

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default=model_name)
	parser.add_argument('--model_path',
	                    help='Path to directory to save the model (default=""../trained_models/"+model_name").',
	                    type=str, default=model_path)
	parser.add_argument('--pretrained_convnet_path',
	                    help='Path to directory that comprises the pre-trained convnet model (default=" ").',
	                    type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
	                    help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
	                    type=str, default=log_dir)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
	                    type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
	                    default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
	                    default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
	                    default='/' + scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
	                    default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
	                    default=warmstart_model)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
	                    default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
	                    default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
	                    default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
	                    default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
	                    default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
	                    default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_size_lstm_grid',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
	                    help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
	                    type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
	                    type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
	                    type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
	                    help='Number of fully connected layer units after LSTM layer (default=64).',
	                    type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
	                    default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
	                    default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
	                    default=output_dim)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
	                    default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
	                    default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
	                    default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
	                    default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
	                    default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
	                    default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
	                    default=pedestrian_radius)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=int, default=0)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
	                    default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
	                    type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
	                    default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
	                    default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
	                    default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
	                    default=0)
	parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=submap_width)
	parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=submap_height)
	parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
	                    default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--relative_info', help='Use relative info for other agents.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
	                    default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
	                    default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
	                    default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	args = parser.parse_args()

	return args


args = parse_args()

if args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf
tfk = tf.keras

args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)

model_parameters = {"args": args}
print(args)
# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()

if "GA3C" in args.scenario:
	data_prep = dh.DataHandler(args)
	# Collect data online
	if os.path.isfile(args.data_path + args.dataset):
		data_prep.load_data()
	else:
		data_prep.start_node()
else:
	data_prep = dhlstm.DataHandlerLSTM(args)
	# Only used to create a map from png
	map_args = {"file_name": 'map.png',
	            "resolution": 0.1,
	            "map_size": np.array([30., 30.]), }
	data_pickle = args.data_path+args.scenario + "/data" + str(args.prediction_horizon) +"_"+str(args.truncated_backprop_length)+".pickle"
	data_prep.processData(data_pickle,**map_args)

	original_scenario = args.scenario
	args.scenario = "ewap_dataset/seq_eth"
	data_prep_2 = dhlstm.DataHandlerLSTM(args)
	data_pickle = args.data_path + args.scenario + "/data_seq_eth" + str(prediction_horizon) +"_"+str(args.truncated_backprop_length)+ ".pickle"
	data_prep_2.processData(data_pickle, **map_args)

	args.scenario = "st"
	data_prep_3 = dhlstm.DataHandlerLSTM(args)
	data_pickle = args.data_path + args.scenario + "/data_" +args.scenario + str(prediction_horizon)+"_"+str(args.truncated_backprop_length) + ".pickle"
	data_prep_3.processData(data_pickle, **map_args)

	args.scenario = "zara_02"
	data_prep_4 = dhlstm.DataHandlerLSTM(args)
	data_pickle = args.data_path + args.scenario + "/data_" +args.scenario + str(prediction_horizon)+"_"+str(args.truncated_backprop_length) + ".pickle"
	data_prep_4.processData(data_pickle, **map_args)

	args.scenario = original_scenario
# Import model
if "noGrid" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_noGrid import *
	else:
		from src.models.VGDNN_noGrid import *
elif "Grid" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_Grid import *
	else:
		from src.models.VGDNN_Grid import *
elif "RNN" in args.model_name:
	if sys.version_info[0] < 3:
		from RNN import *
	else:
		from src.models.RNN import *
elif "KL" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_KL import *
	else:
		from src.models.VGDNN_KL import *
elif "diversity" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_attention_diversity import *
	else:
		from src.models.VGDNN_attention_diversity import *
elif "ped" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_ped_attention import *
	else:
		from src.models.VGDNN_ped_attention import *
elif "attention" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_attention import *
	else:
		from src.models.VGDNN_attention import *
elif "multihead" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_multihead_attention import *
	else:
		from src.models.VGDNN_multihead_attention import *
elif "simple" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_simple import *
	else:
		from src.models.VGDNN_simple import *
elif "pos" in args.model_name:
	if sys.version_info[0] < 3:
		from VGDNN_pos import *
	else:
		from src.models.VGDNN_pos import *
else:
	if sys.version_info[0] < 3:
		from VGDNN import *
	else:
		from src.models.VGDNN import *

model = NetworkModel(args)

config = tf.ConfigProto(
        device_count = {'GPU': 0}
    )


fig, axarr = pl.subplots(1, 2)
pl.show(block=False)

def clear_plot():
  axarr[0].clear()
  axarr[1].clear()
  axarr[0].set_title('Training Sample')
  axarr[1].set_title('Training Output')

with tf.Session(config=config) as sess:
	# Timing
	avg_time_comp = 0.0
	avg_time_total = 0.0

	if args.warmstart_model:
		model.warmstart_model(args, sess)
	else:
		# Initialize all TF variables
		sess.run(tf.global_variables_initializer())
	model.warmstart_convnet(args, sess)

	# if the trainning was interrupted
	try:
		step = int(open(args.model_path + "/tf_log", 'r').read().split('\n')[-2]) + 1
	except:
		step = 1
	epoch = 0
	training_loss = []
	diversity_loss = []
	dataset = random.randint(0, 3)
	while step <= 20000:
		start_time_loop = time.time()
		# if res == None:

		if dataset == 0:
			batch_x, batch_vel, batch_pos, batch_grid, batch_ped_grid, batch_y, other_agents_pos, new_epoch = data_prep.getBatch()
		if dataset == 1:
			batch_x, batch_vel, batch_pos,batch_grid, batch_ped_grid, batch_y, other_agents_pos, new_epoch = data_prep_2.getBatch()
		if dataset == 2:
			batch_x, batch_vel,batch_pos, batch_grid, batch_ped_grid, batch_y, other_agents_pos, new_epoch = data_prep_3.getBatch()
		if dataset == 3:
			batch_x, batch_vel, batch_pos, batch_grid, batch_ped_grid, batch_y, other_agents_pos, new_epoch = data_prep_4.getBatch()

		if new_epoch:
			dataset = random.randint(0, 3)
			epoch += 1

		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		model.reset_cells(data_prep.sequence_reset)

		# Backup variables to plot
		if args.input_state_dim <4:
			batch_x_backup = deepcopy(batch_vel)
		else:
			batch_x_backup = deepcopy(batch_x)
		batch_grid_backup = deepcopy(batch_grid)
		batch_ped_grid_backup = deepcopy(batch_ped_grid)
		if args.predict_positions:
			batch_y_backup = deepcopy(batch_pos)
		else:
			batch_y_backup = deepcopy(batch_y)
		other_agents_pos_backup = other_agents_pos

		start_time_comp = time.time()

		if step > 2000 and args.diversity_update:
			model.update_test_hidden_state()
			feed_test_dic = model.feed_test_dic(batch_x_backup, batch_grid_backup, batch_ped_grid_backup, batch_y_backup)
			y_model_pred, loss, output_decoder = model.predict(sess, feed_test_dic, False)
			batch_y_varible = np.zeros((args.batch_size,args.truncated_backprop_length,args.prediction_horizon*args.output_dim))
			for t_id in range(len(y_model_pred)):
				for batch_id in range(y_model_pred[0].shape[0]):
					for pred_id in range(args.prediction_horizon):
						new_batch_x = y_model_pred[t_id][batch_id][pred_id*args.output_pred_state_dim*args.n_mixtures]
						new_batch_y = y_model_pred[t_id][batch_id][pred_id * args.output_pred_state_dim * args.n_mixtures]
						batch_y_varible[batch_id,t_id,2*pred_id:2*pred_id+args.output_dim] = np.array([new_batch_x,new_batch_y])
					# Assemble feed dict for training
			feed_dict_train = model.feed_dic(batch_x_backup, batch_grid_backup, batch_ped_grid_backup, step,
					                                 batch_y_backup,batch_y_varible)
		else:
			batch_y_varible = batch_y_backup.copy()
			try:
				if "diversity" in args.model_name:
					feed_dict_train = model.feed_dic(batch_x_backup, batch_grid_backup, batch_ped_grid_backup, step,
					                                 batch_y_backup, batch_pos,batch_y_varible)
				else:
					feed_dict_train = model.feed_dic(batch_x_backup, batch_grid_backup, batch_ped_grid_backup, step,
					                                 batch_y_backup, batch_pos)
			except:
				feed_dict_train = model.feed_dic(batch_x_backup, batch_grid_backup, batch_ped_grid_backup, step,
				                                 batch_y_backup)

		batch_loss, kl_loss, _model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = model.run(
			sess, feed_dict_train)

		if np.mean(autoencoder_loss) > 0.01:
			print("Autoencoder loss: " +str(np.mean(autoencoder_loss)))
			model.run_autoencoder(sess, feed_dict_train)
			model.convnet_saver.save(sess, args.pretrained_convnet_path+ '/final-model.ckpt')

		diversity_loss.append(div_loss)
		training_loss.append(batch_loss)

		avg_time_comp += time.time() - start_time_comp
		avg_time_total += time.time() - start_time_loop

		if tensorboard_logging:
			model.summary_writer.add_summary(_summary_str, step)

		# Print training info
		if step % print_freq == 0:
			"""
			for batch_id in range(batch_grid_backup.shape[0]):
				input_grid = batch_grid_backup[batch_id, 0, :, :]
				output_grid = output_decoder[batch_id, :, :, 0]
				clear_plot()
				sup.plotGrid(batch_grid_backup[batch_id, 0, :, :], axarr[0], color='k', alpha=0.5)
				sup.plotGrid(output_grid, axarr[1], color='r', alpha=0.5)
				fig.canvas.draw()
				pl.show(block=False)
			#plot_batch_OpenCV(step, batch_x, batch_grid, batch_ped_grid, batch_y, other_agents_pos, _model_prediction, args)

			# Testing on a trajectory from a different dataset
			"""
			batch_x, batch_vel, batch_pos,batch_grid, other_agents_info, batch_target, other_agents_pos, traj = data_prep.getTrajectoryAsBatch(
				random.randint(0, len(data_prep.trajectory_set) - 1))  # trajectory_set random.randint(0, len(data_prep.dataset) - 1)

			model.reset_test_cells(np.ones((args.batch_size)))
			predictions = []
			for traj_step in range(batch_x.shape[1]):
				# Assemble feed dict for training
				if args.input_state_dim < 4:
					x = np.repeat(np.repeat(np.expand_dims(batch_vel[:, traj_step, :], axis=0),args.truncated_backprop_length,axis=1),args.batch_size,axis=0)
				else:
					x = np.repeat(np.repeat(np.expand_dims(batch_x[:, traj_step, :], axis=0),args.truncated_backprop_length,axis=1),args.batch_size,axis=0)
				agents = np.repeat(np.repeat(np.expand_dims(other_agents_info[:, traj_step, :], axis=0),args.truncated_backprop_length,axis=1),args.batch_size,axis=0)
				grid = np.repeat(np.repeat(np.expand_dims(batch_grid[:, traj_step, :, :], axis=0),args.truncated_backprop_length,axis=1),args.batch_size,axis=0)
				y = np.repeat(np.repeat(np.expand_dims(batch_target[:, traj_step, :], axis=0),args.truncated_backprop_length,axis=1),args.batch_size,axis=0)
				pos = np.repeat(np.repeat(np.expand_dims(batch_pos[:, traj_step, :], axis=0),args.truncated_backprop_length,axis=1),args.batch_size,axis=0)

				try:
					feed_dict_ = model.feed_test_dic(x, grid, agents, y,pos)
				except:
					feed_dict_ = model.feed_test_dic(x, grid, agents, y)

				y_model_pred, loss, output_decoder = model.predict(sess, feed_dict_, True)

				predictions.append([y_model_pred[0]])

			plot_scenario_vel_OpenCV_simdata([traj], [batch_target], [predictions], args,step)
			plot_local_scenario_vel_OpenCV([traj], [batch_target], [predictions], args,step)

			print(bcolors.WARNING + 'Epoch {0}: batch_loss {1:.5f} \t comp time.: {2:.5f}ms \t tot. time.: {3:.5f}ms'.format(
				epoch,
				batch_loss,
				avg_time_comp * 1000 / print_freq,
				avg_time_total * 1000 / print_freq) + bcolors.ENDC)

			print("Step: " + str(step) + " Loss: " + str(batch_loss) + " KL Loss: " + str(kl_loss) +" Beta: " + str(beta) + " Diversity: " + str(np.mean(diversity_loss))+ " Autoencoder loss: " + str(np.mean(autoencoder_loss)))

			with open(args.model_path + "/tf_log", 'a') as f:
				f.write(str(step) + '\n')
			save_path = args.model_path + '/model_ckpt'
			model.full_saver.save(sess, save_path, global_step=step)

			print('Step {}: Saving model under {}'.format(step, save_path))

			print('Epoch {0},Step {0}: Learning_rate {1:.5f} \t'.format(epoch, lr))
			print("Total steps: " + str(step))

		step = step + 1

	write_summary(training_loss[-1], args)
	full_path = args.model_path + '/final-model.ckpt'
	model.full_saver.save(sess, full_path)
	print(bcolors.OKBLUE + 'Final loss: batch_loss ' + str(training_loss[-1]) + bcolors.ENDC)
	print('Saved final model under "{}"'.format(full_path))

	if tensorboard_logging:
		model.summary_writer.close()

sess.close()

print("test")

build/lib/train_multidataset.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
import json
import importlib
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils import DataHandlerRoboat as dhroboat
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
from datetime import datetime
import pickle as pkl
import time
from copy import deepcopy
from multiprocessing.pool import ThreadPool
import colorama
from colorama import Fore, Style

#os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

pretrained_convnet_path = "../trained_models/autoencoder"
exp_num = 6
data_path = '../data/roboat/'
scenario = 'herengracht'
no_stopping = True

# Hyperparameters
n_epochs = 0

batch_size = 16 #128
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 8
prediction_horizon = 12
prev_horizon = 8

# Model params
# rnn_state_size = 32
# rnn_state_size_lstm_grid = 256
# rnn_state_size_lstm_ped = 128
# rnn_state_ped_size = 16
# rnn_state_size_lstm_concat = 512
# prior_size = 512
# latent_space_size = 256
# x_dim = 512
# fc_hidden_unit_size = 256

rnn_state_size = 32
rnn_state_size_lstm_grid = 64
rnn_state_size_lstm_ped = 64 
rnn_state_size_lstm_concat = 128
latent_space_size = 0

# Not used
rnn_state_ped_size = 1
prior_size = 1
x_dim = 1
fc_hidden_unit_size = 1

# Regularization and training params
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 18
tensorboard_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 5

print_freq = 200
save_freq = 500
total_training_steps = 50000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = True
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = False
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
#submap_resolution = 0.081
#submap_width = 4.86
#submap_height = 4.86
submap_size = 60 # In pixels
submap_span_real = 10 # In Map coordinates
relative_covariance = 10.0
diversity_update = False
predict_positions = False
warm_start_convnet = True
debug_plotting = False

# Dataset division
train_set = 0.8

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
						help='Path to directory that comprises the model (default="model_name").',
						type=str, default= "RNN")
	parser.add_argument('--model_path',
						help='Path to directory to save the model (default=""../trained_models/"+model_name").',
						type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
						help='Path to directory that comprises the pre-trained convnet model (default=" ").',
						type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
						help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
						type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
						type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
						default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
						default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
						default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
						default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
						default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
						default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
						default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--total_training_steps', help='Number of training steps (default=20000).', type=int, default=total_training_steps)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
						default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
						default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
						default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
						default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
						default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
						default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
						type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
						default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
						default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
						default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
						help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
						type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
						type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
						type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
						help='Number of fully connected layer units after LSTM layer (default=64).',
						type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
						default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
						default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
						default=output_dim)
	parser.add_argument('--goal_size', help='Goal dimension (default=).', type=int,
						default=2)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
						default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
						default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
						default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
						default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
						default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
						default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
						default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the network.', type=int,
						default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=sup.str2bool, default=debug_plotting)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
						default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
						type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
						default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
						default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
						default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
						default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
						default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
						default=0)
	#parser.add_argument('--submap_width', help='width of occupancy grid (For initialization of arrays)', type=int, default=submap_width)
	#parser.add_argument('--submap_height', help='height of occupancy grid (for initiatization of arrays)', type=int, default=submap_height)
	parser.add_argument('--submap_size', help='size of submap in pixels (60)', type=int, default=submap_size)
	parser.add_argument('--submap_span_real', help='span of submap occupancy grid', type=int, default=submap_span_real)
	#parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	#                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
						default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
						default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
						default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
						default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
						default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
						default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
						default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
						default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	parser.add_argument('--kl_weight', help='Experimenting with additional weight for KL loss', type=float, default=1)
	parser.add_argument('--relative_covariance', help='Changing covariance of relative positions', type=float, default=relative_covariance)
	args = parser.parse_args()

	return args

args = parse_args()


# Cyclical KL Annealing weight
def frange_cycle_linear(n_iter, start=0.0, stop=1.0,  n_cycle=4, ratio=0.5):
	L = np.ones(n_iter) * stop
	period = n_iter/n_cycle
	step = (stop-start)/(period*ratio) # linear schedule

	for c in range(n_cycle):
		v, i = start, 0
		while v <= stop and (int(i+c*period) < n_iter):
			L[int(i+c*period)] = v
			v += step
			i += 1
	return L

# Enable / Disable GPU
if args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

# Create Log and Model Directory to save training model
args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)
args.log_dir = args.model_path + '/log'
if not os.path.exists(args.log_dir):
	os.makedirs(args.log_dir)
args.dataset = '/' + args.scenario + '.pkl'
model_parameters = {"args": args}

# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)

# Save Model Parameters
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()
with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

# Specify the number of datasets / scenarios 
n_scenarios = 4

# Create Datahandler class for Herengracht experiment
original_scenario = args.scenario
data_prep = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.081,
				"map_size": np.array([1715,881]),
				"map_center": np.array([-78,-40])}
data_prep.processData(**map_args)

# Create Datahandler class for Prinsengracht scenario
args.scenario = "prinsengracht"
data_prep_2 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.091,
				"map_size": np.array([816,816]),
				"map_center": np.array([-37.1, -37.1])} 
data_prep_2.processData(**map_args)

args.scenario = "bloemgracht"
data_prep_3 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.22,
				"map_size": np.array([1469,928]),
				"map_center": np.array([-161.59,-102.08])} 
data_prep_3.processData(**map_args)

args.scenario = "open_crossing"
data_prep_4 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.6,
				"map_size": np.array([100,100]),
				"map_center": np.array([-30,-30])} 
data_prep_4.processData(**map_args)

args.scenario = original_scenario

# Initialize main data handler
dpreps = [data_prep, data_prep_2, data_prep_3, data_prep_4]
data_roboat = dhroboat.DataHandlerRoboat(args)

# Combine trajectory sets
data_roboat.CombineTrajectorySets(dpreps)

min = 100
print(len(data_roboat.trajectory_set))
for i in range(len(data_roboat.trajectory_set)):
	if len(data_roboat.trajectory_set[i][1].pose_vec) < min:
		min = len(data_roboat.trajectory_set[i][1].pose_vec)
print(min)
# Remove end of trajectories if at goal, and trajectories that are too short
if no_stopping:
	data_roboat.RemoveStops()

min = 100
print(len(data_roboat.trajectory_set))
for i in range(len(data_roboat.trajectory_set)):
	if len(data_roboat.trajectory_set[i][1].pose_vec) < min:
		min = len(data_roboat.trajectory_set[i][1].pose_vec)
print(min)
# Import Deep Learning model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

# Create Model Graph
model = NetworkModel(args)

config = tf.ConfigProto()

beta_list = frange_cycle_linear(args.total_training_steps)

# Start Training Session
with tf.Session(config=config) as sess:

	# Load a pre-trained model
	if args.warmstart_model:
		model.warmstart_model(args, sess)
	else:
		# Initialize all TF variables
		sess.run(tf.global_variables_initializer())

	# Load Convnet Model
	try:
		if args.warm_start_convnet:
			model.warmstart_convnet(args, sess)
	except:
		print("Failed to initialized Convnet or Convnet does not exist")

	# if the training was interrupted load last training step index
	try:
		initial_step = int(open(args.model_path + "/tf_log", 'r').read().split('\n')[-2]) + 1
	except:
		initial_step = 1

	epoch = 0
	training_loss = []
	diversity_loss = []
	training_loss.append(0)

	# Set up multithreadign for data handler
	pool = ThreadPool(1)
	res = None
	_model_prediction = []
	start_time = time.time()
	best_loss = float('inf')
	avg_training_loss = np.ones(100)


	for step in range(initial_step,args.total_training_steps):
		start_time_loop = time.time()

		# Get Next Batch of Data
		#batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch, train_log = data_roboat.getBatch(dpreps)
		batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch, train_log = data_roboat.getRoboatBatch(dpreps)

		# Create dictionary to feed into the model
		dict = {"batch_x": batch_x,
				"batch_vel": batch_vel,
				"batch_pos": batch_pos,
				"batch_goal": batch_goal,
				  "batch_grid": batch_grid,
				  "batch_ped_grid": batch_ped_grid,
				  "step": step,
				  "beta": beta_list[step],
				  "batch_y": batch_y,
				  "batch_pos_target": batch_pos_target,
				"batch_div": batch_y,
				"other_agents_pos": other_agents_pos
				  }

		feed_dict_train = model.feed_dic(**dict)

		# Here was new epoch before....

		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		model.reset_cells(data_roboat.sequence_reset)

		start_time_training = time.time()

		model_output = model.train_step(sess, feed_dict_train, step)

		avg_training_time = time.time() - start_time_training
		avg_loop_time = time.time() - start_time_loop

		training_loss.append(model_output["batch_loss"])

		# Print training info
		if step % print_freq == 0:
			
			# Get test batch for model validation
			#validation_dict, val_log = data_roboat.getTestBatch(dpreps)
			validation_dict, val_log = data_roboat.getRoboatValidationBatch(dpreps)
			model.reset_test_cells(data_roboat.val_sequence_reset)

			feed_dict_validation = model.feed_val_dic(**validation_dict)
			validation_loss, validation_summary, validation_predictions = model.validation_step(sess, feed_dict_validation)
			#validation_loss, validation_summary, validation_predictions = model.validation_step(sess, feed_dict_train)

			autoencoder_loss = model_output["autoencoder_loss"]

			ellapsed_time = time.time() - start_time

			print(Fore.BLUE + "\n\nEpoch {:d}, Steps: {:d}, Train loss: {:01.2f}, Validation loss: {:01.2f}, AutoEncoder Loss: {:01.2f},  Epoch time: {:01.2f} sec"
				  .format(epoch + 1, step, np.mean(training_loss), validation_loss, autoencoder_loss, ellapsed_time)+Style.RESET_ALL)

			if tensorboard_logging:
				model.summary_writer.add_summary(model_output["summary"], step)
				model.summary_writer.flush()
				model.validation_summary_writer.add_summary(validation_summary, step)
				model.validation_summary_writer.flush()

			with open(args.model_path + "/tf_log", 'a') as f:
				f.write(str(step) + '\n')
				f.write("TrainLog: " + str(train_log) + '\n')
				f.write("ValidationLog: " + str(val_log) + '\n')
			curr_loss = (validation_loss+np.mean(avg_training_loss))/2.0
			if curr_loss < best_loss:
				save_path = args.model_path + '/model_ckpt'
				model.full_saver.save(sess, save_path, global_step=step)
				best_loss = curr_loss
				print(Fore.LIGHTCYAN_EX+'Step {}: Saving model under {}'.format(step, save_path))
		
		# Change this back to the top to test what happens...
		if new_epoch:
			epoch += 1

	write_summary(training_loss[-1], args)
	full_path = args.model_path + '/final-model.ckpt'
	model.full_saver.save(sess, full_path)
	print('Saved final model under "{}"'.format(full_path))

	if tensorboard_logging:
		model.summary_writer.close()

sess.close()



build/lib/train_roboat.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import tensorflow as tf
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
if sys.version_info[0] < 3:
  sys.path.append('../src/data_utils')
  sys.path.append('../src/models')
  import DataHandler as dh
  from seqRNNModel import *
else:
  from src.data_utils import DataHandler as dh
  from src.models.seqRNNModel import *
  from src.data_utils.plot_utils import *

tfk = tf.keras

model_name = "seq_to_seq"
input_dim = 8
latent_dim = 256
output_dim = 2
batch_size = 16
learning_rate_init = 0.01
grads_clip = 10.0
n_epochs = 100
stateful = False
prediction_horizon = 10
prev_horizon = 10
truncated_backprop_length = 10
regularization_weight = 0.0001
dropout = 0.2

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default=model_name)
	parser.add_argument('--data_path', help='Data directory', type=str,
	                    default='../data')
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=int,
	                    default=input_dim)
	parser.add_argument('--latent_dim', help='Latent dimension.', type=int,
	                    default=latent_dim)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--prediction_horizon', help='Prediction Horizon.', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--output_dim', help='Output dimension.', type=int,
	                    default=output_dim)
	parser.add_argument('--batch_size', help='Batch size.', type=int,
	                    default=batch_size)
	parser.add_argument('--learning_rate_init', help='Learning rate.', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--grads_clip', help='Grad clip.', type=float,
	                    default=grads_clip)
	parser.add_argument('--n_epochs', help='Number of epochs.', type=float,
	                    default=n_epochs)
	parser.add_argument('--stateful', help='Stateful.', type=bool,
	                    default=stateful)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--dropout', help='Dropout probability (default=0.005).', type=float,
	                    default=dropout)
	args = parser.parse_args()

	return args


args = parse_args()

data_prep = dh.DataHandler()
# Collect data online
if os.path.isfile(data_prep.args.save_path):
	data_prep.load_data()
else:
	data_prep.start_node()

batch_enc, batch_dec, batch_dec_target, new_epoch = data_prep.getBatchTBP()

model = RNNModel(args)

model.model.compile(optimizer=tfk.optimizers.Adam(lr=args.learning_rate_init,clipnorm=args.grads_clip),
                             loss='mean_squared_error')


history = model.model.fit([batch_enc,batch_dec],batch_dec_target,
                epochs=n_epochs,
                batch_size=16,
								#validation_split=0.2,
                shuffle=True)

# Plot training & validation loss values
plt.plot(history.history['loss'])
plt.title('Model loss during training')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train'], loc='upper left')
plt.show()

batch_enc_seq, batch_dec_seq, traj = data_prep.getTrajectoryAsBatchTBP()

predictions = model.decode_sequence([batch_enc_seq,batch_dec_seq])

#plot_scenario(batch_enc_seq,batch_dec_seq,traj,predictions,data_prep.occupancy_grid,args)
plot_scenario_vel(batch_enc_seq,batch_dec_seq,traj,predictions,data_prep.occupancy_grid,args)
print("test")


build/lib/train_tbp.py
eliatrevisan/IA-MPPI-LBM
from setuptools import setup

setup(
    name='social_vdgnn',
    version='1.0.0',
    description='',
    url='https://github.com/bbrito/',
    author='B. Brito, H. Zhu, and J. Alonso-Mora, TU Delft',  # Optional
    keywords='pedestrian prediction',  # Optional
    python_requires='>=3.0, <4',
)
setup.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/__future__.py
social_vdgnn/lib/python3.6/__future__.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/_bootlocale.py
social_vdgnn/lib/python3.6/_bootlocale.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/_collections_abc.py
social_vdgnn/lib/python3.6/_collections_abc.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/_dummy_thread.py
social_vdgnn/lib/python3.6/_dummy_thread.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/_weakrefset.py
social_vdgnn/lib/python3.6/_weakrefset.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/abc.py
social_vdgnn/lib/python3.6/abc.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/base64.py
social_vdgnn/lib/python3.6/base64.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/bisect.py
social_vdgnn/lib/python3.6/bisect.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/codecs.py
social_vdgnn/lib/python3.6/codecs.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/copy.py
social_vdgnn/lib/python3.6/copy.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/copyreg.py
social_vdgnn/lib/python3.6/copyreg.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/enum.py
social_vdgnn/lib/python3.6/enum.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/fnmatch.py
social_vdgnn/lib/python3.6/fnmatch.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/functools.py
social_vdgnn/lib/python3.6/functools.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/genericpath.py
social_vdgnn/lib/python3.6/genericpath.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/hashlib.py
social_vdgnn/lib/python3.6/hashlib.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/heapq.py
social_vdgnn/lib/python3.6/heapq.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/hmac.py
social_vdgnn/lib/python3.6/hmac.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/imp.py
social_vdgnn/lib/python3.6/imp.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/io.py
social_vdgnn/lib/python3.6/io.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/keyword.py
social_vdgnn/lib/python3.6/keyword.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/linecache.py
social_vdgnn/lib/python3.6/linecache.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/locale.py
social_vdgnn/lib/python3.6/locale.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/ntpath.py
social_vdgnn/lib/python3.6/ntpath.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/operator.py
social_vdgnn/lib/python3.6/operator.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/os.py
social_vdgnn/lib/python3.6/os.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/posixpath.py
social_vdgnn/lib/python3.6/posixpath.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/random.py
social_vdgnn/lib/python3.6/random.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/re.py
social_vdgnn/lib/python3.6/re.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/reprlib.py
social_vdgnn/lib/python3.6/reprlib.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/rlcompleter.py
social_vdgnn/lib/python3.6/rlcompleter.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/shutil.py
social_vdgnn/lib/python3.6/shutil.py
eliatrevisan/IA-MPPI-LBM
'''
OpenCV Python binary extension loader
'''
import os
import importlib
import sys

from .cv2 import *
from .cv2 import _registerMatType
from . import mat_wrapper
from . import gapi
from . import misc
from . import utils
from . import data
from . import version

__all__ = []

try:
    import numpy
    import numpy.core.multiarray
except ImportError:
    print('OpenCV bindings requires "numpy" package.')
    print('Install it via command:')
    print('    pip install numpy')
    raise

# TODO
# is_x64 = sys.maxsize > 2**32


def __load_extra_py_code_for_module(base, name, enable_debug_print=False):
    module_name = "{}.{}".format(__name__, name)
    export_module_name = "{}.{}".format(base, name)
    native_module = sys.modules.pop(module_name, None)
    try:
        py_module = importlib.import_module(module_name)
    except ImportError as err:
        if enable_debug_print:
            print("Can't load Python code for module:", module_name,
                  ". Reason:", err)
        # Extension doesn't contain extra py code
        return False

    if not hasattr(base, name):
        setattr(sys.modules[base], name, py_module)
    sys.modules[export_module_name] = py_module
    # If it is C extension module it is already loaded by cv2 package
    if native_module:
        setattr(py_module, "_native", native_module)
        for k, v in filter(lambda kv: not hasattr(py_module, kv[0]),
                           native_module.__dict__.items()):
            if enable_debug_print: print('    symbol({}): {} = {}'.format(name, k, v))
            setattr(py_module, k, v)
    return True


def __collect_extra_submodules(enable_debug_print=False):
    def modules_filter(module):
        return all((
             # module is not internal
             not module.startswith("_"),
             not module.startswith("python-"),
             # it is not a file
             os.path.isdir(os.path.join(_extra_submodules_init_path, module))
        ))
    if sys.version_info[0] < 3:
        if enable_debug_print:
            print("Extra submodules is loaded only for Python 3")
        return []

    __INIT_FILE_PATH = os.path.abspath(__file__)
    _extra_submodules_init_path = os.path.dirname(__INIT_FILE_PATH)
    return filter(modules_filter, os.listdir(_extra_submodules_init_path))


def bootstrap():
    import sys

    import copy
    save_sys_path = copy.copy(sys.path)

    if hasattr(sys, 'OpenCV_LOADER'):
        print(sys.path)
        raise ImportError('ERROR: recursion is detected during loading of "cv2" binary extensions. Check OpenCV installation.')
    sys.OpenCV_LOADER = True

    DEBUG = False
    if hasattr(sys, 'OpenCV_LOADER_DEBUG'):
        DEBUG = True

    import platform
    if DEBUG: print('OpenCV loader: os.name="{}"  platform.system()="{}"'.format(os.name, str(platform.system())))

    LOADER_DIR = os.path.dirname(os.path.abspath(os.path.realpath(__file__)))

    PYTHON_EXTENSIONS_PATHS = []
    BINARIES_PATHS = []

    g_vars = globals()
    l_vars = locals()

    if sys.version_info[:2] < (3, 0):
        from . load_config_py2 import exec_file_wrapper
    else:
        from . load_config_py3 import exec_file_wrapper

    def load_first_config(fnames, required=True):
        for fname in fnames:
            fpath = os.path.join(LOADER_DIR, fname)
            if not os.path.exists(fpath):
                if DEBUG: print('OpenCV loader: config not found, skip: {}'.format(fpath))
                continue
            if DEBUG: print('OpenCV loader: loading config: {}'.format(fpath))
            exec_file_wrapper(fpath, g_vars, l_vars)
            return True
        if required:
            raise ImportError('OpenCV loader: missing configuration file: {}. Check OpenCV installation.'.format(fnames))

    load_first_config(['config.py'], True)
    load_first_config([
        'config-{}.{}.py'.format(sys.version_info[0], sys.version_info[1]),
        'config-{}.py'.format(sys.version_info[0])
    ], True)

    if DEBUG: print('OpenCV loader: PYTHON_EXTENSIONS_PATHS={}'.format(str(l_vars['PYTHON_EXTENSIONS_PATHS'])))
    if DEBUG: print('OpenCV loader: BINARIES_PATHS={}'.format(str(l_vars['BINARIES_PATHS'])))

    applySysPathWorkaround = False
    if hasattr(sys, 'OpenCV_REPLACE_SYS_PATH_0'):
        applySysPathWorkaround = True
    else:
        try:
            BASE_DIR = os.path.dirname(LOADER_DIR)
            if sys.path[0] == BASE_DIR or os.path.realpath(sys.path[0]) == BASE_DIR:
                applySysPathWorkaround = True
        except:
            if DEBUG: print('OpenCV loader: exception during checking workaround for sys.path[0]')
            pass  # applySysPathWorkaround is False

    for p in reversed(l_vars['PYTHON_EXTENSIONS_PATHS']):
        sys.path.insert(1 if not applySysPathWorkaround else 0, p)

    if os.name == 'nt':
        if sys.version_info[:2] >= (3, 8):  # https://github.com/python/cpython/pull/12302
            for p in l_vars['BINARIES_PATHS']:
                try:
                    os.add_dll_directory(p)
                except Exception as e:
                    if DEBUG: print('Failed os.add_dll_directory(): '+ str(e))
                    pass
        os.environ['PATH'] = ';'.join(l_vars['BINARIES_PATHS']) + ';' + os.environ.get('PATH', '')
        if DEBUG: print('OpenCV loader: PATH={}'.format(str(os.environ['PATH'])))
    else:
        # amending of LD_LIBRARY_PATH works for sub-processes only
        os.environ['LD_LIBRARY_PATH'] = ':'.join(l_vars['BINARIES_PATHS']) + ':' + os.environ.get('LD_LIBRARY_PATH', '')

    if DEBUG: print("Relink everything from native cv2 module to cv2 package")

    py_module = sys.modules.pop("cv2")

    native_module = importlib.import_module("cv2.cv2")

    sys.modules["cv2"] = py_module
    setattr(py_module, "_native", native_module)

    for item_name, item in filter(lambda kv: kv[0] not in ("__file__", "__loader__", "__spec__",
                                                           "__name__", "__package__"),
                                  native_module.__dict__.items()):
        if item_name not in g_vars:
            g_vars[item_name] = item

    sys.path = save_sys_path  # multiprocessing should start from bootstrap code (https://github.com/opencv/opencv/issues/18502)

    try:
        del sys.OpenCV_LOADER
    except Exception as e:
        if DEBUG:
            print("Exception during delete OpenCV_LOADER:", e)

    if DEBUG: print('OpenCV loader: binary extension... OK')

    for submodule in __collect_extra_submodules(DEBUG):
        if __load_extra_py_code_for_module("cv2", submodule, DEBUG):
            if DEBUG: print("Extra Python code for", submodule, "is loaded")

    if DEBUG: print('OpenCV loader: DONE')


bootstrap()

social_vdgnn/lib/python3.6/site-packages/cv2/__init__.py
eliatrevisan/IA-MPPI-LBM
# Protocol Buffers - Google's data interchange format
# Copyright 2008 Google Inc.  All rights reserved.
# https://developers.google.com/protocol-buffers/
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Descriptors essentially contain exactly the information found in a .proto
file, in types that make this information accessible in Python.
"""

__author__ = 'robinson@google.com (Will Robinson)'

import threading
import warnings

from google.protobuf.internal import api_implementation

_USE_C_DESCRIPTORS = False
if api_implementation.Type() == 'cpp':
  # Used by MakeDescriptor in cpp mode
  import binascii
  import os
  from google.protobuf.pyext import _message
  _USE_C_DESCRIPTORS = True


class Error(Exception):
  """Base error for this module."""


class TypeTransformationError(Error):
  """Error transforming between python proto type and corresponding C++ type."""


if _USE_C_DESCRIPTORS:
  # This metaclass allows to override the behavior of code like
  #     isinstance(my_descriptor, FieldDescriptor)
  # and make it return True when the descriptor is an instance of the extension
  # type written in C++.
  class DescriptorMetaclass(type):
    def __instancecheck__(cls, obj):
      if super(DescriptorMetaclass, cls).__instancecheck__(obj):
        return True
      if isinstance(obj, cls._C_DESCRIPTOR_CLASS):
        return True
      return False
else:
  # The standard metaclass; nothing changes.
  DescriptorMetaclass = type


class _Lock(object):
  """Wrapper class of threading.Lock(), which is allowed by 'with'."""

  def __new__(cls):
    self = object.__new__(cls)
    self._lock = threading.Lock()  # pylint: disable=protected-access
    return self

  def __enter__(self):
    self._lock.acquire()

  def __exit__(self, exc_type, exc_value, exc_tb):
    self._lock.release()


_lock = threading.Lock()


def _Deprecated(name):
  if _Deprecated.count > 0:
    _Deprecated.count -= 1
    warnings.warn(
        'Call to deprecated create function %s(). Note: Create unlinked '
        'descriptors is going to go away. Please use get/find descriptors from '
        'generated code or query the descriptor_pool.'
        % name,
        category=DeprecationWarning, stacklevel=3)


# Deprecated warnings will print 100 times at most which should be enough for
# users to notice and do not cause timeout.
_Deprecated.count = 100


_internal_create_key = object()


class DescriptorBase(metaclass=DescriptorMetaclass):

  """Descriptors base class.

  This class is the base of all descriptor classes. It provides common options
  related functionality.

  Attributes:
    has_options:  True if the descriptor has non-default options.  Usually it
        is not necessary to read this -- just call GetOptions() which will
        happily return the default instance.  However, it's sometimes useful
        for efficiency, and also useful inside the protobuf implementation to
        avoid some bootstrapping issues.
  """

  if _USE_C_DESCRIPTORS:
    # The class, or tuple of classes, that are considered as "virtual
    # subclasses" of this descriptor class.
    _C_DESCRIPTOR_CLASS = ()

  def __init__(self, options, serialized_options, options_class_name):
    """Initialize the descriptor given its options message and the name of the
    class of the options message. The name of the class is required in case
    the options message is None and has to be created.
    """
    self._options = options
    self._options_class_name = options_class_name
    self._serialized_options = serialized_options

    # Does this descriptor have non-default options?
    self.has_options = (options is not None) or (serialized_options is not None)

  def _SetOptions(self, options, options_class_name):
    """Sets the descriptor's options

    This function is used in generated proto2 files to update descriptor
    options. It must not be used outside proto2.
    """
    self._options = options
    self._options_class_name = options_class_name

    # Does this descriptor have non-default options?
    self.has_options = options is not None

  def GetOptions(self):
    """Retrieves descriptor options.

    This method returns the options set or creates the default options for the
    descriptor.
    """
    if self._options:
      return self._options

    from google.protobuf import descriptor_pb2
    try:
      options_class = getattr(descriptor_pb2,
                              self._options_class_name)
    except AttributeError:
      raise RuntimeError('Unknown options class name %s!' %
                         (self._options_class_name))

    with _lock:
      if self._serialized_options is None:
        self._options = options_class()
      else:
        self._options = _ParseOptions(options_class(),
                                      self._serialized_options)

      return self._options


class _NestedDescriptorBase(DescriptorBase):
  """Common class for descriptors that can be nested."""

  def __init__(self, options, options_class_name, name, full_name,
               file, containing_type, serialized_start=None,
               serialized_end=None, serialized_options=None):
    """Constructor.

    Args:
      options: Protocol message options or None
        to use default message options.
      options_class_name (str): The class name of the above options.
      name (str): Name of this protocol message type.
      full_name (str): Fully-qualified name of this protocol message type,
        which will include protocol "package" name and the name of any
        enclosing types.
      file (FileDescriptor): Reference to file info.
      containing_type: if provided, this is a nested descriptor, with this
        descriptor as parent, otherwise None.
      serialized_start: The start index (inclusive) in block in the
        file.serialized_pb that describes this descriptor.
      serialized_end: The end index (exclusive) in block in the
        file.serialized_pb that describes this descriptor.
      serialized_options: Protocol message serialized options or None.
    """
    super(_NestedDescriptorBase, self).__init__(
        options, serialized_options, options_class_name)

    self.name = name
    # TODO(falk): Add function to calculate full_name instead of having it in
    #             memory?
    self.full_name = full_name
    self.file = file
    self.containing_type = containing_type

    self._serialized_start = serialized_start
    self._serialized_end = serialized_end

  def CopyToProto(self, proto):
    """Copies this to the matching proto in descriptor_pb2.

    Args:
      proto: An empty proto instance from descriptor_pb2.

    Raises:
      Error: If self couldn't be serialized, due to to few constructor
        arguments.
    """
    if (self.file is not None and
        self._serialized_start is not None and
        self._serialized_end is not None):
      proto.ParseFromString(self.file.serialized_pb[
          self._serialized_start:self._serialized_end])
    else:
      raise Error('Descriptor does not contain serialization.')


class Descriptor(_NestedDescriptorBase):

  """Descriptor for a protocol message type.

  Attributes:
      name (str): Name of this protocol message type.
      full_name (str): Fully-qualified name of this protocol message type,
          which will include protocol "package" name and the name of any
          enclosing types.
      containing_type (Descriptor): Reference to the descriptor of the type
          containing us, or None if this is top-level.
      fields (list[FieldDescriptor]): Field descriptors for all fields in
          this type.
      fields_by_number (dict(int, FieldDescriptor)): Same
          :class:`FieldDescriptor` objects as in :attr:`fields`, but indexed
          by "number" attribute in each FieldDescriptor.
      fields_by_name (dict(str, FieldDescriptor)): Same
          :class:`FieldDescriptor` objects as in :attr:`fields`, but indexed by
          "name" attribute in each :class:`FieldDescriptor`.
      nested_types (list[Descriptor]): Descriptor references
          for all protocol message types nested within this one.
      nested_types_by_name (dict(str, Descriptor)): Same Descriptor
          objects as in :attr:`nested_types`, but indexed by "name" attribute
          in each Descriptor.
      enum_types (list[EnumDescriptor]): :class:`EnumDescriptor` references
          for all enums contained within this type.
      enum_types_by_name (dict(str, EnumDescriptor)): Same
          :class:`EnumDescriptor` objects as in :attr:`enum_types`, but
          indexed by "name" attribute in each EnumDescriptor.
      enum_values_by_name (dict(str, EnumValueDescriptor)): Dict mapping
          from enum value name to :class:`EnumValueDescriptor` for that value.
      extensions (list[FieldDescriptor]): All extensions defined directly
          within this message type (NOT within a nested type).
      extensions_by_name (dict(str, FieldDescriptor)): Same FieldDescriptor
          objects as :attr:`extensions`, but indexed by "name" attribute of each
          FieldDescriptor.
      is_extendable (bool):  Does this type define any extension ranges?
      oneofs (list[OneofDescriptor]): The list of descriptors for oneof fields
          in this message.
      oneofs_by_name (dict(str, OneofDescriptor)): Same objects as in
          :attr:`oneofs`, but indexed by "name" attribute.
      file (FileDescriptor): Reference to file descriptor.

  """

  if _USE_C_DESCRIPTORS:
    _C_DESCRIPTOR_CLASS = _message.Descriptor

    def __new__(
        cls,
        name=None,
        full_name=None,
        filename=None,
        containing_type=None,
        fields=None,
        nested_types=None,
        enum_types=None,
        extensions=None,
        options=None,
        serialized_options=None,
        is_extendable=True,
        extension_ranges=None,
        oneofs=None,
        file=None,  # pylint: disable=redefined-builtin
        serialized_start=None,
        serialized_end=None,
        syntax=None,
        create_key=None):
      _message.Message._CheckCalledFromGeneratedFile()
      return _message.default_pool.FindMessageTypeByName(full_name)

  # NOTE(tmarek): The file argument redefining a builtin is nothing we can
  # fix right now since we don't know how many clients already rely on the
  # name of the argument.
  def __init__(self, name, full_name, filename, containing_type, fields,
               nested_types, enum_types, extensions, options=None,
               serialized_options=None,
               is_extendable=True, extension_ranges=None, oneofs=None,
               file=None, serialized_start=None, serialized_end=None,  # pylint: disable=redefined-builtin
               syntax=None, create_key=None):
    """Arguments to __init__() are as described in the description
    of Descriptor fields above.

    Note that filename is an obsolete argument, that is not used anymore.
    Please use file.name to access this as an attribute.
    """
    if create_key is not _internal_create_key:
      _Deprecated('Descriptor')

    super(Descriptor, self).__init__(
        options, 'MessageOptions', name, full_name, file,
        containing_type, serialized_start=serialized_start,
        serialized_end=serialized_end, serialized_options=serialized_options)

    # We have fields in addition to fields_by_name and fields_by_number,
    # so that:
    #   1. Clients can index fields by "order in which they're listed."
    #   2. Clients can easily iterate over all fields with the terse
    #      syntax: for f in descriptor.fields: ...
    self.fields = fields
    for field in self.fields:
      field.containing_type = self
    self.fields_by_number = dict((f.number, f) for f in fields)
    self.fields_by_name = dict((f.name, f) for f in fields)
    self._fields_by_camelcase_name = None

    self.nested_types = nested_types
    for nested_type in nested_types:
      nested_type.containing_type = self
    self.nested_types_by_name = dict((t.name, t) for t in nested_types)

    self.enum_types = enum_types
    for enum_type in self.enum_types:
      enum_type.containing_type = self
    self.enum_types_by_name = dict((t.name, t) for t in enum_types)
    self.enum_values_by_name = dict(
        (v.name, v) for t in enum_types for v in t.values)

    self.extensions = extensions
    for extension in self.extensions:
      extension.extension_scope = self
    self.extensions_by_name = dict((f.name, f) for f in extensions)
    self.is_extendable = is_extendable
    self.extension_ranges = extension_ranges
    self.oneofs = oneofs if oneofs is not None else []
    self.oneofs_by_name = dict((o.name, o) for o in self.oneofs)
    for oneof in self.oneofs:
      oneof.containing_type = self
    self.syntax = syntax or "proto2"

  @property
  def fields_by_camelcase_name(self):
    """Same FieldDescriptor objects as in :attr:`fields`, but indexed by
    :attr:`FieldDescriptor.camelcase_name`.
    """
    if self._fields_by_camelcase_name is None:
      self._fields_by_camelcase_name = dict(
          (f.camelcase_name, f) for f in self.fields)
    return self._fields_by_camelcase_name

  def EnumValueName(self, enum, value):
    """Returns the string name of an enum value.

    This is just a small helper method to simplify a common operation.

    Args:
      enum: string name of the Enum.
      value: int, value of the enum.

    Returns:
      string name of the enum value.

    Raises:
      KeyError if either the Enum doesn't exist or the value is not a valid
        value for the enum.
    """
    return self.enum_types_by_name[enum].values_by_number[value].name

  def CopyToProto(self, proto):
    """Copies this to a descriptor_pb2.DescriptorProto.

    Args:
      proto: An empty descriptor_pb2.DescriptorProto.
    """
    # This function is overridden to give a better doc comment.
    super(Descriptor, self).CopyToProto(proto)


# TODO(robinson): We should have aggressive checking here,
# for example:
#   * If you specify a repeated field, you should not be allowed
#     to specify a default value.
#   * [Other examples here as needed].
#
# TODO(robinson): for this and other *Descriptor classes, we
# might also want to lock things down aggressively (e.g.,
# prevent clients from setting the attributes).  Having
# stronger invariants here in general will reduce the number
# of runtime checks we must do in reflection.py...
class FieldDescriptor(DescriptorBase):

  """Descriptor for a single field in a .proto file.

  Attributes:
    name (str): Name of this field, exactly as it appears in .proto.
    full_name (str): Name of this field, including containing scope.  This is
      particularly relevant for extensions.
    index (int): Dense, 0-indexed index giving the order that this
      field textually appears within its message in the .proto file.
    number (int): Tag number declared for this field in the .proto file.

    type (int): (One of the TYPE_* constants below) Declared type.
    cpp_type (int): (One of the CPPTYPE_* constants below) C++ type used to
      represent this field.

    label (int): (One of the LABEL_* constants below) Tells whether this
      field is optional, required, or repeated.
    has_default_value (bool): True if this field has a default value defined,
      otherwise false.
    default_value (Varies): Default value of this field.  Only
      meaningful for non-repeated scalar fields.  Repeated fields
      should always set this to [], and non-repeated composite
      fields should always set this to None.

    containing_type (Descriptor): Descriptor of the protocol message
      type that contains this field.  Set by the Descriptor constructor
      if we're passed into one.
      Somewhat confusingly, for extension fields, this is the
      descriptor of the EXTENDED message, not the descriptor
      of the message containing this field.  (See is_extension and
      extension_scope below).
    message_type (Descriptor): If a composite field, a descriptor
      of the message type contained in this field.  Otherwise, this is None.
    enum_type (EnumDescriptor): If this field contains an enum, a
      descriptor of that enum.  Otherwise, this is None.

    is_extension: True iff this describes an extension field.
    extension_scope (Descriptor): Only meaningful if is_extension is True.
      Gives the message that immediately contains this extension field.
      Will be None iff we're a top-level (file-level) extension field.

    options (descriptor_pb2.FieldOptions): Protocol message field options or
      None to use default field options.

    containing_oneof (OneofDescriptor): If the field is a member of a oneof
      union, contains its descriptor. Otherwise, None.

    file (FileDescriptor): Reference to file descriptor.
  """

  # Must be consistent with C++ FieldDescriptor::Type enum in
  # descriptor.h.
  #
  # TODO(robinson): Find a way to eliminate this repetition.
  TYPE_DOUBLE         = 1
  TYPE_FLOAT          = 2
  TYPE_INT64          = 3
  TYPE_UINT64         = 4
  TYPE_INT32          = 5
  TYPE_FIXED64        = 6
  TYPE_FIXED32        = 7
  TYPE_BOOL           = 8
  TYPE_STRING         = 9
  TYPE_GROUP          = 10
  TYPE_MESSAGE        = 11
  TYPE_BYTES          = 12
  TYPE_UINT32         = 13
  TYPE_ENUM           = 14
  TYPE_SFIXED32       = 15
  TYPE_SFIXED64       = 16
  TYPE_SINT32         = 17
  TYPE_SINT64         = 18
  MAX_TYPE            = 18

  # Must be consistent with C++ FieldDescriptor::CppType enum in
  # descriptor.h.
  #
  # TODO(robinson): Find a way to eliminate this repetition.
  CPPTYPE_INT32       = 1
  CPPTYPE_INT64       = 2
  CPPTYPE_UINT32      = 3
  CPPTYPE_UINT64      = 4
  CPPTYPE_DOUBLE      = 5
  CPPTYPE_FLOAT       = 6
  CPPTYPE_BOOL        = 7
  CPPTYPE_ENUM        = 8
  CPPTYPE_STRING      = 9
  CPPTYPE_MESSAGE     = 10
  MAX_CPPTYPE         = 10

  _PYTHON_TO_CPP_PROTO_TYPE_MAP = {
      TYPE_DOUBLE: CPPTYPE_DOUBLE,
      TYPE_FLOAT: CPPTYPE_FLOAT,
      TYPE_ENUM: CPPTYPE_ENUM,
      TYPE_INT64: CPPTYPE_INT64,
      TYPE_SINT64: CPPTYPE_INT64,
      TYPE_SFIXED64: CPPTYPE_INT64,
      TYPE_UINT64: CPPTYPE_UINT64,
      TYPE_FIXED64: CPPTYPE_UINT64,
      TYPE_INT32: CPPTYPE_INT32,
      TYPE_SFIXED32: CPPTYPE_INT32,
      TYPE_SINT32: CPPTYPE_INT32,
      TYPE_UINT32: CPPTYPE_UINT32,
      TYPE_FIXED32: CPPTYPE_UINT32,
      TYPE_BYTES: CPPTYPE_STRING,
      TYPE_STRING: CPPTYPE_STRING,
      TYPE_BOOL: CPPTYPE_BOOL,
      TYPE_MESSAGE: CPPTYPE_MESSAGE,
      TYPE_GROUP: CPPTYPE_MESSAGE
      }

  # Must be consistent with C++ FieldDescriptor::Label enum in
  # descriptor.h.
  #
  # TODO(robinson): Find a way to eliminate this repetition.
  LABEL_OPTIONAL      = 1
  LABEL_REQUIRED      = 2
  LABEL_REPEATED      = 3
  MAX_LABEL           = 3

  # Must be consistent with C++ constants kMaxNumber, kFirstReservedNumber,
  # and kLastReservedNumber in descriptor.h
  MAX_FIELD_NUMBER = (1 << 29) - 1
  FIRST_RESERVED_FIELD_NUMBER = 19000
  LAST_RESERVED_FIELD_NUMBER = 19999

  if _USE_C_DESCRIPTORS:
    _C_DESCRIPTOR_CLASS = _message.FieldDescriptor

    def __new__(cls, name, full_name, index, number, type, cpp_type, label,
                default_value, message_type, enum_type, containing_type,
                is_extension, extension_scope, options=None,
                serialized_options=None,
                has_default_value=True, containing_oneof=None, json_name=None,
                file=None, create_key=None):  # pylint: disable=redefined-builtin
      _message.Message._CheckCalledFromGeneratedFile()
      if is_extension:
        return _message.default_pool.FindExtensionByName(full_name)
      else:
        return _message.default_pool.FindFieldByName(full_name)

  def __init__(self, name, full_name, index, number, type, cpp_type, label,
               default_value, message_type, enum_type, containing_type,
               is_extension, extension_scope, options=None,
               serialized_options=None,
               has_default_value=True, containing_oneof=None, json_name=None,
               file=None, create_key=None):  # pylint: disable=redefined-builtin
    """The arguments are as described in the description of FieldDescriptor
    attributes above.

    Note that containing_type may be None, and may be set later if necessary
    (to deal with circular references between message types, for example).
    Likewise for extension_scope.
    """
    if create_key is not _internal_create_key:
      _Deprecated('FieldDescriptor')

    super(FieldDescriptor, self).__init__(
        options, serialized_options, 'FieldOptions')
    self.name = name
    self.full_name = full_name
    self.file = file
    self._camelcase_name = None
    if json_name is None:
      self.json_name = _ToJsonName(name)
    else:
      self.json_name = json_name
    self.index = index
    self.number = number
    self.type = type
    self.cpp_type = cpp_type
    self.label = label
    self.has_default_value = has_default_value
    self.default_value = default_value
    self.containing_type = containing_type
    self.message_type = message_type
    self.enum_type = enum_type
    self.is_extension = is_extension
    self.extension_scope = extension_scope
    self.containing_oneof = containing_oneof
    if api_implementation.Type() == 'cpp':
      if is_extension:
        self._cdescriptor = _message.default_pool.FindExtensionByName(full_name)
      else:
        self._cdescriptor = _message.default_pool.FindFieldByName(full_name)
    else:
      self._cdescriptor = None

  @property
  def camelcase_name(self):
    """Camelcase name of this field.

    Returns:
      str: the name in CamelCase.
    """
    if self._camelcase_name is None:
      self._camelcase_name = _ToCamelCase(self.name)
    return self._camelcase_name

  @staticmethod
  def ProtoTypeToCppProtoType(proto_type):
    """Converts from a Python proto type to a C++ Proto Type.

    The Python ProtocolBuffer classes specify both the 'Python' datatype and the
    'C++' datatype - and they're not the same. This helper method should
    translate from one to another.

    Args:
      proto_type: the Python proto type (descriptor.FieldDescriptor.TYPE_*)
    Returns:
      int: descriptor.FieldDescriptor.CPPTYPE_*, the C++ type.
    Raises:
      TypeTransformationError: when the Python proto type isn't known.
    """
    try:
      return FieldDescriptor._PYTHON_TO_CPP_PROTO_TYPE_MAP[proto_type]
    except KeyError:
      raise TypeTransformationError('Unknown proto_type: %s' % proto_type)


class EnumDescriptor(_NestedDescriptorBase):

  """Descriptor for an enum defined in a .proto file.

  Attributes:
    name (str): Name of the enum type.
    full_name (str): Full name of the type, including package name
      and any enclosing type(s).

    values (list[EnumValueDescriptors]): List of the values
      in this enum.
    values_by_name (dict(str, EnumValueDescriptor)): Same as :attr:`values`,
      but indexed by the "name" field of each EnumValueDescriptor.
    values_by_number (dict(int, EnumValueDescriptor)): Same as :attr:`values`,
      but indexed by the "number" field of each EnumValueDescriptor.
    containing_type (Descriptor): Descriptor of the immediate containing
      type of this enum, or None if this is an enum defined at the
      top level in a .proto file.  Set by Descriptor's constructor
      if we're passed into one.
    file (FileDescriptor): Reference to file descriptor.
    options (descriptor_pb2.EnumOptions): Enum options message or
      None to use default enum options.
  """

  if _USE_C_DESCRIPTORS:
    _C_DESCRIPTOR_CLASS = _message.EnumDescriptor

    def __new__(cls, name, full_name, filename, values,
                containing_type=None, options=None,
                serialized_options=None, file=None,  # pylint: disable=redefined-builtin
                serialized_start=None, serialized_end=None, create_key=None):
      _message.Message._CheckCalledFromGeneratedFile()
      return _message.default_pool.FindEnumTypeByName(full_name)

  def __init__(self, name, full_name, filename, values,
               containing_type=None, options=None,
               serialized_options=None, file=None,  # pylint: disable=redefined-builtin
               serialized_start=None, serialized_end=None, create_key=None):
    """Arguments are as described in the attribute description above.

    Note that filename is an obsolete argument, that is not used anymore.
    Please use file.name to access this as an attribute.
    """
    if create_key is not _internal_create_key:
      _Deprecated('EnumDescriptor')

    super(EnumDescriptor, self).__init__(
        options, 'EnumOptions', name, full_name, file,
        containing_type, serialized_start=serialized_start,
        serialized_end=serialized_end, serialized_options=serialized_options)

    self.values = values
    for value in self.values:
      value.type = self
    self.values_by_name = dict((v.name, v) for v in values)
    # Values are reversed to ensure that the first alias is retained.
    self.values_by_number = dict((v.number, v) for v in reversed(values))

  def CopyToProto(self, proto):
    """Copies this to a descriptor_pb2.EnumDescriptorProto.

    Args:
      proto (descriptor_pb2.EnumDescriptorProto): An empty descriptor proto.
    """
    # This function is overridden to give a better doc comment.
    super(EnumDescriptor, self).CopyToProto(proto)


class EnumValueDescriptor(DescriptorBase):

  """Descriptor for a single value within an enum.

  Attributes:
    name (str): Name of this value.
    index (int): Dense, 0-indexed index giving the order that this
      value appears textually within its enum in the .proto file.
    number (int): Actual number assigned to this enum value.
    type (EnumDescriptor): :class:`EnumDescriptor` to which this value
      belongs.  Set by :class:`EnumDescriptor`'s constructor if we're
      passed into one.
    options (descriptor_pb2.EnumValueOptions): Enum value options message or
      None to use default enum value options options.
  """

  if _USE_C_DESCRIPTORS:
    _C_DESCRIPTOR_CLASS = _message.EnumValueDescriptor

    def __new__(cls, name, index, number,
                type=None,  # pylint: disable=redefined-builtin
                options=None, serialized_options=None, create_key=None):
      _message.Message._CheckCalledFromGeneratedFile()
      # There is no way we can build a complete EnumValueDescriptor with the
      # given parameters (the name of the Enum is not known, for example).
      # Fortunately generated files just pass it to the EnumDescriptor()
      # constructor, which will ignore it, so returning None is good enough.
      return None

  def __init__(self, name, index, number,
               type=None,  # pylint: disable=redefined-builtin
               options=None, serialized_options=None, create_key=None):
    """Arguments are as described in the attribute description above."""
    if create_key is not _internal_create_key:
      _Deprecated('EnumValueDescriptor')

    super(EnumValueDescriptor, self).__init__(
        options, serialized_options, 'EnumValueOptions')
    self.name = name
    self.index = index
    self.number = number
    self.type = type


class OneofDescriptor(DescriptorBase):
  """Descriptor for a oneof field.

  Attributes:
    name (str): Name of the oneof field.
    full_name (str): Full name of the oneof field, including package name.
    index (int): 0-based index giving the order of the oneof field inside
      its containing type.
    containing_type (Descriptor): :class:`Descriptor` of the protocol message
      type that contains this field.  Set by the :class:`Descriptor` constructor
      if we're passed into one.
    fields (list[FieldDescriptor]): The list of field descriptors this
      oneof can contain.
  """

  if _USE_C_DESCRIPTORS:
    _C_DESCRIPTOR_CLASS = _message.OneofDescriptor

    def __new__(
        cls, name, full_name, index, containing_type, fields, options=None,
        serialized_options=None, create_key=None):
      _message.Message._CheckCalledFromGeneratedFile()
      return _message.default_pool.FindOneofByName(full_name)

  def __init__(
      self, name, full_name, index, containing_type, fields, options=None,
      serialized_options=None, create_key=None):
    """Arguments are as described in the attribute description above."""
    if create_key is not _internal_create_key:
      _Deprecated('OneofDescriptor')

    super(OneofDescriptor, self).__init__(
        options, serialized_options, 'OneofOptions')
    self.name = name
    self.full_name = full_name
    self.index = index
    self.containing_type = containing_type
    self.fields = fields


class ServiceDescriptor(_NestedDescriptorBase):

  """Descriptor for a service.

  Attributes:
    name (str): Name of the service.
    full_name (str): Full name of the service, including package name.
    index (int): 0-indexed index giving the order that this services
      definition appears within the .proto file.
    methods (list[MethodDescriptor]): List of methods provided by this
      service.
    methods_by_name (dict(str, MethodDescriptor)): Same
      :class:`MethodDescriptor` objects as in :attr:`methods_by_name`, but
      indexed by "name" attribute in each :class:`MethodDescriptor`.
    options (descriptor_pb2.ServiceOptions): Service options message or
      None to use default service options.
    file (FileDescriptor): Reference to file info.
  """

  if _USE_C_DESCRIPTORS:
    _C_DESCRIPTOR_CLASS = _message.ServiceDescriptor

    def __new__(
        cls,
        name=None,
        full_name=None,
        index=None,
        methods=None,
        options=None,
        serialized_options=None,
        file=None,  # pylint: disable=redefined-builtin
        serialized_start=None,
        serialized_end=None,
        create_key=None):
      _message.Message._CheckCalledFromGeneratedFile()  # pylint: disable=protected-access
      return _message.default_pool.FindServiceByName(full_name)

  def __init__(self, name, full_name, index, methods, options=None,
               serialized_options=None, file=None,  # pylint: disable=redefined-builtin
               serialized_start=None, serialized_end=None, create_key=None):
    if create_key is not _internal_create_key:
      _Deprecated('ServiceDescriptor')

    super(ServiceDescriptor, self).__init__(
        options, 'ServiceOptions', name, full_name, file,
        None, serialized_start=serialized_start,
        serialized_end=serialized_end, serialized_options=serialized_options)
    self.index = index
    self.methods = methods
    self.methods_by_name = dict((m.name, m) for m in methods)
    # Set the containing service for each method in this service.
    for method in self.methods:
      method.containing_service = self

  def FindMethodByName(self, name):
    """Searches for the specified method, and returns its descriptor.

    Args:
      name (str): Name of the method.
    Returns:
      MethodDescriptor or None: the descriptor for the requested method, if
      found.
    """
    return self.methods_by_name.get(name, None)

  def CopyToProto(self, proto):
    """Copies this to a descriptor_pb2.ServiceDescriptorProto.

    Args:
      proto (descriptor_pb2.ServiceDescriptorProto): An empty descriptor proto.
    """
    # This function is overridden to give a better doc comment.
    super(ServiceDescriptor, self).CopyToProto(proto)


class MethodDescriptor(DescriptorBase):

  """Descriptor for a method in a service.

  Attributes:
    name (str): Name of the method within the service.
    full_name (str): Full name of method.
    index (int): 0-indexed index of the method inside the service.
    containing_service (ServiceDescriptor): The service that contains this
      method.
    input_type (Descriptor): The descriptor of the message that this method
      accepts.
    output_type (Descriptor): The descriptor of the message that this method
      returns.
    options (descriptor_pb2.MethodOptions or None): Method options message, or
      None to use default method options.
  """

  if _USE_C_DESCRIPTORS:
    _C_DESCRIPTOR_CLASS = _message.MethodDescriptor

    def __new__(cls, name, full_name, index, containing_service,
                input_type, output_type, options=None, serialized_options=None,
                create_key=None):
      _message.Message._CheckCalledFromGeneratedFile()  # pylint: disable=protected-access
      return _message.default_pool.FindMethodByName(full_name)

  def __init__(self, name, full_name, index, containing_service,
               input_type, output_type, options=None, serialized_options=None,
               create_key=None):
    """The arguments are as described in the description of MethodDescriptor
    attributes above.

    Note that containing_service may be None, and may be set later if necessary.
    """
    if create_key is not _internal_create_key:
      _Deprecated('MethodDescriptor')

    super(MethodDescriptor, self).__init__(
        options, serialized_options, 'MethodOptions')
    self.name = name
    self.full_name = full_name
    self.index = index
    self.containing_service = containing_service
    self.input_type = input_type
    self.output_type = output_type

  def CopyToProto(self, proto):
    """Copies this to a descriptor_pb2.MethodDescriptorProto.

    Args:
      proto (descriptor_pb2.MethodDescriptorProto): An empty descriptor proto.

    Raises:
      Error: If self couldn't be serialized, due to too few constructor
        arguments.
    """
    if self.containing_service is not None:
      from google.protobuf import descriptor_pb2
      service_proto = descriptor_pb2.ServiceDescriptorProto()
      self.containing_service.CopyToProto(service_proto)
      proto.CopyFrom(service_proto.method[self.index])
    else:
      raise Error('Descriptor does not contain a service.')


class FileDescriptor(DescriptorBase):
  """Descriptor for a file. Mimics the descriptor_pb2.FileDescriptorProto.

  Note that :attr:`enum_types_by_name`, :attr:`extensions_by_name`, and
  :attr:`dependencies` fields are only set by the
  :py:mod:`google.protobuf.message_factory` module, and not by the generated
  proto code.

  Attributes:
    name (str): Name of file, relative to root of source tree.
    package (str): Name of the package
    syntax (str): string indicating syntax of the file (can be "proto2" or
      "proto3")
    serialized_pb (bytes): Byte string of serialized
      :class:`descriptor_pb2.FileDescriptorProto`.
    dependencies (list[FileDescriptor]): List of other :class:`FileDescriptor`
      objects this :class:`FileDescriptor` depends on.
    public_dependencies (list[FileDescriptor]): A subset of
      :attr:`dependencies`, which were declared as "public".
    message_types_by_name (dict(str, Descriptor)): Mapping from message names
      to their :class:`Descriptor`.
    enum_types_by_name (dict(str, EnumDescriptor)): Mapping from enum names to
      their :class:`EnumDescriptor`.
    extensions_by_name (dict(str, FieldDescriptor)): Mapping from extension
      names declared at file scope to their :class:`FieldDescriptor`.
    services_by_name (dict(str, ServiceDescriptor)): Mapping from services'
      names to their :class:`ServiceDescriptor`.
    pool (DescriptorPool): The pool this descriptor belongs to.  When not
      passed to the constructor, the global default pool is used.
  """

  if _USE_C_DESCRIPTORS:
    _C_DESCRIPTOR_CLASS = _message.FileDescriptor

    def __new__(cls, name, package, options=None,
                serialized_options=None, serialized_pb=None,
                dependencies=None, public_dependencies=None,
                syntax=None, pool=None, create_key=None):
      # FileDescriptor() is called from various places, not only from generated
      # files, to register dynamic proto files and messages.
      # pylint: disable=g-explicit-bool-comparison
      if serialized_pb == b'':
        # Cpp generated code must be linked in if serialized_pb is ''
        try:
          return _message.default_pool.FindFileByName(name)
        except KeyError:
          raise RuntimeError('Please link in cpp generated lib for %s' % (name))
      elif serialized_pb:
        return _message.default_pool.AddSerializedFile(serialized_pb)
      else:
        return super(FileDescriptor, cls).__new__(cls)

  def __init__(self, name, package, options=None,
               serialized_options=None, serialized_pb=None,
               dependencies=None, public_dependencies=None,
               syntax=None, pool=None, create_key=None):
    """Constructor."""
    if create_key is not _internal_create_key:
      _Deprecated('FileDescriptor')

    super(FileDescriptor, self).__init__(
        options, serialized_options, 'FileOptions')

    if pool is None:
      from google.protobuf import descriptor_pool
      pool = descriptor_pool.Default()
    self.pool = pool
    self.message_types_by_name = {}
    self.name = name
    self.package = package
    self.syntax = syntax or "proto2"
    self.serialized_pb = serialized_pb

    self.enum_types_by_name = {}
    self.extensions_by_name = {}
    self.services_by_name = {}
    self.dependencies = (dependencies or [])
    self.public_dependencies = (public_dependencies or [])

  def CopyToProto(self, proto):
    """Copies this to a descriptor_pb2.FileDescriptorProto.

    Args:
      proto: An empty descriptor_pb2.FileDescriptorProto.
    """
    proto.ParseFromString(self.serialized_pb)


def _ParseOptions(message, string):
  """Parses serialized options.

  This helper function is used to parse serialized options in generated
  proto2 files. It must not be used outside proto2.
  """
  message.ParseFromString(string)
  return message


def _ToCamelCase(name):
  """Converts name to camel-case and returns it."""
  capitalize_next = False
  result = []

  for c in name:
    if c == '_':
      if result:
        capitalize_next = True
    elif capitalize_next:
      result.append(c.upper())
      capitalize_next = False
    else:
      result += c

  # Lower-case the first letter.
  if result and result[0].isupper():
    result[0] = result[0].lower()
  return ''.join(result)


def _OptionsOrNone(descriptor_proto):
  """Returns the value of the field `options`, or None if it is not set."""
  if descriptor_proto.HasField('options'):
    return descriptor_proto.options
  else:
    return None


def _ToJsonName(name):
  """Converts name to Json name and returns it."""
  capitalize_next = False
  result = []

  for c in name:
    if c == '_':
      capitalize_next = True
    elif capitalize_next:
      result.append(c.upper())
      capitalize_next = False
    else:
      result += c

  return ''.join(result)


def MakeDescriptor(desc_proto, package='', build_file_if_cpp=True,
                   syntax=None):
  """Make a protobuf Descriptor given a DescriptorProto protobuf.

  Handles nested descriptors. Note that this is limited to the scope of defining
  a message inside of another message. Composite fields can currently only be
  resolved if the message is defined in the same scope as the field.

  Args:
    desc_proto: The descriptor_pb2.DescriptorProto protobuf message.
    package: Optional package name for the new message Descriptor (string).
    build_file_if_cpp: Update the C++ descriptor pool if api matches.
                       Set to False on recursion, so no duplicates are created.
    syntax: The syntax/semantics that should be used.  Set to "proto3" to get
            proto3 field presence semantics.
  Returns:
    A Descriptor for protobuf messages.
  """
  if api_implementation.Type() == 'cpp' and build_file_if_cpp:
    # The C++ implementation requires all descriptors to be backed by the same
    # definition in the C++ descriptor pool. To do this, we build a
    # FileDescriptorProto with the same definition as this descriptor and build
    # it into the pool.
    from google.protobuf import descriptor_pb2
    file_descriptor_proto = descriptor_pb2.FileDescriptorProto()
    file_descriptor_proto.message_type.add().MergeFrom(desc_proto)

    # Generate a random name for this proto file to prevent conflicts with any
    # imported ones. We need to specify a file name so the descriptor pool
    # accepts our FileDescriptorProto, but it is not important what that file
    # name is actually set to.
    proto_name = binascii.hexlify(os.urandom(16)).decode('ascii')

    if package:
      file_descriptor_proto.name = os.path.join(package.replace('.', '/'),
                                                proto_name + '.proto')
      file_descriptor_proto.package = package
    else:
      file_descriptor_proto.name = proto_name + '.proto'

    _message.default_pool.Add(file_descriptor_proto)
    result = _message.default_pool.FindFileByName(file_descriptor_proto.name)

    if _USE_C_DESCRIPTORS:
      return result.message_types_by_name[desc_proto.name]

  full_message_name = [desc_proto.name]
  if package: full_message_name.insert(0, package)

  # Create Descriptors for enum types
  enum_types = {}
  for enum_proto in desc_proto.enum_type:
    full_name = '.'.join(full_message_name + [enum_proto.name])
    enum_desc = EnumDescriptor(
        enum_proto.name, full_name, None, [
            EnumValueDescriptor(enum_val.name, ii, enum_val.number,
                                create_key=_internal_create_key)
            for ii, enum_val in enumerate(enum_proto.value)],
        create_key=_internal_create_key)
    enum_types[full_name] = enum_desc

  # Create Descriptors for nested types
  nested_types = {}
  for nested_proto in desc_proto.nested_type:
    full_name = '.'.join(full_message_name + [nested_proto.name])
    # Nested types are just those defined inside of the message, not all types
    # used by fields in the message, so no loops are possible here.
    nested_desc = MakeDescriptor(nested_proto,
                                 package='.'.join(full_message_name),
                                 build_file_if_cpp=False,
                                 syntax=syntax)
    nested_types[full_name] = nested_desc

  fields = []
  for field_proto in desc_proto.field:
    full_name = '.'.join(full_message_name + [field_proto.name])
    enum_desc = None
    nested_desc = None
    if field_proto.json_name:
      json_name = field_proto.json_name
    else:
      json_name = None
    if field_proto.HasField('type_name'):
      type_name = field_proto.type_name
      full_type_name = '.'.join(full_message_name +
                                [type_name[type_name.rfind('.')+1:]])
      if full_type_name in nested_types:
        nested_desc = nested_types[full_type_name]
      elif full_type_name in enum_types:
        enum_desc = enum_types[full_type_name]
      # Else type_name references a non-local type, which isn't implemented
    field = FieldDescriptor(
        field_proto.name, full_name, field_proto.number - 1,
        field_proto.number, field_proto.type,
        FieldDescriptor.ProtoTypeToCppProtoType(field_proto.type),
        field_proto.label, None, nested_desc, enum_desc, None, False, None,
        options=_OptionsOrNone(field_proto), has_default_value=False,
        json_name=json_name, create_key=_internal_create_key)
    fields.append(field)

  desc_name = '.'.join(full_message_name)
  return Descriptor(desc_proto.name, desc_name, None, None, fields,
                    list(nested_types.values()), list(enum_types.values()), [],
                    options=_OptionsOrNone(desc_proto),
                    create_key=_internal_create_key)

social_vdgnn/lib/python3.6/site-packages/google/protobuf/descriptor.py
eliatrevisan/IA-MPPI-LBM
# Protocol Buffers - Google's data interchange format
# Copyright 2008 Google Inc.  All rights reserved.
# https://developers.google.com/protocol-buffers/
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Provides DescriptorPool to use as a container for proto2 descriptors.

The DescriptorPool is used in conjection with a DescriptorDatabase to maintain
a collection of protocol buffer descriptors for use when dynamically creating
message types at runtime.

For most applications protocol buffers should be used via modules generated by
the protocol buffer compiler tool. This should only be used when the type of
protocol buffers used in an application or library cannot be predetermined.

Below is a straightforward example on how to use this class::

  pool = DescriptorPool()
  file_descriptor_protos = [ ... ]
  for file_descriptor_proto in file_descriptor_protos:
    pool.Add(file_descriptor_proto)
  my_message_descriptor = pool.FindMessageTypeByName('some.package.MessageType')

The message descriptor can be used in conjunction with the message_factory
module in order to create a protocol buffer class that can be encoded and
decoded.

If you want to get a Python class for the specified proto, use the
helper functions inside google.protobuf.message_factory
directly instead of this class.
"""

__author__ = 'matthewtoia@google.com (Matt Toia)'

import collections
import warnings

from google.protobuf import descriptor
from google.protobuf import descriptor_database
from google.protobuf import text_encoding


_USE_C_DESCRIPTORS = descriptor._USE_C_DESCRIPTORS  # pylint: disable=protected-access


def _Deprecated(func):
  """Mark functions as deprecated."""

  def NewFunc(*args, **kwargs):
    warnings.warn(
        'Call to deprecated function %s(). Note: Do add unlinked descriptors '
        'to descriptor_pool is wrong. Use Add() or AddSerializedFile() '
        'instead.' % func.__name__,
        category=DeprecationWarning)
    return func(*args, **kwargs)
  NewFunc.__name__ = func.__name__
  NewFunc.__doc__ = func.__doc__
  NewFunc.__dict__.update(func.__dict__)
  return NewFunc


def _NormalizeFullyQualifiedName(name):
  """Remove leading period from fully-qualified type name.

  Due to b/13860351 in descriptor_database.py, types in the root namespace are
  generated with a leading period. This function removes that prefix.

  Args:
    name (str): The fully-qualified symbol name.

  Returns:
    str: The normalized fully-qualified symbol name.
  """
  return name.lstrip('.')


def _OptionsOrNone(descriptor_proto):
  """Returns the value of the field `options`, or None if it is not set."""
  if descriptor_proto.HasField('options'):
    return descriptor_proto.options
  else:
    return None


def _IsMessageSetExtension(field):
  return (field.is_extension and
          field.containing_type.has_options and
          field.containing_type.GetOptions().message_set_wire_format and
          field.type == descriptor.FieldDescriptor.TYPE_MESSAGE and
          field.label == descriptor.FieldDescriptor.LABEL_OPTIONAL)


class DescriptorPool(object):
  """A collection of protobufs dynamically constructed by descriptor protos."""

  if _USE_C_DESCRIPTORS:

    def __new__(cls, descriptor_db=None):
      # pylint: disable=protected-access
      return descriptor._message.DescriptorPool(descriptor_db)

  def __init__(self, descriptor_db=None):
    """Initializes a Pool of proto buffs.

    The descriptor_db argument to the constructor is provided to allow
    specialized file descriptor proto lookup code to be triggered on demand. An
    example would be an implementation which will read and compile a file
    specified in a call to FindFileByName() and not require the call to Add()
    at all. Results from this database will be cached internally here as well.

    Args:
      descriptor_db: A secondary source of file descriptors.
    """

    self._internal_db = descriptor_database.DescriptorDatabase()
    self._descriptor_db = descriptor_db
    self._descriptors = {}
    self._enum_descriptors = {}
    self._service_descriptors = {}
    self._file_descriptors = {}
    self._toplevel_extensions = {}
    # TODO(jieluo): Remove _file_desc_by_toplevel_extension after
    # maybe year 2020 for compatibility issue (with 3.4.1 only).
    self._file_desc_by_toplevel_extension = {}
    self._top_enum_values = {}
    # We store extensions in two two-level mappings: The first key is the
    # descriptor of the message being extended, the second key is the extension
    # full name or its tag number.
    self._extensions_by_name = collections.defaultdict(dict)
    self._extensions_by_number = collections.defaultdict(dict)

  def _CheckConflictRegister(self, desc, desc_name, file_name):
    """Check if the descriptor name conflicts with another of the same name.

    Args:
      desc: Descriptor of a message, enum, service, extension or enum value.
      desc_name (str): the full name of desc.
      file_name (str): The file name of descriptor.
    """
    for register, descriptor_type in [
        (self._descriptors, descriptor.Descriptor),
        (self._enum_descriptors, descriptor.EnumDescriptor),
        (self._service_descriptors, descriptor.ServiceDescriptor),
        (self._toplevel_extensions, descriptor.FieldDescriptor),
        (self._top_enum_values, descriptor.EnumValueDescriptor)]:
      if desc_name in register:
        old_desc = register[desc_name]
        if isinstance(old_desc, descriptor.EnumValueDescriptor):
          old_file = old_desc.type.file.name
        else:
          old_file = old_desc.file.name

        if not isinstance(desc, descriptor_type) or (
            old_file != file_name):
          error_msg = ('Conflict register for file "' + file_name +
                       '": ' + desc_name +
                       ' is already defined in file "' +
                       old_file + '". Please fix the conflict by adding '
                       'package name on the proto file, or use different '
                       'name for the duplication.')
          if isinstance(desc, descriptor.EnumValueDescriptor):
            error_msg += ('\nNote: enum values appear as '
                          'siblings of the enum type instead of '
                          'children of it.')

          raise TypeError(error_msg)

        return

  def Add(self, file_desc_proto):
    """Adds the FileDescriptorProto and its types to this pool.

    Args:
      file_desc_proto (FileDescriptorProto): The file descriptor to add.
    """

    self._internal_db.Add(file_desc_proto)

  def AddSerializedFile(self, serialized_file_desc_proto):
    """Adds the FileDescriptorProto and its types to this pool.

    Args:
      serialized_file_desc_proto (bytes): A bytes string, serialization of the
        :class:`FileDescriptorProto` to add.

    Returns:
      FileDescriptor: Descriptor for the added file.
    """

    # pylint: disable=g-import-not-at-top
    from google.protobuf import descriptor_pb2
    file_desc_proto = descriptor_pb2.FileDescriptorProto.FromString(
        serialized_file_desc_proto)
    file_desc = self._ConvertFileProtoToFileDescriptor(file_desc_proto)
    file_desc.serialized_pb = serialized_file_desc_proto
    return file_desc

  # Add Descriptor to descriptor pool is dreprecated. Please use Add()
  # or AddSerializedFile() to add a FileDescriptorProto instead.
  @_Deprecated
  def AddDescriptor(self, desc):
    self._AddDescriptor(desc)

  # Never call this method. It is for internal usage only.
  def _AddDescriptor(self, desc):
    """Adds a Descriptor to the pool, non-recursively.

    If the Descriptor contains nested messages or enums, the caller must
    explicitly register them. This method also registers the FileDescriptor
    associated with the message.

    Args:
      desc: A Descriptor.
    """
    if not isinstance(desc, descriptor.Descriptor):
      raise TypeError('Expected instance of descriptor.Descriptor.')

    self._CheckConflictRegister(desc, desc.full_name, desc.file.name)

    self._descriptors[desc.full_name] = desc
    self._AddFileDescriptor(desc.file)

  # Add EnumDescriptor to descriptor pool is dreprecated. Please use Add()
  # or AddSerializedFile() to add a FileDescriptorProto instead.
  @_Deprecated
  def AddEnumDescriptor(self, enum_desc):
    self._AddEnumDescriptor(enum_desc)

  # Never call this method. It is for internal usage only.
  def _AddEnumDescriptor(self, enum_desc):
    """Adds an EnumDescriptor to the pool.

    This method also registers the FileDescriptor associated with the enum.

    Args:
      enum_desc: An EnumDescriptor.
    """

    if not isinstance(enum_desc, descriptor.EnumDescriptor):
      raise TypeError('Expected instance of descriptor.EnumDescriptor.')

    file_name = enum_desc.file.name
    self._CheckConflictRegister(enum_desc, enum_desc.full_name, file_name)
    self._enum_descriptors[enum_desc.full_name] = enum_desc

    # Top enum values need to be indexed.
    # Count the number of dots to see whether the enum is toplevel or nested
    # in a message. We cannot use enum_desc.containing_type at this stage.
    if enum_desc.file.package:
      top_level = (enum_desc.full_name.count('.')
                   - enum_desc.file.package.count('.') == 1)
    else:
      top_level = enum_desc.full_name.count('.') == 0
    if top_level:
      file_name = enum_desc.file.name
      package = enum_desc.file.package
      for enum_value in enum_desc.values:
        full_name = _NormalizeFullyQualifiedName(
            '.'.join((package, enum_value.name)))
        self._CheckConflictRegister(enum_value, full_name, file_name)
        self._top_enum_values[full_name] = enum_value
    self._AddFileDescriptor(enum_desc.file)

  # Add ServiceDescriptor to descriptor pool is dreprecated. Please use Add()
  # or AddSerializedFile() to add a FileDescriptorProto instead.
  @_Deprecated
  def AddServiceDescriptor(self, service_desc):
    self._AddServiceDescriptor(service_desc)

  # Never call this method. It is for internal usage only.
  def _AddServiceDescriptor(self, service_desc):
    """Adds a ServiceDescriptor to the pool.

    Args:
      service_desc: A ServiceDescriptor.
    """

    if not isinstance(service_desc, descriptor.ServiceDescriptor):
      raise TypeError('Expected instance of descriptor.ServiceDescriptor.')

    self._CheckConflictRegister(service_desc, service_desc.full_name,
                                service_desc.file.name)
    self._service_descriptors[service_desc.full_name] = service_desc

  # Add ExtensionDescriptor to descriptor pool is dreprecated. Please use Add()
  # or AddSerializedFile() to add a FileDescriptorProto instead.
  @_Deprecated
  def AddExtensionDescriptor(self, extension):
    self._AddExtensionDescriptor(extension)

  # Never call this method. It is for internal usage only.
  def _AddExtensionDescriptor(self, extension):
    """Adds a FieldDescriptor describing an extension to the pool.

    Args:
      extension: A FieldDescriptor.

    Raises:
      AssertionError: when another extension with the same number extends the
        same message.
      TypeError: when the specified extension is not a
        descriptor.FieldDescriptor.
    """
    if not (isinstance(extension, descriptor.FieldDescriptor) and
            extension.is_extension):
      raise TypeError('Expected an extension descriptor.')

    if extension.extension_scope is None:
      self._toplevel_extensions[extension.full_name] = extension

    try:
      existing_desc = self._extensions_by_number[
          extension.containing_type][extension.number]
    except KeyError:
      pass
    else:
      if extension is not existing_desc:
        raise AssertionError(
            'Extensions "%s" and "%s" both try to extend message type "%s" '
            'with field number %d.' %
            (extension.full_name, existing_desc.full_name,
             extension.containing_type.full_name, extension.number))

    self._extensions_by_number[extension.containing_type][
        extension.number] = extension
    self._extensions_by_name[extension.containing_type][
        extension.full_name] = extension

    # Also register MessageSet extensions with the type name.
    if _IsMessageSetExtension(extension):
      self._extensions_by_name[extension.containing_type][
          extension.message_type.full_name] = extension

  @_Deprecated
  def AddFileDescriptor(self, file_desc):
    self._InternalAddFileDescriptor(file_desc)

  # Never call this method. It is for internal usage only.
  def _InternalAddFileDescriptor(self, file_desc):
    """Adds a FileDescriptor to the pool, non-recursively.

    If the FileDescriptor contains messages or enums, the caller must explicitly
    register them.

    Args:
      file_desc: A FileDescriptor.
    """

    self._AddFileDescriptor(file_desc)
    # TODO(jieluo): This is a temporary solution for FieldDescriptor.file.
    # FieldDescriptor.file is added in code gen. Remove this solution after
    # maybe 2020 for compatibility reason (with 3.4.1 only).
    for extension in file_desc.extensions_by_name.values():
      self._file_desc_by_toplevel_extension[
          extension.full_name] = file_desc

  def _AddFileDescriptor(self, file_desc):
    """Adds a FileDescriptor to the pool, non-recursively.

    If the FileDescriptor contains messages or enums, the caller must explicitly
    register them.

    Args:
      file_desc: A FileDescriptor.
    """

    if not isinstance(file_desc, descriptor.FileDescriptor):
      raise TypeError('Expected instance of descriptor.FileDescriptor.')
    self._file_descriptors[file_desc.name] = file_desc

  def FindFileByName(self, file_name):
    """Gets a FileDescriptor by file name.

    Args:
      file_name (str): The path to the file to get a descriptor for.

    Returns:
      FileDescriptor: The descriptor for the named file.

    Raises:
      KeyError: if the file cannot be found in the pool.
    """

    try:
      return self._file_descriptors[file_name]
    except KeyError:
      pass

    try:
      file_proto = self._internal_db.FindFileByName(file_name)
    except KeyError as error:
      if self._descriptor_db:
        file_proto = self._descriptor_db.FindFileByName(file_name)
      else:
        raise error
    if not file_proto:
      raise KeyError('Cannot find a file named %s' % file_name)
    return self._ConvertFileProtoToFileDescriptor(file_proto)

  def FindFileContainingSymbol(self, symbol):
    """Gets the FileDescriptor for the file containing the specified symbol.

    Args:
      symbol (str): The name of the symbol to search for.

    Returns:
      FileDescriptor: Descriptor for the file that contains the specified
      symbol.

    Raises:
      KeyError: if the file cannot be found in the pool.
    """

    symbol = _NormalizeFullyQualifiedName(symbol)
    try:
      return self._InternalFindFileContainingSymbol(symbol)
    except KeyError:
      pass

    try:
      # Try fallback database. Build and find again if possible.
      self._FindFileContainingSymbolInDb(symbol)
      return self._InternalFindFileContainingSymbol(symbol)
    except KeyError:
      raise KeyError('Cannot find a file containing %s' % symbol)

  def _InternalFindFileContainingSymbol(self, symbol):
    """Gets the already built FileDescriptor containing the specified symbol.

    Args:
      symbol (str): The name of the symbol to search for.

    Returns:
      FileDescriptor: Descriptor for the file that contains the specified
      symbol.

    Raises:
      KeyError: if the file cannot be found in the pool.
    """
    try:
      return self._descriptors[symbol].file
    except KeyError:
      pass

    try:
      return self._enum_descriptors[symbol].file
    except KeyError:
      pass

    try:
      return self._service_descriptors[symbol].file
    except KeyError:
      pass

    try:
      return self._top_enum_values[symbol].type.file
    except KeyError:
      pass

    try:
      return self._file_desc_by_toplevel_extension[symbol]
    except KeyError:
      pass

    # Try fields, enum values and nested extensions inside a message.
    top_name, _, sub_name = symbol.rpartition('.')
    try:
      message = self.FindMessageTypeByName(top_name)
      assert (sub_name in message.extensions_by_name or
              sub_name in message.fields_by_name or
              sub_name in message.enum_values_by_name)
      return message.file
    except (KeyError, AssertionError):
      raise KeyError('Cannot find a file containing %s' % symbol)

  def FindMessageTypeByName(self, full_name):
    """Loads the named descriptor from the pool.

    Args:
      full_name (str): The full name of the descriptor to load.

    Returns:
      Descriptor: The descriptor for the named type.

    Raises:
      KeyError: if the message cannot be found in the pool.
    """

    full_name = _NormalizeFullyQualifiedName(full_name)
    if full_name not in self._descriptors:
      self._FindFileContainingSymbolInDb(full_name)
    return self._descriptors[full_name]

  def FindEnumTypeByName(self, full_name):
    """Loads the named enum descriptor from the pool.

    Args:
      full_name (str): The full name of the enum descriptor to load.

    Returns:
      EnumDescriptor: The enum descriptor for the named type.

    Raises:
      KeyError: if the enum cannot be found in the pool.
    """

    full_name = _NormalizeFullyQualifiedName(full_name)
    if full_name not in self._enum_descriptors:
      self._FindFileContainingSymbolInDb(full_name)
    return self._enum_descriptors[full_name]

  def FindFieldByName(self, full_name):
    """Loads the named field descriptor from the pool.

    Args:
      full_name (str): The full name of the field descriptor to load.

    Returns:
      FieldDescriptor: The field descriptor for the named field.

    Raises:
      KeyError: if the field cannot be found in the pool.
    """
    full_name = _NormalizeFullyQualifiedName(full_name)
    message_name, _, field_name = full_name.rpartition('.')
    message_descriptor = self.FindMessageTypeByName(message_name)
    return message_descriptor.fields_by_name[field_name]

  def FindOneofByName(self, full_name):
    """Loads the named oneof descriptor from the pool.

    Args:
      full_name (str): The full name of the oneof descriptor to load.

    Returns:
      OneofDescriptor: The oneof descriptor for the named oneof.

    Raises:
      KeyError: if the oneof cannot be found in the pool.
    """
    full_name = _NormalizeFullyQualifiedName(full_name)
    message_name, _, oneof_name = full_name.rpartition('.')
    message_descriptor = self.FindMessageTypeByName(message_name)
    return message_descriptor.oneofs_by_name[oneof_name]

  def FindExtensionByName(self, full_name):
    """Loads the named extension descriptor from the pool.

    Args:
      full_name (str): The full name of the extension descriptor to load.

    Returns:
      FieldDescriptor: The field descriptor for the named extension.

    Raises:
      KeyError: if the extension cannot be found in the pool.
    """
    full_name = _NormalizeFullyQualifiedName(full_name)
    try:
      # The proto compiler does not give any link between the FileDescriptor
      # and top-level extensions unless the FileDescriptorProto is added to
      # the DescriptorDatabase, but this can impact memory usage.
      # So we registered these extensions by name explicitly.
      return self._toplevel_extensions[full_name]
    except KeyError:
      pass
    message_name, _, extension_name = full_name.rpartition('.')
    try:
      # Most extensions are nested inside a message.
      scope = self.FindMessageTypeByName(message_name)
    except KeyError:
      # Some extensions are defined at file scope.
      scope = self._FindFileContainingSymbolInDb(full_name)
    return scope.extensions_by_name[extension_name]

  def FindExtensionByNumber(self, message_descriptor, number):
    """Gets the extension of the specified message with the specified number.

    Extensions have to be registered to this pool by calling :func:`Add` or
    :func:`AddExtensionDescriptor`.

    Args:
      message_descriptor (Descriptor): descriptor of the extended message.
      number (int): Number of the extension field.

    Returns:
      FieldDescriptor: The descriptor for the extension.

    Raises:
      KeyError: when no extension with the given number is known for the
        specified message.
    """
    try:
      return self._extensions_by_number[message_descriptor][number]
    except KeyError:
      self._TryLoadExtensionFromDB(message_descriptor, number)
      return self._extensions_by_number[message_descriptor][number]

  def FindAllExtensions(self, message_descriptor):
    """Gets all the known extensions of a given message.

    Extensions have to be registered to this pool by build related
    :func:`Add` or :func:`AddExtensionDescriptor`.

    Args:
      message_descriptor (Descriptor): Descriptor of the extended message.

    Returns:
      list[FieldDescriptor]: Field descriptors describing the extensions.
    """
    # Fallback to descriptor db if FindAllExtensionNumbers is provided.
    if self._descriptor_db and hasattr(
        self._descriptor_db, 'FindAllExtensionNumbers'):
      full_name = message_descriptor.full_name
      all_numbers = self._descriptor_db.FindAllExtensionNumbers(full_name)
      for number in all_numbers:
        if number in self._extensions_by_number[message_descriptor]:
          continue
        self._TryLoadExtensionFromDB(message_descriptor, number)

    return list(self._extensions_by_number[message_descriptor].values())

  def _TryLoadExtensionFromDB(self, message_descriptor, number):
    """Try to Load extensions from descriptor db.

    Args:
      message_descriptor: descriptor of the extended message.
      number: the extension number that needs to be loaded.
    """
    if not self._descriptor_db:
      return
    # Only supported when FindFileContainingExtension is provided.
    if not hasattr(
        self._descriptor_db, 'FindFileContainingExtension'):
      return

    full_name = message_descriptor.full_name
    file_proto = self._descriptor_db.FindFileContainingExtension(
        full_name, number)

    if file_proto is None:
      return

    try:
      self._ConvertFileProtoToFileDescriptor(file_proto)
    except:
      warn_msg = ('Unable to load proto file %s for extension number %d.' %
                  (file_proto.name, number))
      warnings.warn(warn_msg, RuntimeWarning)

  def FindServiceByName(self, full_name):
    """Loads the named service descriptor from the pool.

    Args:
      full_name (str): The full name of the service descriptor to load.

    Returns:
      ServiceDescriptor: The service descriptor for the named service.

    Raises:
      KeyError: if the service cannot be found in the pool.
    """
    full_name = _NormalizeFullyQualifiedName(full_name)
    if full_name not in self._service_descriptors:
      self._FindFileContainingSymbolInDb(full_name)
    return self._service_descriptors[full_name]

  def FindMethodByName(self, full_name):
    """Loads the named service method descriptor from the pool.

    Args:
      full_name (str): The full name of the method descriptor to load.

    Returns:
      MethodDescriptor: The method descriptor for the service method.

    Raises:
      KeyError: if the method cannot be found in the pool.
    """
    full_name = _NormalizeFullyQualifiedName(full_name)
    service_name, _, method_name = full_name.rpartition('.')
    service_descriptor = self.FindServiceByName(service_name)
    return service_descriptor.methods_by_name[method_name]

  def _FindFileContainingSymbolInDb(self, symbol):
    """Finds the file in descriptor DB containing the specified symbol.

    Args:
      symbol (str): The name of the symbol to search for.

    Returns:
      FileDescriptor: The file that contains the specified symbol.

    Raises:
      KeyError: if the file cannot be found in the descriptor database.
    """
    try:
      file_proto = self._internal_db.FindFileContainingSymbol(symbol)
    except KeyError as error:
      if self._descriptor_db:
        file_proto = self._descriptor_db.FindFileContainingSymbol(symbol)
      else:
        raise error
    if not file_proto:
      raise KeyError('Cannot find a file containing %s' % symbol)
    return self._ConvertFileProtoToFileDescriptor(file_proto)

  def _ConvertFileProtoToFileDescriptor(self, file_proto):
    """Creates a FileDescriptor from a proto or returns a cached copy.

    This method also has the side effect of loading all the symbols found in
    the file into the appropriate dictionaries in the pool.

    Args:
      file_proto: The proto to convert.

    Returns:
      A FileDescriptor matching the passed in proto.
    """
    if file_proto.name not in self._file_descriptors:
      built_deps = list(self._GetDeps(file_proto.dependency))
      direct_deps = [self.FindFileByName(n) for n in file_proto.dependency]
      public_deps = [direct_deps[i] for i in file_proto.public_dependency]

      file_descriptor = descriptor.FileDescriptor(
          pool=self,
          name=file_proto.name,
          package=file_proto.package,
          syntax=file_proto.syntax,
          options=_OptionsOrNone(file_proto),
          serialized_pb=file_proto.SerializeToString(),
          dependencies=direct_deps,
          public_dependencies=public_deps,
          # pylint: disable=protected-access
          create_key=descriptor._internal_create_key)
      scope = {}

      # This loop extracts all the message and enum types from all the
      # dependencies of the file_proto. This is necessary to create the
      # scope of available message types when defining the passed in
      # file proto.
      for dependency in built_deps:
        scope.update(self._ExtractSymbols(
            dependency.message_types_by_name.values()))
        scope.update((_PrefixWithDot(enum.full_name), enum)
                     for enum in dependency.enum_types_by_name.values())

      for message_type in file_proto.message_type:
        message_desc = self._ConvertMessageDescriptor(
            message_type, file_proto.package, file_descriptor, scope,
            file_proto.syntax)
        file_descriptor.message_types_by_name[message_desc.name] = (
            message_desc)

      for enum_type in file_proto.enum_type:
        file_descriptor.enum_types_by_name[enum_type.name] = (
            self._ConvertEnumDescriptor(enum_type, file_proto.package,
                                        file_descriptor, None, scope, True))

      for index, extension_proto in enumerate(file_proto.extension):
        extension_desc = self._MakeFieldDescriptor(
            extension_proto, file_proto.package, index, file_descriptor,
            is_extension=True)
        extension_desc.containing_type = self._GetTypeFromScope(
            file_descriptor.package, extension_proto.extendee, scope)
        self._SetFieldType(extension_proto, extension_desc,
                           file_descriptor.package, scope)
        file_descriptor.extensions_by_name[extension_desc.name] = (
            extension_desc)
        self._file_desc_by_toplevel_extension[extension_desc.full_name] = (
            file_descriptor)

      for desc_proto in file_proto.message_type:
        self._SetAllFieldTypes(file_proto.package, desc_proto, scope)

      if file_proto.package:
        desc_proto_prefix = _PrefixWithDot(file_proto.package)
      else:
        desc_proto_prefix = ''

      for desc_proto in file_proto.message_type:
        desc = self._GetTypeFromScope(
            desc_proto_prefix, desc_proto.name, scope)
        file_descriptor.message_types_by_name[desc_proto.name] = desc

      for index, service_proto in enumerate(file_proto.service):
        file_descriptor.services_by_name[service_proto.name] = (
            self._MakeServiceDescriptor(service_proto, index, scope,
                                        file_proto.package, file_descriptor))

      self._file_descriptors[file_proto.name] = file_descriptor

    # Add extensions to the pool
    file_desc = self._file_descriptors[file_proto.name]
    for extension in file_desc.extensions_by_name.values():
      self._AddExtensionDescriptor(extension)
    for message_type in file_desc.message_types_by_name.values():
      for extension in message_type.extensions:
        self._AddExtensionDescriptor(extension)

    return file_desc

  def _ConvertMessageDescriptor(self, desc_proto, package=None, file_desc=None,
                                scope=None, syntax=None):
    """Adds the proto to the pool in the specified package.

    Args:
      desc_proto: The descriptor_pb2.DescriptorProto protobuf message.
      package: The package the proto should be located in.
      file_desc: The file containing this message.
      scope: Dict mapping short and full symbols to message and enum types.
      syntax: string indicating syntax of the file ("proto2" or "proto3")

    Returns:
      The added descriptor.
    """

    if package:
      desc_name = '.'.join((package, desc_proto.name))
    else:
      desc_name = desc_proto.name

    if file_desc is None:
      file_name = None
    else:
      file_name = file_desc.name

    if scope is None:
      scope = {}

    nested = [
        self._ConvertMessageDescriptor(
            nested, desc_name, file_desc, scope, syntax)
        for nested in desc_proto.nested_type]
    enums = [
        self._ConvertEnumDescriptor(enum, desc_name, file_desc, None,
                                    scope, False)
        for enum in desc_proto.enum_type]
    fields = [self._MakeFieldDescriptor(field, desc_name, index, file_desc)
              for index, field in enumerate(desc_proto.field)]
    extensions = [
        self._MakeFieldDescriptor(extension, desc_name, index, file_desc,
                                  is_extension=True)
        for index, extension in enumerate(desc_proto.extension)]
    oneofs = [
        # pylint: disable=g-complex-comprehension
        descriptor.OneofDescriptor(
            desc.name,
            '.'.join((desc_name, desc.name)),
            index,
            None,
            [],
            _OptionsOrNone(desc),
            # pylint: disable=protected-access
            create_key=descriptor._internal_create_key)
        for index, desc in enumerate(desc_proto.oneof_decl)
    ]
    extension_ranges = [(r.start, r.end) for r in desc_proto.extension_range]
    if extension_ranges:
      is_extendable = True
    else:
      is_extendable = False
    desc = descriptor.Descriptor(
        name=desc_proto.name,
        full_name=desc_name,
        filename=file_name,
        containing_type=None,
        fields=fields,
        oneofs=oneofs,
        nested_types=nested,
        enum_types=enums,
        extensions=extensions,
        options=_OptionsOrNone(desc_proto),
        is_extendable=is_extendable,
        extension_ranges=extension_ranges,
        file=file_desc,
        serialized_start=None,
        serialized_end=None,
        syntax=syntax,
        # pylint: disable=protected-access
        create_key=descriptor._internal_create_key)
    for nested in desc.nested_types:
      nested.containing_type = desc
    for enum in desc.enum_types:
      enum.containing_type = desc
    for field_index, field_desc in enumerate(desc_proto.field):
      if field_desc.HasField('oneof_index'):
        oneof_index = field_desc.oneof_index
        oneofs[oneof_index].fields.append(fields[field_index])
        fields[field_index].containing_oneof = oneofs[oneof_index]

    scope[_PrefixWithDot(desc_name)] = desc
    self._CheckConflictRegister(desc, desc.full_name, desc.file.name)
    self._descriptors[desc_name] = desc
    return desc

  def _ConvertEnumDescriptor(self, enum_proto, package=None, file_desc=None,
                             containing_type=None, scope=None, top_level=False):
    """Make a protobuf EnumDescriptor given an EnumDescriptorProto protobuf.

    Args:
      enum_proto: The descriptor_pb2.EnumDescriptorProto protobuf message.
      package: Optional package name for the new message EnumDescriptor.
      file_desc: The file containing the enum descriptor.
      containing_type: The type containing this enum.
      scope: Scope containing available types.
      top_level: If True, the enum is a top level symbol. If False, the enum
          is defined inside a message.

    Returns:
      The added descriptor
    """

    if package:
      enum_name = '.'.join((package, enum_proto.name))
    else:
      enum_name = enum_proto.name

    if file_desc is None:
      file_name = None
    else:
      file_name = file_desc.name

    values = [self._MakeEnumValueDescriptor(value, index)
              for index, value in enumerate(enum_proto.value)]
    desc = descriptor.EnumDescriptor(name=enum_proto.name,
                                     full_name=enum_name,
                                     filename=file_name,
                                     file=file_desc,
                                     values=values,
                                     containing_type=containing_type,
                                     options=_OptionsOrNone(enum_proto),
                                     # pylint: disable=protected-access
                                     create_key=descriptor._internal_create_key)
    scope['.%s' % enum_name] = desc
    self._CheckConflictRegister(desc, desc.full_name, desc.file.name)
    self._enum_descriptors[enum_name] = desc

    # Add top level enum values.
    if top_level:
      for value in values:
        full_name = _NormalizeFullyQualifiedName(
            '.'.join((package, value.name)))
        self._CheckConflictRegister(value, full_name, file_name)
        self._top_enum_values[full_name] = value

    return desc

  def _MakeFieldDescriptor(self, field_proto, message_name, index,
                           file_desc, is_extension=False):
    """Creates a field descriptor from a FieldDescriptorProto.

    For message and enum type fields, this method will do a look up
    in the pool for the appropriate descriptor for that type. If it
    is unavailable, it will fall back to the _source function to
    create it. If this type is still unavailable, construction will
    fail.

    Args:
      field_proto: The proto describing the field.
      message_name: The name of the containing message.
      index: Index of the field
      file_desc: The file containing the field descriptor.
      is_extension: Indication that this field is for an extension.

    Returns:
      An initialized FieldDescriptor object
    """

    if message_name:
      full_name = '.'.join((message_name, field_proto.name))
    else:
      full_name = field_proto.name

    if field_proto.json_name:
      json_name = field_proto.json_name
    else:
      json_name = None

    return descriptor.FieldDescriptor(
        name=field_proto.name,
        full_name=full_name,
        index=index,
        number=field_proto.number,
        type=field_proto.type,
        cpp_type=None,
        message_type=None,
        enum_type=None,
        containing_type=None,
        label=field_proto.label,
        has_default_value=False,
        default_value=None,
        is_extension=is_extension,
        extension_scope=None,
        options=_OptionsOrNone(field_proto),
        json_name=json_name,
        file=file_desc,
        # pylint: disable=protected-access
        create_key=descriptor._internal_create_key)

  def _SetAllFieldTypes(self, package, desc_proto, scope):
    """Sets all the descriptor's fields's types.

    This method also sets the containing types on any extensions.

    Args:
      package: The current package of desc_proto.
      desc_proto: The message descriptor to update.
      scope: Enclosing scope of available types.
    """

    package = _PrefixWithDot(package)

    main_desc = self._GetTypeFromScope(package, desc_proto.name, scope)

    if package == '.':
      nested_package = _PrefixWithDot(desc_proto.name)
    else:
      nested_package = '.'.join([package, desc_proto.name])

    for field_proto, field_desc in zip(desc_proto.field, main_desc.fields):
      self._SetFieldType(field_proto, field_desc, nested_package, scope)

    for extension_proto, extension_desc in (
        zip(desc_proto.extension, main_desc.extensions)):
      extension_desc.containing_type = self._GetTypeFromScope(
          nested_package, extension_proto.extendee, scope)
      self._SetFieldType(extension_proto, extension_desc, nested_package, scope)

    for nested_type in desc_proto.nested_type:
      self._SetAllFieldTypes(nested_package, nested_type, scope)

  def _SetFieldType(self, field_proto, field_desc, package, scope):
    """Sets the field's type, cpp_type, message_type and enum_type.

    Args:
      field_proto: Data about the field in proto format.
      field_desc: The descriptor to modify.
      package: The package the field's container is in.
      scope: Enclosing scope of available types.
    """
    if field_proto.type_name:
      desc = self._GetTypeFromScope(package, field_proto.type_name, scope)
    else:
      desc = None

    if not field_proto.HasField('type'):
      if isinstance(desc, descriptor.Descriptor):
        field_proto.type = descriptor.FieldDescriptor.TYPE_MESSAGE
      else:
        field_proto.type = descriptor.FieldDescriptor.TYPE_ENUM

    field_desc.cpp_type = descriptor.FieldDescriptor.ProtoTypeToCppProtoType(
        field_proto.type)

    if (field_proto.type == descriptor.FieldDescriptor.TYPE_MESSAGE
        or field_proto.type == descriptor.FieldDescriptor.TYPE_GROUP):
      field_desc.message_type = desc

    if field_proto.type == descriptor.FieldDescriptor.TYPE_ENUM:
      field_desc.enum_type = desc

    if field_proto.label == descriptor.FieldDescriptor.LABEL_REPEATED:
      field_desc.has_default_value = False
      field_desc.default_value = []
    elif field_proto.HasField('default_value'):
      field_desc.has_default_value = True
      if (field_proto.type == descriptor.FieldDescriptor.TYPE_DOUBLE or
          field_proto.type == descriptor.FieldDescriptor.TYPE_FLOAT):
        field_desc.default_value = float(field_proto.default_value)
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_STRING:
        field_desc.default_value = field_proto.default_value
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_BOOL:
        field_desc.default_value = field_proto.default_value.lower() == 'true'
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_ENUM:
        field_desc.default_value = field_desc.enum_type.values_by_name[
            field_proto.default_value].number
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_BYTES:
        field_desc.default_value = text_encoding.CUnescape(
            field_proto.default_value)
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_MESSAGE:
        field_desc.default_value = None
      else:
        # All other types are of the "int" type.
        field_desc.default_value = int(field_proto.default_value)
    else:
      field_desc.has_default_value = False
      if (field_proto.type == descriptor.FieldDescriptor.TYPE_DOUBLE or
          field_proto.type == descriptor.FieldDescriptor.TYPE_FLOAT):
        field_desc.default_value = 0.0
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_STRING:
        field_desc.default_value = u''
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_BOOL:
        field_desc.default_value = False
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_ENUM:
        field_desc.default_value = field_desc.enum_type.values[0].number
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_BYTES:
        field_desc.default_value = b''
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_MESSAGE:
        field_desc.default_value = None
      elif field_proto.type == descriptor.FieldDescriptor.TYPE_GROUP:
        field_desc.default_value = None
      else:
        # All other types are of the "int" type.
        field_desc.default_value = 0

    field_desc.type = field_proto.type

  def _MakeEnumValueDescriptor(self, value_proto, index):
    """Creates a enum value descriptor object from a enum value proto.

    Args:
      value_proto: The proto describing the enum value.
      index: The index of the enum value.

    Returns:
      An initialized EnumValueDescriptor object.
    """

    return descriptor.EnumValueDescriptor(
        name=value_proto.name,
        index=index,
        number=value_proto.number,
        options=_OptionsOrNone(value_proto),
        type=None,
        # pylint: disable=protected-access
        create_key=descriptor._internal_create_key)

  def _MakeServiceDescriptor(self, service_proto, service_index, scope,
                             package, file_desc):
    """Make a protobuf ServiceDescriptor given a ServiceDescriptorProto.

    Args:
      service_proto: The descriptor_pb2.ServiceDescriptorProto protobuf message.
      service_index: The index of the service in the File.
      scope: Dict mapping short and full symbols to message and enum types.
      package: Optional package name for the new message EnumDescriptor.
      file_desc: The file containing the service descriptor.

    Returns:
      The added descriptor.
    """

    if package:
      service_name = '.'.join((package, service_proto.name))
    else:
      service_name = service_proto.name

    methods = [self._MakeMethodDescriptor(method_proto, service_name, package,
                                          scope, index)
               for index, method_proto in enumerate(service_proto.method)]
    desc = descriptor.ServiceDescriptor(
        name=service_proto.name,
        full_name=service_name,
        index=service_index,
        methods=methods,
        options=_OptionsOrNone(service_proto),
        file=file_desc,
        # pylint: disable=protected-access
        create_key=descriptor._internal_create_key)
    self._CheckConflictRegister(desc, desc.full_name, desc.file.name)
    self._service_descriptors[service_name] = desc
    return desc

  def _MakeMethodDescriptor(self, method_proto, service_name, package, scope,
                            index):
    """Creates a method descriptor from a MethodDescriptorProto.

    Args:
      method_proto: The proto describing the method.
      service_name: The name of the containing service.
      package: Optional package name to look up for types.
      scope: Scope containing available types.
      index: Index of the method in the service.

    Returns:
      An initialized MethodDescriptor object.
    """
    full_name = '.'.join((service_name, method_proto.name))
    input_type = self._GetTypeFromScope(
        package, method_proto.input_type, scope)
    output_type = self._GetTypeFromScope(
        package, method_proto.output_type, scope)
    return descriptor.MethodDescriptor(
        name=method_proto.name,
        full_name=full_name,
        index=index,
        containing_service=None,
        input_type=input_type,
        output_type=output_type,
        options=_OptionsOrNone(method_proto),
        # pylint: disable=protected-access
        create_key=descriptor._internal_create_key)

  def _ExtractSymbols(self, descriptors):
    """Pulls out all the symbols from descriptor protos.

    Args:
      descriptors: The messages to extract descriptors from.
    Yields:
      A two element tuple of the type name and descriptor object.
    """

    for desc in descriptors:
      yield (_PrefixWithDot(desc.full_name), desc)
      for symbol in self._ExtractSymbols(desc.nested_types):
        yield symbol
      for enum in desc.enum_types:
        yield (_PrefixWithDot(enum.full_name), enum)

  def _GetDeps(self, dependencies):
    """Recursively finds dependencies for file protos.

    Args:
      dependencies: The names of the files being depended on.

    Yields:
      Each direct and indirect dependency.
    """

    for dependency in dependencies:
      dep_desc = self.FindFileByName(dependency)
      yield dep_desc
      for parent_dep in dep_desc.dependencies:
        yield parent_dep

  def _GetTypeFromScope(self, package, type_name, scope):
    """Finds a given type name in the current scope.

    Args:
      package: The package the proto should be located in.
      type_name: The name of the type to be found in the scope.
      scope: Dict mapping short and full symbols to message and enum types.

    Returns:
      The descriptor for the requested type.
    """
    if type_name not in scope:
      components = _PrefixWithDot(package).split('.')
      while components:
        possible_match = '.'.join(components + [type_name])
        if possible_match in scope:
          type_name = possible_match
          break
        else:
          components.pop(-1)
    return scope[type_name]


def _PrefixWithDot(name):
  return name if name.startswith('.') else '.%s' % name


if _USE_C_DESCRIPTORS:
  # TODO(amauryfa): This pool could be constructed from Python code, when we
  # support a flag like 'use_cpp_generated_pool=True'.
  # pylint: disable=protected-access
  _DEFAULT = descriptor._message.default_pool
else:
  _DEFAULT = DescriptorPool()


def Default():
  return _DEFAULT

social_vdgnn/lib/python3.6/site-packages/google/protobuf/descriptor_pool.py
eliatrevisan/IA-MPPI-LBM
# Protocol Buffers - Google's data interchange format
# Copyright 2008 Google Inc.  All rights reserved.
# https://developers.google.com/protocol-buffers/
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Contains container classes to represent different protocol buffer types.

This file defines container classes which represent categories of protocol
buffer field types which need extra maintenance. Currently these categories
are:

-   Repeated scalar fields - These are all repeated fields which aren't
    composite (e.g. they are of simple types like int32, string, etc).
-   Repeated composite fields - Repeated fields which are composite. This
    includes groups and nested messages.
"""

__author__ = 'petar@google.com (Petar Petrov)'

import collections.abc


class BaseContainer(object):

  """Base container class."""

  # Minimizes memory usage and disallows assignment to other attributes.
  __slots__ = ['_message_listener', '_values']

  def __init__(self, message_listener):
    """
    Args:
      message_listener: A MessageListener implementation.
        The RepeatedScalarFieldContainer will call this object's
        Modified() method when it is modified.
    """
    self._message_listener = message_listener
    self._values = []

  def __getitem__(self, key):
    """Retrieves item by the specified key."""
    return self._values[key]

  def __len__(self):
    """Returns the number of elements in the container."""
    return len(self._values)

  def __ne__(self, other):
    """Checks if another instance isn't equal to this one."""
    # The concrete classes should define __eq__.
    return not self == other

  def __hash__(self):
    raise TypeError('unhashable object')

  def __repr__(self):
    return repr(self._values)

  def sort(self, *args, **kwargs):
    # Continue to support the old sort_function keyword argument.
    # This is expected to be a rare occurrence, so use LBYL to avoid
    # the overhead of actually catching KeyError.
    if 'sort_function' in kwargs:
      kwargs['cmp'] = kwargs.pop('sort_function')
    self._values.sort(*args, **kwargs)

  def reverse(self):
    self._values.reverse()


collections.abc.MutableSequence.register(BaseContainer)


class RepeatedScalarFieldContainer(BaseContainer):
  """Simple, type-checked, list-like container for holding repeated scalars."""

  # Disallows assignment to other attributes.
  __slots__ = ['_type_checker']

  def __init__(self, message_listener, type_checker):
    """Args:

      message_listener: A MessageListener implementation. The
      RepeatedScalarFieldContainer will call this object's Modified() method
      when it is modified.
      type_checker: A type_checkers.ValueChecker instance to run on elements
      inserted into this container.
    """
    super(RepeatedScalarFieldContainer, self).__init__(message_listener)
    self._type_checker = type_checker

  def append(self, value):
    """Appends an item to the list. Similar to list.append()."""
    self._values.append(self._type_checker.CheckValue(value))
    if not self._message_listener.dirty:
      self._message_listener.Modified()

  def insert(self, key, value):
    """Inserts the item at the specified position. Similar to list.insert()."""
    self._values.insert(key, self._type_checker.CheckValue(value))
    if not self._message_listener.dirty:
      self._message_listener.Modified()

  def extend(self, elem_seq):
    """Extends by appending the given iterable. Similar to list.extend()."""

    if elem_seq is None:
      return
    try:
      elem_seq_iter = iter(elem_seq)
    except TypeError:
      if not elem_seq:
        # silently ignore falsy inputs :-/.
        # TODO(ptucker): Deprecate this behavior. b/18413862
        return
      raise

    new_values = [self._type_checker.CheckValue(elem) for elem in elem_seq_iter]
    if new_values:
      self._values.extend(new_values)
    self._message_listener.Modified()

  def MergeFrom(self, other):
    """Appends the contents of another repeated field of the same type to this
    one. We do not check the types of the individual fields.
    """
    self._values.extend(other._values)
    self._message_listener.Modified()

  def remove(self, elem):
    """Removes an item from the list. Similar to list.remove()."""
    self._values.remove(elem)
    self._message_listener.Modified()

  def pop(self, key=-1):
    """Removes and returns an item at a given index. Similar to list.pop()."""
    value = self._values[key]
    self.__delitem__(key)
    return value

  def __setitem__(self, key, value):
    """Sets the item on the specified position."""
    if isinstance(key, slice):  # PY3
      if key.step is not None:
        raise ValueError('Extended slices not supported')
      self.__setslice__(key.start, key.stop, value)
    else:
      self._values[key] = self._type_checker.CheckValue(value)
      self._message_listener.Modified()

  def __getslice__(self, start, stop):
    """Retrieves the subset of items from between the specified indices."""
    return self._values[start:stop]

  def __setslice__(self, start, stop, values):
    """Sets the subset of items from between the specified indices."""
    new_values = []
    for value in values:
      new_values.append(self._type_checker.CheckValue(value))
    self._values[start:stop] = new_values
    self._message_listener.Modified()

  def __delitem__(self, key):
    """Deletes the item at the specified position."""
    del self._values[key]
    self._message_listener.Modified()

  def __delslice__(self, start, stop):
    """Deletes the subset of items from between the specified indices."""
    del self._values[start:stop]
    self._message_listener.Modified()

  def __eq__(self, other):
    """Compares the current instance with another one."""
    if self is other:
      return True
    # Special case for the same type which should be common and fast.
    if isinstance(other, self.__class__):
      return other._values == self._values
    # We are presumably comparing against some other sequence type.
    return other == self._values


class RepeatedCompositeFieldContainer(BaseContainer):

  """Simple, list-like container for holding repeated composite fields."""

  # Disallows assignment to other attributes.
  __slots__ = ['_message_descriptor']

  def __init__(self, message_listener, message_descriptor):
    """
    Note that we pass in a descriptor instead of the generated directly,
    since at the time we construct a _RepeatedCompositeFieldContainer we
    haven't yet necessarily initialized the type that will be contained in the
    container.

    Args:
      message_listener: A MessageListener implementation.
        The RepeatedCompositeFieldContainer will call this object's
        Modified() method when it is modified.
      message_descriptor: A Descriptor instance describing the protocol type
        that should be present in this container.  We'll use the
        _concrete_class field of this descriptor when the client calls add().
    """
    super(RepeatedCompositeFieldContainer, self).__init__(message_listener)
    self._message_descriptor = message_descriptor

  def add(self, **kwargs):
    """Adds a new element at the end of the list and returns it. Keyword
    arguments may be used to initialize the element.
    """
    new_element = self._message_descriptor._concrete_class(**kwargs)
    new_element._SetListener(self._message_listener)
    self._values.append(new_element)
    if not self._message_listener.dirty:
      self._message_listener.Modified()
    return new_element

  def append(self, value):
    """Appends one element by copying the message."""
    new_element = self._message_descriptor._concrete_class()
    new_element._SetListener(self._message_listener)
    new_element.CopyFrom(value)
    self._values.append(new_element)
    if not self._message_listener.dirty:
      self._message_listener.Modified()

  def insert(self, key, value):
    """Inserts the item at the specified position by copying."""
    new_element = self._message_descriptor._concrete_class()
    new_element._SetListener(self._message_listener)
    new_element.CopyFrom(value)
    self._values.insert(key, new_element)
    if not self._message_listener.dirty:
      self._message_listener.Modified()

  def extend(self, elem_seq):
    """Extends by appending the given sequence of elements of the same type

    as this one, copying each individual message.
    """
    message_class = self._message_descriptor._concrete_class
    listener = self._message_listener
    values = self._values
    for message in elem_seq:
      new_element = message_class()
      new_element._SetListener(listener)
      new_element.MergeFrom(message)
      values.append(new_element)
    listener.Modified()

  def MergeFrom(self, other):
    """Appends the contents of another repeated field of the same type to this
    one, copying each individual message.
    """
    self.extend(other._values)

  def remove(self, elem):
    """Removes an item from the list. Similar to list.remove()."""
    self._values.remove(elem)
    self._message_listener.Modified()

  def pop(self, key=-1):
    """Removes and returns an item at a given index. Similar to list.pop()."""
    value = self._values[key]
    self.__delitem__(key)
    return value

  def __getslice__(self, start, stop):
    """Retrieves the subset of items from between the specified indices."""
    return self._values[start:stop]

  def __delitem__(self, key):
    """Deletes the item at the specified position."""
    del self._values[key]
    self._message_listener.Modified()

  def __delslice__(self, start, stop):
    """Deletes the subset of items from between the specified indices."""
    del self._values[start:stop]
    self._message_listener.Modified()

  def __eq__(self, other):
    """Compares the current instance with another one."""
    if self is other:
      return True
    if not isinstance(other, self.__class__):
      raise TypeError('Can only compare repeated composite fields against '
                      'other repeated composite fields.')
    return self._values == other._values


class ScalarMap(collections.abc.MutableMapping):

  """Simple, type-checked, dict-like container for holding repeated scalars."""

  # Disallows assignment to other attributes.
  __slots__ = ['_key_checker', '_value_checker', '_values', '_message_listener',
               '_entry_descriptor']

  def __init__(self, message_listener, key_checker, value_checker,
               entry_descriptor):
    """
    Args:
      message_listener: A MessageListener implementation.
        The ScalarMap will call this object's Modified() method when it
        is modified.
      key_checker: A type_checkers.ValueChecker instance to run on keys
        inserted into this container.
      value_checker: A type_checkers.ValueChecker instance to run on values
        inserted into this container.
      entry_descriptor: The MessageDescriptor of a map entry: key and value.
    """
    self._message_listener = message_listener
    self._key_checker = key_checker
    self._value_checker = value_checker
    self._entry_descriptor = entry_descriptor
    self._values = {}

  def __getitem__(self, key):
    try:
      return self._values[key]
    except KeyError:
      key = self._key_checker.CheckValue(key)
      val = self._value_checker.DefaultValue()
      self._values[key] = val
      return val

  def __contains__(self, item):
    # We check the key's type to match the strong-typing flavor of the API.
    # Also this makes it easier to match the behavior of the C++ implementation.
    self._key_checker.CheckValue(item)
    return item in self._values

  # We need to override this explicitly, because our defaultdict-like behavior
  # will make the default implementation (from our base class) always insert
  # the key.
  def get(self, key, default=None):
    if key in self:
      return self[key]
    else:
      return default

  def __setitem__(self, key, value):
    checked_key = self._key_checker.CheckValue(key)
    checked_value = self._value_checker.CheckValue(value)
    self._values[checked_key] = checked_value
    self._message_listener.Modified()

  def __delitem__(self, key):
    del self._values[key]
    self._message_listener.Modified()

  def __len__(self):
    return len(self._values)

  def __iter__(self):
    return iter(self._values)

  def __repr__(self):
    return repr(self._values)

  def MergeFrom(self, other):
    self._values.update(other._values)
    self._message_listener.Modified()

  def InvalidateIterators(self):
    # It appears that the only way to reliably invalidate iterators to
    # self._values is to ensure that its size changes.
    original = self._values
    self._values = original.copy()
    original[None] = None

  # This is defined in the abstract base, but we can do it much more cheaply.
  def clear(self):
    self._values.clear()
    self._message_listener.Modified()

  def GetEntryClass(self):
    return self._entry_descriptor._concrete_class


class MessageMap(collections.abc.MutableMapping):

  """Simple, type-checked, dict-like container for with submessage values."""

  # Disallows assignment to other attributes.
  __slots__ = ['_key_checker', '_values', '_message_listener',
               '_message_descriptor', '_entry_descriptor']

  def __init__(self, message_listener, message_descriptor, key_checker,
               entry_descriptor):
    """
    Args:
      message_listener: A MessageListener implementation.
        The ScalarMap will call this object's Modified() method when it
        is modified.
      key_checker: A type_checkers.ValueChecker instance to run on keys
        inserted into this container.
      value_checker: A type_checkers.ValueChecker instance to run on values
        inserted into this container.
      entry_descriptor: The MessageDescriptor of a map entry: key and value.
    """
    self._message_listener = message_listener
    self._message_descriptor = message_descriptor
    self._key_checker = key_checker
    self._entry_descriptor = entry_descriptor
    self._values = {}

  def __getitem__(self, key):
    key = self._key_checker.CheckValue(key)
    try:
      return self._values[key]
    except KeyError:
      new_element = self._message_descriptor._concrete_class()
      new_element._SetListener(self._message_listener)
      self._values[key] = new_element
      self._message_listener.Modified()

      return new_element

  def get_or_create(self, key):
    """get_or_create() is an alias for getitem (ie. map[key]).

    Args:
      key: The key to get or create in the map.

    This is useful in cases where you want to be explicit that the call is
    mutating the map.  This can avoid lint errors for statements like this
    that otherwise would appear to be pointless statements:

      msg.my_map[key]
    """
    return self[key]

  # We need to override this explicitly, because our defaultdict-like behavior
  # will make the default implementation (from our base class) always insert
  # the key.
  def get(self, key, default=None):
    if key in self:
      return self[key]
    else:
      return default

  def __contains__(self, item):
    item = self._key_checker.CheckValue(item)
    return item in self._values

  def __setitem__(self, key, value):
    raise ValueError('May not set values directly, call my_map[key].foo = 5')

  def __delitem__(self, key):
    key = self._key_checker.CheckValue(key)
    del self._values[key]
    self._message_listener.Modified()

  def __len__(self):
    return len(self._values)

  def __iter__(self):
    return iter(self._values)

  def __repr__(self):
    return repr(self._values)

  def MergeFrom(self, other):
    # pylint: disable=protected-access
    for key in other._values:
      # According to documentation: "When parsing from the wire or when merging,
      # if there are duplicate map keys the last key seen is used".
      if key in self:
        del self[key]
      self[key].CopyFrom(other[key])
    # self._message_listener.Modified() not required here, because
    # mutations to submessages already propagate.

  def InvalidateIterators(self):
    # It appears that the only way to reliably invalidate iterators to
    # self._values is to ensure that its size changes.
    original = self._values
    self._values = original.copy()
    original[None] = None

  # This is defined in the abstract base, but we can do it much more cheaply.
  def clear(self):
    self._values.clear()
    self._message_listener.Modified()

  def GetEntryClass(self):
    return self._entry_descriptor._concrete_class


class _UnknownField(object):

  """A parsed unknown field."""

  # Disallows assignment to other attributes.
  __slots__ = ['_field_number', '_wire_type', '_data']

  def __init__(self, field_number, wire_type, data):
    self._field_number = field_number
    self._wire_type = wire_type
    self._data = data
    return

  def __lt__(self, other):
    # pylint: disable=protected-access
    return self._field_number < other._field_number

  def __eq__(self, other):
    if self is other:
      return True
    # pylint: disable=protected-access
    return (self._field_number == other._field_number and
            self._wire_type == other._wire_type and
            self._data == other._data)


class UnknownFieldRef(object):

  def __init__(self, parent, index):
    self._parent = parent
    self._index = index
    return

  def _check_valid(self):
    if not self._parent:
      raise ValueError('UnknownField does not exist. '
                       'The parent message might be cleared.')
    if self._index >= len(self._parent):
      raise ValueError('UnknownField does not exist. '
                       'The parent message might be cleared.')

  @property
  def field_number(self):
    self._check_valid()
    # pylint: disable=protected-access
    return self._parent._internal_get(self._index)._field_number

  @property
  def wire_type(self):
    self._check_valid()
    # pylint: disable=protected-access
    return self._parent._internal_get(self._index)._wire_type

  @property
  def data(self):
    self._check_valid()
    # pylint: disable=protected-access
    return self._parent._internal_get(self._index)._data


class UnknownFieldSet(object):

  """UnknownField container"""

  # Disallows assignment to other attributes.
  __slots__ = ['_values']

  def __init__(self):
    self._values = []

  def __getitem__(self, index):
    if self._values is None:
      raise ValueError('UnknownFields does not exist. '
                       'The parent message might be cleared.')
    size = len(self._values)
    if index < 0:
      index += size
    if index < 0 or index >= size:
      raise IndexError('index %d out of range'.index)

    return UnknownFieldRef(self, index)

  def _internal_get(self, index):
    return self._values[index]

  def __len__(self):
    if self._values is None:
      raise ValueError('UnknownFields does not exist. '
                       'The parent message might be cleared.')
    return len(self._values)

  def _add(self, field_number, wire_type, data):
    unknown_field = _UnknownField(field_number, wire_type, data)
    self._values.append(unknown_field)
    return unknown_field

  def __iter__(self):
    for i in range(len(self)):
      yield UnknownFieldRef(self, i)

  def _extend(self, other):
    if other is None:
      return
    # pylint: disable=protected-access
    self._values.extend(other._values)

  def __eq__(self, other):
    if self is other:
      return True
    # Sort unknown fields because their order shouldn't
    # affect equality test.
    values = list(self._values)
    if other is None:
      return not values
    values.sort()
    # pylint: disable=protected-access
    other_values = sorted(other._values)
    return values == other_values

  def _clear(self):
    for value in self._values:
      # pylint: disable=protected-access
      if isinstance(value._data, UnknownFieldSet):
        value._data._clear()  # pylint: disable=protected-access
    self._values = None

social_vdgnn/lib/python3.6/site-packages/google/protobuf/internal/containers.py
eliatrevisan/IA-MPPI-LBM
# Protocol Buffers - Google's data interchange format
# Copyright 2008 Google Inc.  All rights reserved.
# https://developers.google.com/protocol-buffers/
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Contains well known classes.

This files defines well known classes which need extra maintenance including:
  - Any
  - Duration
  - FieldMask
  - Struct
  - Timestamp
"""

__author__ = 'jieluo@google.com (Jie Luo)'

import calendar
import collections.abc
from datetime import datetime
from datetime import timedelta

from google.protobuf.descriptor import FieldDescriptor

_TIMESTAMPFOMAT = '%Y-%m-%dT%H:%M:%S'
_NANOS_PER_SECOND = 1000000000
_NANOS_PER_MILLISECOND = 1000000
_NANOS_PER_MICROSECOND = 1000
_MILLIS_PER_SECOND = 1000
_MICROS_PER_SECOND = 1000000
_SECONDS_PER_DAY = 24 * 3600
_DURATION_SECONDS_MAX = 315576000000


class Any(object):
  """Class for Any Message type."""

  __slots__ = ()

  def Pack(self, msg, type_url_prefix='type.googleapis.com/',
           deterministic=None):
    """Packs the specified message into current Any message."""
    if len(type_url_prefix) < 1 or type_url_prefix[-1] != '/':
      self.type_url = '%s/%s' % (type_url_prefix, msg.DESCRIPTOR.full_name)
    else:
      self.type_url = '%s%s' % (type_url_prefix, msg.DESCRIPTOR.full_name)
    self.value = msg.SerializeToString(deterministic=deterministic)

  def Unpack(self, msg):
    """Unpacks the current Any message into specified message."""
    descriptor = msg.DESCRIPTOR
    if not self.Is(descriptor):
      return False
    msg.ParseFromString(self.value)
    return True

  def TypeName(self):
    """Returns the protobuf type name of the inner message."""
    # Only last part is to be used: b/25630112
    return self.type_url.split('/')[-1]

  def Is(self, descriptor):
    """Checks if this Any represents the given protobuf type."""
    return '/' in self.type_url and self.TypeName() == descriptor.full_name


_EPOCH_DATETIME = datetime.utcfromtimestamp(0)


class Timestamp(object):
  """Class for Timestamp message type."""

  __slots__ = ()

  def ToJsonString(self):
    """Converts Timestamp to RFC 3339 date string format.

    Returns:
      A string converted from timestamp. The string is always Z-normalized
      and uses 3, 6 or 9 fractional digits as required to represent the
      exact time. Example of the return format: '1972-01-01T10:00:20.021Z'
    """
    nanos = self.nanos % _NANOS_PER_SECOND
    total_sec = self.seconds + (self.nanos - nanos) // _NANOS_PER_SECOND
    seconds = total_sec % _SECONDS_PER_DAY
    days = (total_sec - seconds) // _SECONDS_PER_DAY
    dt = datetime(1970, 1, 1) + timedelta(days, seconds)

    result = dt.isoformat()
    if (nanos % 1e9) == 0:
      # If there are 0 fractional digits, the fractional
      # point '.' should be omitted when serializing.
      return result + 'Z'
    if (nanos % 1e6) == 0:
      # Serialize 3 fractional digits.
      return result + '.%03dZ' % (nanos / 1e6)
    if (nanos % 1e3) == 0:
      # Serialize 6 fractional digits.
      return result + '.%06dZ' % (nanos / 1e3)
    # Serialize 9 fractional digits.
    return result + '.%09dZ' % nanos

  def FromJsonString(self, value):
    """Parse a RFC 3339 date string format to Timestamp.

    Args:
      value: A date string. Any fractional digits (or none) and any offset are
          accepted as long as they fit into nano-seconds precision.
          Example of accepted format: '1972-01-01T10:00:20.021-05:00'

    Raises:
      ValueError: On parsing problems.
    """
    if not isinstance(value, str):
      raise ValueError('Timestamp JSON value not a string: {!r}'.format(value))
    timezone_offset = value.find('Z')
    if timezone_offset == -1:
      timezone_offset = value.find('+')
    if timezone_offset == -1:
      timezone_offset = value.rfind('-')
    if timezone_offset == -1:
      raise ValueError(
          'Failed to parse timestamp: missing valid timezone offset.')
    time_value = value[0:timezone_offset]
    # Parse datetime and nanos.
    point_position = time_value.find('.')
    if point_position == -1:
      second_value = time_value
      nano_value = ''
    else:
      second_value = time_value[:point_position]
      nano_value = time_value[point_position + 1:]
    if 't' in second_value:
      raise ValueError(
          'time data \'{0}\' does not match format \'%Y-%m-%dT%H:%M:%S\', '
          'lowercase \'t\' is not accepted'.format(second_value))
    date_object = datetime.strptime(second_value, _TIMESTAMPFOMAT)
    td = date_object - datetime(1970, 1, 1)
    seconds = td.seconds + td.days * _SECONDS_PER_DAY
    if len(nano_value) > 9:
      raise ValueError(
          'Failed to parse Timestamp: nanos {0} more than '
          '9 fractional digits.'.format(nano_value))
    if nano_value:
      nanos = round(float('0.' + nano_value) * 1e9)
    else:
      nanos = 0
    # Parse timezone offsets.
    if value[timezone_offset] == 'Z':
      if len(value) != timezone_offset + 1:
        raise ValueError('Failed to parse timestamp: invalid trailing'
                         ' data {0}.'.format(value))
    else:
      timezone = value[timezone_offset:]
      pos = timezone.find(':')
      if pos == -1:
        raise ValueError(
            'Invalid timezone offset value: {0}.'.format(timezone))
      if timezone[0] == '+':
        seconds -= (int(timezone[1:pos])*60+int(timezone[pos+1:]))*60
      else:
        seconds += (int(timezone[1:pos])*60+int(timezone[pos+1:]))*60
    # Set seconds and nanos
    self.seconds = int(seconds)
    self.nanos = int(nanos)

  def GetCurrentTime(self):
    """Get the current UTC into Timestamp."""
    self.FromDatetime(datetime.utcnow())

  def ToNanoseconds(self):
    """Converts Timestamp to nanoseconds since epoch."""
    return self.seconds * _NANOS_PER_SECOND + self.nanos

  def ToMicroseconds(self):
    """Converts Timestamp to microseconds since epoch."""
    return (self.seconds * _MICROS_PER_SECOND +
            self.nanos // _NANOS_PER_MICROSECOND)

  def ToMilliseconds(self):
    """Converts Timestamp to milliseconds since epoch."""
    return (self.seconds * _MILLIS_PER_SECOND +
            self.nanos // _NANOS_PER_MILLISECOND)

  def ToSeconds(self):
    """Converts Timestamp to seconds since epoch."""
    return self.seconds

  def FromNanoseconds(self, nanos):
    """Converts nanoseconds since epoch to Timestamp."""
    self.seconds = nanos // _NANOS_PER_SECOND
    self.nanos = nanos % _NANOS_PER_SECOND

  def FromMicroseconds(self, micros):
    """Converts microseconds since epoch to Timestamp."""
    self.seconds = micros // _MICROS_PER_SECOND
    self.nanos = (micros % _MICROS_PER_SECOND) * _NANOS_PER_MICROSECOND

  def FromMilliseconds(self, millis):
    """Converts milliseconds since epoch to Timestamp."""
    self.seconds = millis // _MILLIS_PER_SECOND
    self.nanos = (millis % _MILLIS_PER_SECOND) * _NANOS_PER_MILLISECOND

  def FromSeconds(self, seconds):
    """Converts seconds since epoch to Timestamp."""
    self.seconds = seconds
    self.nanos = 0

  def ToDatetime(self):
    """Converts Timestamp to datetime."""
    return _EPOCH_DATETIME + timedelta(
        seconds=self.seconds, microseconds=_RoundTowardZero(
            self.nanos, _NANOS_PER_MICROSECOND))

  def FromDatetime(self, dt):
    """Converts datetime to Timestamp."""
    # Using this guide: http://wiki.python.org/moin/WorkingWithTime
    # And this conversion guide: http://docs.python.org/library/time.html

    # Turn the date parameter into a tuple (struct_time) that can then be
    # manipulated into a long value of seconds.  During the conversion from
    # struct_time to long, the source date in UTC, and so it follows that the
    # correct transformation is calendar.timegm()
    self.seconds = calendar.timegm(dt.utctimetuple())
    self.nanos = dt.microsecond * _NANOS_PER_MICROSECOND


class Duration(object):
  """Class for Duration message type."""

  __slots__ = ()

  def ToJsonString(self):
    """Converts Duration to string format.

    Returns:
      A string converted from self. The string format will contains
      3, 6, or 9 fractional digits depending on the precision required to
      represent the exact Duration value. For example: "1s", "1.010s",
      "1.000000100s", "-3.100s"
    """
    _CheckDurationValid(self.seconds, self.nanos)
    if self.seconds < 0 or self.nanos < 0:
      result = '-'
      seconds = - self.seconds + int((0 - self.nanos) // 1e9)
      nanos = (0 - self.nanos) % 1e9
    else:
      result = ''
      seconds = self.seconds + int(self.nanos // 1e9)
      nanos = self.nanos % 1e9
    result += '%d' % seconds
    if (nanos % 1e9) == 0:
      # If there are 0 fractional digits, the fractional
      # point '.' should be omitted when serializing.
      return result + 's'
    if (nanos % 1e6) == 0:
      # Serialize 3 fractional digits.
      return result + '.%03ds' % (nanos / 1e6)
    if (nanos % 1e3) == 0:
      # Serialize 6 fractional digits.
      return result + '.%06ds' % (nanos / 1e3)
    # Serialize 9 fractional digits.
    return result + '.%09ds' % nanos

  def FromJsonString(self, value):
    """Converts a string to Duration.

    Args:
      value: A string to be converted. The string must end with 's'. Any
          fractional digits (or none) are accepted as long as they fit into
          precision. For example: "1s", "1.01s", "1.0000001s", "-3.100s

    Raises:
      ValueError: On parsing problems.
    """
    if not isinstance(value, str):
      raise ValueError('Duration JSON value not a string: {!r}'.format(value))
    if len(value) < 1 or value[-1] != 's':
      raise ValueError(
          'Duration must end with letter "s": {0}.'.format(value))
    try:
      pos = value.find('.')
      if pos == -1:
        seconds = int(value[:-1])
        nanos = 0
      else:
        seconds = int(value[:pos])
        if value[0] == '-':
          nanos = int(round(float('-0{0}'.format(value[pos: -1])) *1e9))
        else:
          nanos = int(round(float('0{0}'.format(value[pos: -1])) *1e9))
      _CheckDurationValid(seconds, nanos)
      self.seconds = seconds
      self.nanos = nanos
    except ValueError as e:
      raise ValueError(
          'Couldn\'t parse duration: {0} : {1}.'.format(value, e))

  def ToNanoseconds(self):
    """Converts a Duration to nanoseconds."""
    return self.seconds * _NANOS_PER_SECOND + self.nanos

  def ToMicroseconds(self):
    """Converts a Duration to microseconds."""
    micros = _RoundTowardZero(self.nanos, _NANOS_PER_MICROSECOND)
    return self.seconds * _MICROS_PER_SECOND + micros

  def ToMilliseconds(self):
    """Converts a Duration to milliseconds."""
    millis = _RoundTowardZero(self.nanos, _NANOS_PER_MILLISECOND)
    return self.seconds * _MILLIS_PER_SECOND + millis

  def ToSeconds(self):
    """Converts a Duration to seconds."""
    return self.seconds

  def FromNanoseconds(self, nanos):
    """Converts nanoseconds to Duration."""
    self._NormalizeDuration(nanos // _NANOS_PER_SECOND,
                            nanos % _NANOS_PER_SECOND)

  def FromMicroseconds(self, micros):
    """Converts microseconds to Duration."""
    self._NormalizeDuration(
        micros // _MICROS_PER_SECOND,
        (micros % _MICROS_PER_SECOND) * _NANOS_PER_MICROSECOND)

  def FromMilliseconds(self, millis):
    """Converts milliseconds to Duration."""
    self._NormalizeDuration(
        millis // _MILLIS_PER_SECOND,
        (millis % _MILLIS_PER_SECOND) * _NANOS_PER_MILLISECOND)

  def FromSeconds(self, seconds):
    """Converts seconds to Duration."""
    self.seconds = seconds
    self.nanos = 0

  def ToTimedelta(self):
    """Converts Duration to timedelta."""
    return timedelta(
        seconds=self.seconds, microseconds=_RoundTowardZero(
            self.nanos, _NANOS_PER_MICROSECOND))

  def FromTimedelta(self, td):
    """Converts timedelta to Duration."""
    self._NormalizeDuration(td.seconds + td.days * _SECONDS_PER_DAY,
                            td.microseconds * _NANOS_PER_MICROSECOND)

  def _NormalizeDuration(self, seconds, nanos):
    """Set Duration by seconds and nanos."""
    # Force nanos to be negative if the duration is negative.
    if seconds < 0 and nanos > 0:
      seconds += 1
      nanos -= _NANOS_PER_SECOND
    self.seconds = seconds
    self.nanos = nanos


def _CheckDurationValid(seconds, nanos):
  if seconds < -_DURATION_SECONDS_MAX or seconds > _DURATION_SECONDS_MAX:
    raise ValueError(
        'Duration is not valid: Seconds {0} must be in range '
        '[-315576000000, 315576000000].'.format(seconds))
  if nanos <= -_NANOS_PER_SECOND or nanos >= _NANOS_PER_SECOND:
    raise ValueError(
        'Duration is not valid: Nanos {0} must be in range '
        '[-999999999, 999999999].'.format(nanos))
  if (nanos < 0 and seconds > 0) or (nanos > 0 and seconds < 0):
    raise ValueError(
        'Duration is not valid: Sign mismatch.')


def _RoundTowardZero(value, divider):
  """Truncates the remainder part after division."""
  # For some languages, the sign of the remainder is implementation
  # dependent if any of the operands is negative. Here we enforce
  # "rounded toward zero" semantics. For example, for (-5) / 2 an
  # implementation may give -3 as the result with the remainder being
  # 1. This function ensures we always return -2 (closer to zero).
  result = value // divider
  remainder = value % divider
  if result < 0 and remainder > 0:
    return result + 1
  else:
    return result


class FieldMask(object):
  """Class for FieldMask message type."""

  __slots__ = ()

  def ToJsonString(self):
    """Converts FieldMask to string according to proto3 JSON spec."""
    camelcase_paths = []
    for path in self.paths:
      camelcase_paths.append(_SnakeCaseToCamelCase(path))
    return ','.join(camelcase_paths)

  def FromJsonString(self, value):
    """Converts string to FieldMask according to proto3 JSON spec."""
    if not isinstance(value, str):
      raise ValueError('FieldMask JSON value not a string: {!r}'.format(value))
    self.Clear()
    if value:
      for path in value.split(','):
        self.paths.append(_CamelCaseToSnakeCase(path))

  def IsValidForDescriptor(self, message_descriptor):
    """Checks whether the FieldMask is valid for Message Descriptor."""
    for path in self.paths:
      if not _IsValidPath(message_descriptor, path):
        return False
    return True

  def AllFieldsFromDescriptor(self, message_descriptor):
    """Gets all direct fields of Message Descriptor to FieldMask."""
    self.Clear()
    for field in message_descriptor.fields:
      self.paths.append(field.name)

  def CanonicalFormFromMask(self, mask):
    """Converts a FieldMask to the canonical form.

    Removes paths that are covered by another path. For example,
    "foo.bar" is covered by "foo" and will be removed if "foo"
    is also in the FieldMask. Then sorts all paths in alphabetical order.

    Args:
      mask: The original FieldMask to be converted.
    """
    tree = _FieldMaskTree(mask)
    tree.ToFieldMask(self)

  def Union(self, mask1, mask2):
    """Merges mask1 and mask2 into this FieldMask."""
    _CheckFieldMaskMessage(mask1)
    _CheckFieldMaskMessage(mask2)
    tree = _FieldMaskTree(mask1)
    tree.MergeFromFieldMask(mask2)
    tree.ToFieldMask(self)

  def Intersect(self, mask1, mask2):
    """Intersects mask1 and mask2 into this FieldMask."""
    _CheckFieldMaskMessage(mask1)
    _CheckFieldMaskMessage(mask2)
    tree = _FieldMaskTree(mask1)
    intersection = _FieldMaskTree()
    for path in mask2.paths:
      tree.IntersectPath(path, intersection)
    intersection.ToFieldMask(self)

  def MergeMessage(
      self, source, destination,
      replace_message_field=False, replace_repeated_field=False):
    """Merges fields specified in FieldMask from source to destination.

    Args:
      source: Source message.
      destination: The destination message to be merged into.
      replace_message_field: Replace message field if True. Merge message
          field if False.
      replace_repeated_field: Replace repeated field if True. Append
          elements of repeated field if False.
    """
    tree = _FieldMaskTree(self)
    tree.MergeMessage(
        source, destination, replace_message_field, replace_repeated_field)


def _IsValidPath(message_descriptor, path):
  """Checks whether the path is valid for Message Descriptor."""
  parts = path.split('.')
  last = parts.pop()
  for name in parts:
    field = message_descriptor.fields_by_name.get(name)
    if (field is None or
        field.label == FieldDescriptor.LABEL_REPEATED or
        field.type != FieldDescriptor.TYPE_MESSAGE):
      return False
    message_descriptor = field.message_type
  return last in message_descriptor.fields_by_name


def _CheckFieldMaskMessage(message):
  """Raises ValueError if message is not a FieldMask."""
  message_descriptor = message.DESCRIPTOR
  if (message_descriptor.name != 'FieldMask' or
      message_descriptor.file.name != 'google/protobuf/field_mask.proto'):
    raise ValueError('Message {0} is not a FieldMask.'.format(
        message_descriptor.full_name))


def _SnakeCaseToCamelCase(path_name):
  """Converts a path name from snake_case to camelCase."""
  result = []
  after_underscore = False
  for c in path_name:
    if c.isupper():
      raise ValueError(
          'Fail to print FieldMask to Json string: Path name '
          '{0} must not contain uppercase letters.'.format(path_name))
    if after_underscore:
      if c.islower():
        result.append(c.upper())
        after_underscore = False
      else:
        raise ValueError(
            'Fail to print FieldMask to Json string: The '
            'character after a "_" must be a lowercase letter '
            'in path name {0}.'.format(path_name))
    elif c == '_':
      after_underscore = True
    else:
      result += c

  if after_underscore:
    raise ValueError('Fail to print FieldMask to Json string: Trailing "_" '
                     'in path name {0}.'.format(path_name))
  return ''.join(result)


def _CamelCaseToSnakeCase(path_name):
  """Converts a field name from camelCase to snake_case."""
  result = []
  for c in path_name:
    if c == '_':
      raise ValueError('Fail to parse FieldMask: Path name '
                       '{0} must not contain "_"s.'.format(path_name))
    if c.isupper():
      result += '_'
      result += c.lower()
    else:
      result += c
  return ''.join(result)


class _FieldMaskTree(object):
  """Represents a FieldMask in a tree structure.

  For example, given a FieldMask "foo.bar,foo.baz,bar.baz",
  the FieldMaskTree will be:
      [_root] -+- foo -+- bar
            |       |
            |       +- baz
            |
            +- bar --- baz
  In the tree, each leaf node represents a field path.
  """

  __slots__ = ('_root',)

  def __init__(self, field_mask=None):
    """Initializes the tree by FieldMask."""
    self._root = {}
    if field_mask:
      self.MergeFromFieldMask(field_mask)

  def MergeFromFieldMask(self, field_mask):
    """Merges a FieldMask to the tree."""
    for path in field_mask.paths:
      self.AddPath(path)

  def AddPath(self, path):
    """Adds a field path into the tree.

    If the field path to add is a sub-path of an existing field path
    in the tree (i.e., a leaf node), it means the tree already matches
    the given path so nothing will be added to the tree. If the path
    matches an existing non-leaf node in the tree, that non-leaf node
    will be turned into a leaf node with all its children removed because
    the path matches all the node's children. Otherwise, a new path will
    be added.

    Args:
      path: The field path to add.
    """
    node = self._root
    for name in path.split('.'):
      if name not in node:
        node[name] = {}
      elif not node[name]:
        # Pre-existing empty node implies we already have this entire tree.
        return
      node = node[name]
    # Remove any sub-trees we might have had.
    node.clear()

  def ToFieldMask(self, field_mask):
    """Converts the tree to a FieldMask."""
    field_mask.Clear()
    _AddFieldPaths(self._root, '', field_mask)

  def IntersectPath(self, path, intersection):
    """Calculates the intersection part of a field path with this tree.

    Args:
      path: The field path to calculates.
      intersection: The out tree to record the intersection part.
    """
    node = self._root
    for name in path.split('.'):
      if name not in node:
        return
      elif not node[name]:
        intersection.AddPath(path)
        return
      node = node[name]
    intersection.AddLeafNodes(path, node)

  def AddLeafNodes(self, prefix, node):
    """Adds leaf nodes begin with prefix to this tree."""
    if not node:
      self.AddPath(prefix)
    for name in node:
      child_path = prefix + '.' + name
      self.AddLeafNodes(child_path, node[name])

  def MergeMessage(
      self, source, destination,
      replace_message, replace_repeated):
    """Merge all fields specified by this tree from source to destination."""
    _MergeMessage(
        self._root, source, destination, replace_message, replace_repeated)


def _StrConvert(value):
  """Converts value to str if it is not."""
  # This file is imported by c extension and some methods like ClearField
  # requires string for the field name. py2/py3 has different text
  # type and may use unicode.
  if not isinstance(value, str):
    return value.encode('utf-8')
  return value


def _MergeMessage(
    node, source, destination, replace_message, replace_repeated):
  """Merge all fields specified by a sub-tree from source to destination."""
  source_descriptor = source.DESCRIPTOR
  for name in node:
    child = node[name]
    field = source_descriptor.fields_by_name[name]
    if field is None:
      raise ValueError('Error: Can\'t find field {0} in message {1}.'.format(
          name, source_descriptor.full_name))
    if child:
      # Sub-paths are only allowed for singular message fields.
      if (field.label == FieldDescriptor.LABEL_REPEATED or
          field.cpp_type != FieldDescriptor.CPPTYPE_MESSAGE):
        raise ValueError('Error: Field {0} in message {1} is not a singular '
                         'message field and cannot have sub-fields.'.format(
                             name, source_descriptor.full_name))
      if source.HasField(name):
        _MergeMessage(
            child, getattr(source, name), getattr(destination, name),
            replace_message, replace_repeated)
      continue
    if field.label == FieldDescriptor.LABEL_REPEATED:
      if replace_repeated:
        destination.ClearField(_StrConvert(name))
      repeated_source = getattr(source, name)
      repeated_destination = getattr(destination, name)
      repeated_destination.MergeFrom(repeated_source)
    else:
      if field.cpp_type == FieldDescriptor.CPPTYPE_MESSAGE:
        if replace_message:
          destination.ClearField(_StrConvert(name))
        if source.HasField(name):
          getattr(destination, name).MergeFrom(getattr(source, name))
      else:
        setattr(destination, name, getattr(source, name))


def _AddFieldPaths(node, prefix, field_mask):
  """Adds the field paths descended from node to field_mask."""
  if not node and prefix:
    field_mask.paths.append(prefix)
    return
  for name in sorted(node):
    if prefix:
      child_path = prefix + '.' + name
    else:
      child_path = name
    _AddFieldPaths(node[name], child_path, field_mask)


def _SetStructValue(struct_value, value):
  if value is None:
    struct_value.null_value = 0
  elif isinstance(value, bool):
    # Note: this check must come before the number check because in Python
    # True and False are also considered numbers.
    struct_value.bool_value = value
  elif isinstance(value, str):
    struct_value.string_value = value
  elif isinstance(value, (int, float)):
    struct_value.number_value = value
  elif isinstance(value, (dict, Struct)):
    struct_value.struct_value.Clear()
    struct_value.struct_value.update(value)
  elif isinstance(value, (list, ListValue)):
    struct_value.list_value.Clear()
    struct_value.list_value.extend(value)
  else:
    raise ValueError('Unexpected type')


def _GetStructValue(struct_value):
  which = struct_value.WhichOneof('kind')
  if which == 'struct_value':
    return struct_value.struct_value
  elif which == 'null_value':
    return None
  elif which == 'number_value':
    return struct_value.number_value
  elif which == 'string_value':
    return struct_value.string_value
  elif which == 'bool_value':
    return struct_value.bool_value
  elif which == 'list_value':
    return struct_value.list_value
  elif which is None:
    raise ValueError('Value not set')


class Struct(object):
  """Class for Struct message type."""

  __slots__ = ()

  def __getitem__(self, key):
    return _GetStructValue(self.fields[key])

  def __contains__(self, item):
    return item in self.fields

  def __setitem__(self, key, value):
    _SetStructValue(self.fields[key], value)

  def __delitem__(self, key):
    del self.fields[key]

  def __len__(self):
    return len(self.fields)

  def __iter__(self):
    return iter(self.fields)

  def keys(self):  # pylint: disable=invalid-name
    return self.fields.keys()

  def values(self):  # pylint: disable=invalid-name
    return [self[key] for key in self]

  def items(self):  # pylint: disable=invalid-name
    return [(key, self[key]) for key in self]

  def get_or_create_list(self, key):
    """Returns a list for this key, creating if it didn't exist already."""
    if not self.fields[key].HasField('list_value'):
      # Clear will mark list_value modified which will indeed create a list.
      self.fields[key].list_value.Clear()
    return self.fields[key].list_value

  def get_or_create_struct(self, key):
    """Returns a struct for this key, creating if it didn't exist already."""
    if not self.fields[key].HasField('struct_value'):
      # Clear will mark struct_value modified which will indeed create a struct.
      self.fields[key].struct_value.Clear()
    return self.fields[key].struct_value

  def update(self, dictionary):  # pylint: disable=invalid-name
    for key, value in dictionary.items():
      _SetStructValue(self.fields[key], value)

collections.abc.MutableMapping.register(Struct)


class ListValue(object):
  """Class for ListValue message type."""

  __slots__ = ()

  def __len__(self):
    return len(self.values)

  def append(self, value):
    _SetStructValue(self.values.add(), value)

  def extend(self, elem_seq):
    for value in elem_seq:
      self.append(value)

  def __getitem__(self, index):
    """Retrieves item by the specified index."""
    return _GetStructValue(self.values.__getitem__(index))

  def __setitem__(self, index, value):
    _SetStructValue(self.values.__getitem__(index), value)

  def __delitem__(self, key):
    del self.values[key]

  def items(self):
    for i in range(len(self)):
      yield self[i]

  def add_struct(self):
    """Appends and returns a struct value as the next value in the list."""
    struct_value = self.values.add().struct_value
    # Clear will mark struct_value modified which will indeed create a struct.
    struct_value.Clear()
    return struct_value

  def add_list(self):
    """Appends and returns a list value as the next value in the list."""
    list_value = self.values.add().list_value
    # Clear will mark list_value modified which will indeed create a list.
    list_value.Clear()
    return list_value

collections.abc.MutableSequence.register(ListValue)


WKTBASES = {
    'google.protobuf.Any': Any,
    'google.protobuf.Duration': Duration,
    'google.protobuf.FieldMask': FieldMask,
    'google.protobuf.ListValue': ListValue,
    'google.protobuf.Struct': Struct,
    'google.protobuf.Timestamp': Timestamp,
}

social_vdgnn/lib/python3.6/site-packages/google/protobuf/internal/well_known_types.py
eliatrevisan/IA-MPPI-LBM
# Protocol Buffers - Google's data interchange format
# Copyright 2008 Google Inc.  All rights reserved.
# https://developers.google.com/protocol-buffers/
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Contains routines for printing protocol messages in JSON format.

Simple usage example:

  # Create a proto object and serialize it to a json format string.
  message = my_proto_pb2.MyMessage(foo='bar')
  json_string = json_format.MessageToJson(message)

  # Parse a json format string to proto object.
  message = json_format.Parse(json_string, my_proto_pb2.MyMessage())
"""

__author__ = 'jieluo@google.com (Jie Luo)'


import base64
from collections import OrderedDict
import json
import math
from operator import methodcaller
import re
import sys

from google.protobuf.internal import type_checkers
from google.protobuf import descriptor
from google.protobuf import symbol_database


_TIMESTAMPFOMAT = '%Y-%m-%dT%H:%M:%S'
_INT_TYPES = frozenset([descriptor.FieldDescriptor.CPPTYPE_INT32,
                        descriptor.FieldDescriptor.CPPTYPE_UINT32,
                        descriptor.FieldDescriptor.CPPTYPE_INT64,
                        descriptor.FieldDescriptor.CPPTYPE_UINT64])
_INT64_TYPES = frozenset([descriptor.FieldDescriptor.CPPTYPE_INT64,
                          descriptor.FieldDescriptor.CPPTYPE_UINT64])
_FLOAT_TYPES = frozenset([descriptor.FieldDescriptor.CPPTYPE_FLOAT,
                          descriptor.FieldDescriptor.CPPTYPE_DOUBLE])
_INFINITY = 'Infinity'
_NEG_INFINITY = '-Infinity'
_NAN = 'NaN'

_UNPAIRED_SURROGATE_PATTERN = re.compile(
    u'[\ud800-\udbff](?![\udc00-\udfff])|(?<![\ud800-\udbff])[\udc00-\udfff]')

_VALID_EXTENSION_NAME = re.compile(r'\[[a-zA-Z0-9\._]*\]$')


class Error(Exception):
  """Top-level module error for json_format."""


class SerializeToJsonError(Error):
  """Thrown if serialization to JSON fails."""


class ParseError(Error):
  """Thrown in case of parsing error."""


def MessageToJson(
    message,
    including_default_value_fields=False,
    preserving_proto_field_name=False,
    indent=2,
    sort_keys=False,
    use_integers_for_enums=False,
    descriptor_pool=None,
    float_precision=None):
  """Converts protobuf message to JSON format.

  Args:
    message: The protocol buffers message instance to serialize.
    including_default_value_fields: If True, singular primitive fields,
        repeated fields, and map fields will always be serialized.  If
        False, only serialize non-empty fields.  Singular message fields
        and oneof fields are not affected by this option.
    preserving_proto_field_name: If True, use the original proto field
        names as defined in the .proto file. If False, convert the field
        names to lowerCamelCase.
    indent: The JSON object will be pretty-printed with this indent level.
        An indent level of 0 or negative will only insert newlines.
    sort_keys: If True, then the output will be sorted by field names.
    use_integers_for_enums: If true, print integers instead of enum names.
    descriptor_pool: A Descriptor Pool for resolving types. If None use the
        default.
    float_precision: If set, use this to specify float field valid digits.

  Returns:
    A string containing the JSON formatted protocol buffer message.
  """
  printer = _Printer(
      including_default_value_fields,
      preserving_proto_field_name,
      use_integers_for_enums,
      descriptor_pool,
      float_precision=float_precision)
  return printer.ToJsonString(message, indent, sort_keys)


def MessageToDict(
    message,
    including_default_value_fields=False,
    preserving_proto_field_name=False,
    use_integers_for_enums=False,
    descriptor_pool=None,
    float_precision=None):
  """Converts protobuf message to a dictionary.

  When the dictionary is encoded to JSON, it conforms to proto3 JSON spec.

  Args:
    message: The protocol buffers message instance to serialize.
    including_default_value_fields: If True, singular primitive fields,
        repeated fields, and map fields will always be serialized.  If
        False, only serialize non-empty fields.  Singular message fields
        and oneof fields are not affected by this option.
    preserving_proto_field_name: If True, use the original proto field
        names as defined in the .proto file. If False, convert the field
        names to lowerCamelCase.
    use_integers_for_enums: If true, print integers instead of enum names.
    descriptor_pool: A Descriptor Pool for resolving types. If None use the
        default.
    float_precision: If set, use this to specify float field valid digits.

  Returns:
    A dict representation of the protocol buffer message.
  """
  printer = _Printer(
      including_default_value_fields,
      preserving_proto_field_name,
      use_integers_for_enums,
      descriptor_pool,
      float_precision=float_precision)
  # pylint: disable=protected-access
  return printer._MessageToJsonObject(message)


def _IsMapEntry(field):
  return (field.type == descriptor.FieldDescriptor.TYPE_MESSAGE and
          field.message_type.has_options and
          field.message_type.GetOptions().map_entry)


class _Printer(object):
  """JSON format printer for protocol message."""

  def __init__(
      self,
      including_default_value_fields=False,
      preserving_proto_field_name=False,
      use_integers_for_enums=False,
      descriptor_pool=None,
      float_precision=None):
    self.including_default_value_fields = including_default_value_fields
    self.preserving_proto_field_name = preserving_proto_field_name
    self.use_integers_for_enums = use_integers_for_enums
    self.descriptor_pool = descriptor_pool
    if float_precision:
      self.float_format = '.{}g'.format(float_precision)
    else:
      self.float_format = None

  def ToJsonString(self, message, indent, sort_keys):
    js = self._MessageToJsonObject(message)
    return json.dumps(js, indent=indent, sort_keys=sort_keys)

  def _MessageToJsonObject(self, message):
    """Converts message to an object according to Proto3 JSON Specification."""
    message_descriptor = message.DESCRIPTOR
    full_name = message_descriptor.full_name
    if _IsWrapperMessage(message_descriptor):
      return self._WrapperMessageToJsonObject(message)
    if full_name in _WKTJSONMETHODS:
      return methodcaller(_WKTJSONMETHODS[full_name][0], message)(self)
    js = {}
    return self._RegularMessageToJsonObject(message, js)

  def _RegularMessageToJsonObject(self, message, js):
    """Converts normal message according to Proto3 JSON Specification."""
    fields = message.ListFields()

    try:
      for field, value in fields:
        if self.preserving_proto_field_name:
          name = field.name
        else:
          name = field.json_name
        if _IsMapEntry(field):
          # Convert a map field.
          v_field = field.message_type.fields_by_name['value']
          js_map = {}
          for key in value:
            if isinstance(key, bool):
              if key:
                recorded_key = 'true'
              else:
                recorded_key = 'false'
            else:
              recorded_key = str(key)
            js_map[recorded_key] = self._FieldToJsonObject(
                v_field, value[key])
          js[name] = js_map
        elif field.label == descriptor.FieldDescriptor.LABEL_REPEATED:
          # Convert a repeated field.
          js[name] = [self._FieldToJsonObject(field, k)
                      for k in value]
        elif field.is_extension:
          name = '[%s]' % field.full_name
          js[name] = self._FieldToJsonObject(field, value)
        else:
          js[name] = self._FieldToJsonObject(field, value)

      # Serialize default value if including_default_value_fields is True.
      if self.including_default_value_fields:
        message_descriptor = message.DESCRIPTOR
        for field in message_descriptor.fields:
          # Singular message fields and oneof fields will not be affected.
          if ((field.label != descriptor.FieldDescriptor.LABEL_REPEATED and
               field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_MESSAGE) or
              field.containing_oneof):
            continue
          if self.preserving_proto_field_name:
            name = field.name
          else:
            name = field.json_name
          if name in js:
            # Skip the field which has been serialized already.
            continue
          if _IsMapEntry(field):
            js[name] = {}
          elif field.label == descriptor.FieldDescriptor.LABEL_REPEATED:
            js[name] = []
          else:
            js[name] = self._FieldToJsonObject(field, field.default_value)

    except ValueError as e:
      raise SerializeToJsonError(
          'Failed to serialize {0} field: {1}.'.format(field.name, e))

    return js

  def _FieldToJsonObject(self, field, value):
    """Converts field value according to Proto3 JSON Specification."""
    if field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_MESSAGE:
      return self._MessageToJsonObject(value)
    elif field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_ENUM:
      if self.use_integers_for_enums:
        return value
      if field.enum_type.full_name == 'google.protobuf.NullValue':
        return None
      enum_value = field.enum_type.values_by_number.get(value, None)
      if enum_value is not None:
        return enum_value.name
      else:
        if field.file.syntax == 'proto3':
          return value
        raise SerializeToJsonError('Enum field contains an integer value '
                                   'which can not mapped to an enum value.')
    elif field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_STRING:
      if field.type == descriptor.FieldDescriptor.TYPE_BYTES:
        # Use base64 Data encoding for bytes
        return base64.b64encode(value).decode('utf-8')
      else:
        return value
    elif field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_BOOL:
      return bool(value)
    elif field.cpp_type in _INT64_TYPES:
      return str(value)
    elif field.cpp_type in _FLOAT_TYPES:
      if math.isinf(value):
        if value < 0.0:
          return _NEG_INFINITY
        else:
          return _INFINITY
      if math.isnan(value):
        return _NAN
      if field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_FLOAT:
        if self.float_format:
          return float(format(value, self.float_format))
        else:
          return type_checkers.ToShortestFloat(value)

    return value

  def _AnyMessageToJsonObject(self, message):
    """Converts Any message according to Proto3 JSON Specification."""
    if not message.ListFields():
      return {}
    # Must print @type first, use OrderedDict instead of {}
    js = OrderedDict()
    type_url = message.type_url
    js['@type'] = type_url
    sub_message = _CreateMessageFromTypeUrl(type_url, self.descriptor_pool)
    sub_message.ParseFromString(message.value)
    message_descriptor = sub_message.DESCRIPTOR
    full_name = message_descriptor.full_name
    if _IsWrapperMessage(message_descriptor):
      js['value'] = self._WrapperMessageToJsonObject(sub_message)
      return js
    if full_name in _WKTJSONMETHODS:
      js['value'] = methodcaller(_WKTJSONMETHODS[full_name][0],
                                 sub_message)(self)
      return js
    return self._RegularMessageToJsonObject(sub_message, js)

  def _GenericMessageToJsonObject(self, message):
    """Converts message according to Proto3 JSON Specification."""
    # Duration, Timestamp and FieldMask have ToJsonString method to do the
    # convert. Users can also call the method directly.
    return message.ToJsonString()

  def _ValueMessageToJsonObject(self, message):
    """Converts Value message according to Proto3 JSON Specification."""
    which = message.WhichOneof('kind')
    # If the Value message is not set treat as null_value when serialize
    # to JSON. The parse back result will be different from original message.
    if which is None or which == 'null_value':
      return None
    if which == 'list_value':
      return self._ListValueMessageToJsonObject(message.list_value)
    if which == 'struct_value':
      value = message.struct_value
    else:
      value = getattr(message, which)
    oneof_descriptor = message.DESCRIPTOR.fields_by_name[which]
    return self._FieldToJsonObject(oneof_descriptor, value)

  def _ListValueMessageToJsonObject(self, message):
    """Converts ListValue message according to Proto3 JSON Specification."""
    return [self._ValueMessageToJsonObject(value)
            for value in message.values]

  def _StructMessageToJsonObject(self, message):
    """Converts Struct message according to Proto3 JSON Specification."""
    fields = message.fields
    ret = {}
    for key in fields:
      ret[key] = self._ValueMessageToJsonObject(fields[key])
    return ret

  def _WrapperMessageToJsonObject(self, message):
    return self._FieldToJsonObject(
        message.DESCRIPTOR.fields_by_name['value'], message.value)


def _IsWrapperMessage(message_descriptor):
  return message_descriptor.file.name == 'google/protobuf/wrappers.proto'


def _DuplicateChecker(js):
  result = {}
  for name, value in js:
    if name in result:
      raise ParseError('Failed to load JSON: duplicate key {0}.'.format(name))
    result[name] = value
  return result


def _CreateMessageFromTypeUrl(type_url, descriptor_pool):
  """Creates a message from a type URL."""
  db = symbol_database.Default()
  pool = db.pool if descriptor_pool is None else descriptor_pool
  type_name = type_url.split('/')[-1]
  try:
    message_descriptor = pool.FindMessageTypeByName(type_name)
  except KeyError:
    raise TypeError(
        'Can not find message descriptor by type_url: {0}.'.format(type_url))
  message_class = db.GetPrototype(message_descriptor)
  return message_class()


def Parse(text, message, ignore_unknown_fields=False, descriptor_pool=None):
  """Parses a JSON representation of a protocol message into a message.

  Args:
    text: Message JSON representation.
    message: A protocol buffer message to merge into.
    ignore_unknown_fields: If True, do not raise errors for unknown fields.
    descriptor_pool: A Descriptor Pool for resolving types. If None use the
        default.

  Returns:
    The same message passed as argument.

  Raises::
    ParseError: On JSON parsing problems.
  """
  if not isinstance(text, str):
    text = text.decode('utf-8')
  try:
    js = json.loads(text, object_pairs_hook=_DuplicateChecker)
  except ValueError as e:
    raise ParseError('Failed to load JSON: {0}.'.format(str(e)))
  return ParseDict(js, message, ignore_unknown_fields, descriptor_pool)


def ParseDict(js_dict,
              message,
              ignore_unknown_fields=False,
              descriptor_pool=None):
  """Parses a JSON dictionary representation into a message.

  Args:
    js_dict: Dict representation of a JSON message.
    message: A protocol buffer message to merge into.
    ignore_unknown_fields: If True, do not raise errors for unknown fields.
    descriptor_pool: A Descriptor Pool for resolving types. If None use the
      default.

  Returns:
    The same message passed as argument.
  """
  parser = _Parser(ignore_unknown_fields, descriptor_pool)
  parser.ConvertMessage(js_dict, message)
  return message


_INT_OR_FLOAT = (int, float)


class _Parser(object):
  """JSON format parser for protocol message."""

  def __init__(self, ignore_unknown_fields, descriptor_pool):
    self.ignore_unknown_fields = ignore_unknown_fields
    self.descriptor_pool = descriptor_pool

  def ConvertMessage(self, value, message):
    """Convert a JSON object into a message.

    Args:
      value: A JSON object.
      message: A WKT or regular protocol message to record the data.

    Raises:
      ParseError: In case of convert problems.
    """
    message_descriptor = message.DESCRIPTOR
    full_name = message_descriptor.full_name
    if _IsWrapperMessage(message_descriptor):
      self._ConvertWrapperMessage(value, message)
    elif full_name in _WKTJSONMETHODS:
      methodcaller(_WKTJSONMETHODS[full_name][1], value, message)(self)
    else:
      self._ConvertFieldValuePair(value, message)

  def _ConvertFieldValuePair(self, js, message):
    """Convert field value pairs into regular message.

    Args:
      js: A JSON object to convert the field value pairs.
      message: A regular protocol message to record the data.

    Raises:
      ParseError: In case of problems converting.
    """
    names = []
    message_descriptor = message.DESCRIPTOR
    fields_by_json_name = dict((f.json_name, f)
                               for f in message_descriptor.fields)
    for name in js:
      try:
        field = fields_by_json_name.get(name, None)
        if not field:
          field = message_descriptor.fields_by_name.get(name, None)
        if not field and _VALID_EXTENSION_NAME.match(name):
          if not message_descriptor.is_extendable:
            raise ParseError('Message type {0} does not have extensions'.format(
                message_descriptor.full_name))
          identifier = name[1:-1]  # strip [] brackets
          # pylint: disable=protected-access
          field = message.Extensions._FindExtensionByName(identifier)
          # pylint: enable=protected-access
          if not field:
            # Try looking for extension by the message type name, dropping the
            # field name following the final . separator in full_name.
            identifier = '.'.join(identifier.split('.')[:-1])
            # pylint: disable=protected-access
            field = message.Extensions._FindExtensionByName(identifier)
            # pylint: enable=protected-access
        if not field:
          if self.ignore_unknown_fields:
            continue
          raise ParseError(
              ('Message type "{0}" has no field named "{1}".\n'
               ' Available Fields(except extensions): {2}').format(
                   message_descriptor.full_name, name,
                   [f.json_name for f in message_descriptor.fields]))
        if name in names:
          raise ParseError('Message type "{0}" should not have multiple '
                           '"{1}" fields.'.format(
                               message.DESCRIPTOR.full_name, name))
        names.append(name)
        value = js[name]
        # Check no other oneof field is parsed.
        if field.containing_oneof is not None and value is not None:
          oneof_name = field.containing_oneof.name
          if oneof_name in names:
            raise ParseError('Message type "{0}" should not have multiple '
                             '"{1}" oneof fields.'.format(
                                 message.DESCRIPTOR.full_name, oneof_name))
          names.append(oneof_name)

        if value is None:
          if (field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_MESSAGE
              and field.message_type.full_name == 'google.protobuf.Value'):
            sub_message = getattr(message, field.name)
            sub_message.null_value = 0
          elif (field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_ENUM
                and field.enum_type.full_name == 'google.protobuf.NullValue'):
            setattr(message, field.name, 0)
          else:
            message.ClearField(field.name)
          continue

        # Parse field value.
        if _IsMapEntry(field):
          message.ClearField(field.name)
          self._ConvertMapFieldValue(value, message, field)
        elif field.label == descriptor.FieldDescriptor.LABEL_REPEATED:
          message.ClearField(field.name)
          if not isinstance(value, list):
            raise ParseError('repeated field {0} must be in [] which is '
                             '{1}.'.format(name, value))
          if field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_MESSAGE:
            # Repeated message field.
            for item in value:
              sub_message = getattr(message, field.name).add()
              # None is a null_value in Value.
              if (item is None and
                  sub_message.DESCRIPTOR.full_name != 'google.protobuf.Value'):
                raise ParseError('null is not allowed to be used as an element'
                                 ' in a repeated field.')
              self.ConvertMessage(item, sub_message)
          else:
            # Repeated scalar field.
            for item in value:
              if item is None:
                raise ParseError('null is not allowed to be used as an element'
                                 ' in a repeated field.')
              getattr(message, field.name).append(
                  _ConvertScalarFieldValue(item, field))
        elif field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_MESSAGE:
          if field.is_extension:
            sub_message = message.Extensions[field]
          else:
            sub_message = getattr(message, field.name)
          sub_message.SetInParent()
          self.ConvertMessage(value, sub_message)
        else:
          if field.is_extension:
            message.Extensions[field] = _ConvertScalarFieldValue(value, field)
          else:
            setattr(message, field.name, _ConvertScalarFieldValue(value, field))
      except ParseError as e:
        if field and field.containing_oneof is None:
          raise ParseError('Failed to parse {0} field: {1}.'.format(name, e))
        else:
          raise ParseError(str(e))
      except ValueError as e:
        raise ParseError('Failed to parse {0} field: {1}.'.format(name, e))
      except TypeError as e:
        raise ParseError('Failed to parse {0} field: {1}.'.format(name, e))

  def _ConvertAnyMessage(self, value, message):
    """Convert a JSON representation into Any message."""
    if isinstance(value, dict) and not value:
      return
    try:
      type_url = value['@type']
    except KeyError:
      raise ParseError('@type is missing when parsing any message.')

    sub_message = _CreateMessageFromTypeUrl(type_url, self.descriptor_pool)
    message_descriptor = sub_message.DESCRIPTOR
    full_name = message_descriptor.full_name
    if _IsWrapperMessage(message_descriptor):
      self._ConvertWrapperMessage(value['value'], sub_message)
    elif full_name in _WKTJSONMETHODS:
      methodcaller(
          _WKTJSONMETHODS[full_name][1], value['value'], sub_message)(self)
    else:
      del value['@type']
      self._ConvertFieldValuePair(value, sub_message)
      value['@type'] = type_url
    # Sets Any message
    message.value = sub_message.SerializeToString()
    message.type_url = type_url

  def _ConvertGenericMessage(self, value, message):
    """Convert a JSON representation into message with FromJsonString."""
    # Duration, Timestamp, FieldMask have a FromJsonString method to do the
    # conversion. Users can also call the method directly.
    try:
      message.FromJsonString(value)
    except ValueError as e:
      raise ParseError(e)

  def _ConvertValueMessage(self, value, message):
    """Convert a JSON representation into Value message."""
    if isinstance(value, dict):
      self._ConvertStructMessage(value, message.struct_value)
    elif isinstance(value, list):
      self. _ConvertListValueMessage(value, message.list_value)
    elif value is None:
      message.null_value = 0
    elif isinstance(value, bool):
      message.bool_value = value
    elif isinstance(value, str):
      message.string_value = value
    elif isinstance(value, _INT_OR_FLOAT):
      message.number_value = value
    else:
      raise ParseError('Value {0} has unexpected type {1}.'.format(
          value, type(value)))

  def _ConvertListValueMessage(self, value, message):
    """Convert a JSON representation into ListValue message."""
    if not isinstance(value, list):
      raise ParseError(
          'ListValue must be in [] which is {0}.'.format(value))
    message.ClearField('values')
    for item in value:
      self._ConvertValueMessage(item, message.values.add())

  def _ConvertStructMessage(self, value, message):
    """Convert a JSON representation into Struct message."""
    if not isinstance(value, dict):
      raise ParseError(
          'Struct must be in a dict which is {0}.'.format(value))
    # Clear will mark the struct as modified so it will be created even if
    # there are no values.
    message.Clear()
    for key in value:
      self._ConvertValueMessage(value[key], message.fields[key])
    return

  def _ConvertWrapperMessage(self, value, message):
    """Convert a JSON representation into Wrapper message."""
    field = message.DESCRIPTOR.fields_by_name['value']
    setattr(message, 'value', _ConvertScalarFieldValue(value, field))

  def _ConvertMapFieldValue(self, value, message, field):
    """Convert map field value for a message map field.

    Args:
      value: A JSON object to convert the map field value.
      message: A protocol message to record the converted data.
      field: The descriptor of the map field to be converted.

    Raises:
      ParseError: In case of convert problems.
    """
    if not isinstance(value, dict):
      raise ParseError(
          'Map field {0} must be in a dict which is {1}.'.format(
              field.name, value))
    key_field = field.message_type.fields_by_name['key']
    value_field = field.message_type.fields_by_name['value']
    for key in value:
      key_value = _ConvertScalarFieldValue(key, key_field, True)
      if value_field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_MESSAGE:
        self.ConvertMessage(value[key], getattr(
            message, field.name)[key_value])
      else:
        getattr(message, field.name)[key_value] = _ConvertScalarFieldValue(
            value[key], value_field)


def _ConvertScalarFieldValue(value, field, require_str=False):
  """Convert a single scalar field value.

  Args:
    value: A scalar value to convert the scalar field value.
    field: The descriptor of the field to convert.
    require_str: If True, the field value must be a str.

  Returns:
    The converted scalar field value

  Raises:
    ParseError: In case of convert problems.
  """
  if field.cpp_type in _INT_TYPES:
    return _ConvertInteger(value)
  elif field.cpp_type in _FLOAT_TYPES:
    return _ConvertFloat(value, field)
  elif field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_BOOL:
    return _ConvertBool(value, require_str)
  elif field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_STRING:
    if field.type == descriptor.FieldDescriptor.TYPE_BYTES:
      if isinstance(value, str):
        encoded = value.encode('utf-8')
      else:
        encoded = value
      # Add extra padding '='
      padded_value = encoded + b'=' * (4 - len(encoded) % 4)
      return base64.urlsafe_b64decode(padded_value)
    else:
      # Checking for unpaired surrogates appears to be unreliable,
      # depending on the specific Python version, so we check manually.
      if _UNPAIRED_SURROGATE_PATTERN.search(value):
        raise ParseError('Unpaired surrogate')
      return value
  elif field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_ENUM:
    # Convert an enum value.
    enum_value = field.enum_type.values_by_name.get(value, None)
    if enum_value is None:
      try:
        number = int(value)
        enum_value = field.enum_type.values_by_number.get(number, None)
      except ValueError:
        raise ParseError('Invalid enum value {0} for enum type {1}.'.format(
            value, field.enum_type.full_name))
      if enum_value is None:
        if field.file.syntax == 'proto3':
          # Proto3 accepts unknown enums.
          return number
        raise ParseError('Invalid enum value {0} for enum type {1}.'.format(
            value, field.enum_type.full_name))
    return enum_value.number


def _ConvertInteger(value):
  """Convert an integer.

  Args:
    value: A scalar value to convert.

  Returns:
    The integer value.

  Raises:
    ParseError: If an integer couldn't be consumed.
  """
  if isinstance(value, float) and not value.is_integer():
    raise ParseError('Couldn\'t parse integer: {0}.'.format(value))

  if isinstance(value, str) and value.find(' ') != -1:
    raise ParseError('Couldn\'t parse integer: "{0}".'.format(value))

  if isinstance(value, bool):
    raise ParseError('Bool value {0} is not acceptable for '
                     'integer field.'.format(value))

  return int(value)


def _ConvertFloat(value, field):
  """Convert an floating point number."""
  if isinstance(value, float):
    if math.isnan(value):
      raise ParseError('Couldn\'t parse NaN, use quoted "NaN" instead.')
    if math.isinf(value):
      if value > 0:
        raise ParseError('Couldn\'t parse Infinity or value too large, '
                         'use quoted "Infinity" instead.')
      else:
        raise ParseError('Couldn\'t parse -Infinity or value too small, '
                         'use quoted "-Infinity" instead.')
    if field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_FLOAT:
      # pylint: disable=protected-access
      if value > type_checkers._FLOAT_MAX:
        raise ParseError('Float value too large')
      # pylint: disable=protected-access
      if value < type_checkers._FLOAT_MIN:
        raise ParseError('Float value too small')
  if value == 'nan':
    raise ParseError('Couldn\'t parse float "nan", use "NaN" instead.')
  try:
    # Assume Python compatible syntax.
    return float(value)
  except ValueError:
    # Check alternative spellings.
    if value == _NEG_INFINITY:
      return float('-inf')
    elif value == _INFINITY:
      return float('inf')
    elif value == _NAN:
      return float('nan')
    else:
      raise ParseError('Couldn\'t parse float: {0}.'.format(value))


def _ConvertBool(value, require_str):
  """Convert a boolean value.

  Args:
    value: A scalar value to convert.
    require_str: If True, value must be a str.

  Returns:
    The bool parsed.

  Raises:
    ParseError: If a boolean value couldn't be consumed.
  """
  if require_str:
    if value == 'true':
      return True
    elif value == 'false':
      return False
    else:
      raise ParseError('Expected "true" or "false", not {0}.'.format(value))

  if not isinstance(value, bool):
    raise ParseError('Expected true or false without quotes.')
  return value

_WKTJSONMETHODS = {
    'google.protobuf.Any': ['_AnyMessageToJsonObject',
                            '_ConvertAnyMessage'],
    'google.protobuf.Duration': ['_GenericMessageToJsonObject',
                                 '_ConvertGenericMessage'],
    'google.protobuf.FieldMask': ['_GenericMessageToJsonObject',
                                  '_ConvertGenericMessage'],
    'google.protobuf.ListValue': ['_ListValueMessageToJsonObject',
                                  '_ConvertListValueMessage'],
    'google.protobuf.Struct': ['_StructMessageToJsonObject',
                               '_ConvertStructMessage'],
    'google.protobuf.Timestamp': ['_GenericMessageToJsonObject',
                                  '_ConvertGenericMessage'],
    'google.protobuf.Value': ['_ValueMessageToJsonObject',
                              '_ConvertValueMessage']
}

social_vdgnn/lib/python3.6/site-packages/google/protobuf/json_format.py
eliatrevisan/IA-MPPI-LBM
# Protocol Buffers - Google's data interchange format
# Copyright 2008 Google Inc.  All rights reserved.
# https://developers.google.com/protocol-buffers/
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# TODO(robinson): We should just make these methods all "pure-virtual" and move
# all implementation out, into reflection.py for now.


"""Contains an abstract base class for protocol messages."""

__author__ = 'robinson@google.com (Will Robinson)'

class Error(Exception):
  """Base error type for this module."""
  pass


class DecodeError(Error):
  """Exception raised when deserializing messages."""
  pass


class EncodeError(Error):
  """Exception raised when serializing messages."""
  pass


class Message(object):

  """Abstract base class for protocol messages.

  Protocol message classes are almost always generated by the protocol
  compiler.  These generated types subclass Message and implement the methods
  shown below.
  """

  # TODO(robinson): Link to an HTML document here.

  # TODO(robinson): Document that instances of this class will also
  # have an Extensions attribute with __getitem__ and __setitem__.
  # Again, not sure how to best convey this.

  # TODO(robinson): Document that the class must also have a static
  #   RegisterExtension(extension_field) method.
  #   Not sure how to best express at this point.

  # TODO(robinson): Document these fields and methods.

  __slots__ = []

  #: The :class:`google.protobuf.descriptor.Descriptor` for this message type.
  DESCRIPTOR = None

  def __deepcopy__(self, memo=None):
    clone = type(self)()
    clone.MergeFrom(self)
    return clone

  def __eq__(self, other_msg):
    """Recursively compares two messages by value and structure."""
    raise NotImplementedError

  def __ne__(self, other_msg):
    # Can't just say self != other_msg, since that would infinitely recurse. :)
    return not self == other_msg

  def __hash__(self):
    raise TypeError('unhashable object')

  def __str__(self):
    """Outputs a human-readable representation of the message."""
    raise NotImplementedError

  def __unicode__(self):
    """Outputs a human-readable representation of the message."""
    raise NotImplementedError

  def MergeFrom(self, other_msg):
    """Merges the contents of the specified message into current message.

    This method merges the contents of the specified message into the current
    message. Singular fields that are set in the specified message overwrite
    the corresponding fields in the current message. Repeated fields are
    appended. Singular sub-messages and groups are recursively merged.

    Args:
      other_msg (Message): A message to merge into the current message.
    """
    raise NotImplementedError

  def CopyFrom(self, other_msg):
    """Copies the content of the specified message into the current message.

    The method clears the current message and then merges the specified
    message using MergeFrom.

    Args:
      other_msg (Message): A message to copy into the current one.
    """
    if self is other_msg:
      return
    self.Clear()
    self.MergeFrom(other_msg)

  def Clear(self):
    """Clears all data that was set in the message."""
    raise NotImplementedError

  def SetInParent(self):
    """Mark this as present in the parent.

    This normally happens automatically when you assign a field of a
    sub-message, but sometimes you want to make the sub-message
    present while keeping it empty.  If you find yourself using this,
    you may want to reconsider your design.
    """
    raise NotImplementedError

  def IsInitialized(self):
    """Checks if the message is initialized.

    Returns:
      bool: The method returns True if the message is initialized (i.e. all of
      its required fields are set).
    """
    raise NotImplementedError

  # TODO(robinson): MergeFromString() should probably return None and be
  # implemented in terms of a helper that returns the # of bytes read.  Our
  # deserialization routines would use the helper when recursively
  # deserializing, but the end user would almost always just want the no-return
  # MergeFromString().

  def MergeFromString(self, serialized):
    """Merges serialized protocol buffer data into this message.

    When we find a field in `serialized` that is already present
    in this message:

    -   If it's a "repeated" field, we append to the end of our list.
    -   Else, if it's a scalar, we overwrite our field.
    -   Else, (it's a nonrepeated composite), we recursively merge
        into the existing composite.

    Args:
      serialized (bytes): Any object that allows us to call
        ``memoryview(serialized)`` to access a string of bytes using the
        buffer interface.

    Returns:
      int: The number of bytes read from `serialized`.
      For non-group messages, this will always be `len(serialized)`,
      but for messages which are actually groups, this will
      generally be less than `len(serialized)`, since we must
      stop when we reach an ``END_GROUP`` tag.  Note that if
      we *do* stop because of an ``END_GROUP`` tag, the number
      of bytes returned does not include the bytes
      for the ``END_GROUP`` tag information.

    Raises:
      DecodeError: if the input cannot be parsed.
    """
    # TODO(robinson): Document handling of unknown fields.
    # TODO(robinson): When we switch to a helper, this will return None.
    raise NotImplementedError

  def ParseFromString(self, serialized):
    """Parse serialized protocol buffer data into this message.

    Like :func:`MergeFromString()`, except we clear the object first.
    """
    self.Clear()
    return self.MergeFromString(serialized)

  def SerializeToString(self, **kwargs):
    """Serializes the protocol message to a binary string.

    Keyword Args:
      deterministic (bool): If true, requests deterministic serialization
        of the protobuf, with predictable ordering of map keys.

    Returns:
      A binary string representation of the message if all of the required
      fields in the message are set (i.e. the message is initialized).

    Raises:
      EncodeError: if the message isn't initialized (see :func:`IsInitialized`).
    """
    raise NotImplementedError

  def SerializePartialToString(self, **kwargs):
    """Serializes the protocol message to a binary string.

    This method is similar to SerializeToString but doesn't check if the
    message is initialized.

    Keyword Args:
      deterministic (bool): If true, requests deterministic serialization
        of the protobuf, with predictable ordering of map keys.

    Returns:
      bytes: A serialized representation of the partial message.
    """
    raise NotImplementedError

  # TODO(robinson): Decide whether we like these better
  # than auto-generated has_foo() and clear_foo() methods
  # on the instances themselves.  This way is less consistent
  # with C++, but it makes reflection-type access easier and
  # reduces the number of magically autogenerated things.
  #
  # TODO(robinson): Be sure to document (and test) exactly
  # which field names are accepted here.  Are we case-sensitive?
  # What do we do with fields that share names with Python keywords
  # like 'lambda' and 'yield'?
  #
  # nnorwitz says:
  # """
  # Typically (in python), an underscore is appended to names that are
  # keywords. So they would become lambda_ or yield_.
  # """
  def ListFields(self):
    """Returns a list of (FieldDescriptor, value) tuples for present fields.

    A message field is non-empty if HasField() would return true. A singular
    primitive field is non-empty if HasField() would return true in proto2 or it
    is non zero in proto3. A repeated field is non-empty if it contains at least
    one element. The fields are ordered by field number.

    Returns:
      list[tuple(FieldDescriptor, value)]: field descriptors and values
      for all fields in the message which are not empty. The values vary by
      field type.
    """
    raise NotImplementedError

  def HasField(self, field_name):
    """Checks if a certain field is set for the message.

    For a oneof group, checks if any field inside is set. Note that if the
    field_name is not defined in the message descriptor, :exc:`ValueError` will
    be raised.

    Args:
      field_name (str): The name of the field to check for presence.

    Returns:
      bool: Whether a value has been set for the named field.

    Raises:
      ValueError: if the `field_name` is not a member of this message.
    """
    raise NotImplementedError

  def ClearField(self, field_name):
    """Clears the contents of a given field.

    Inside a oneof group, clears the field set. If the name neither refers to a
    defined field or oneof group, :exc:`ValueError` is raised.

    Args:
      field_name (str): The name of the field to check for presence.

    Raises:
      ValueError: if the `field_name` is not a member of this message.
    """
    raise NotImplementedError

  def WhichOneof(self, oneof_group):
    """Returns the name of the field that is set inside a oneof group.

    If no field is set, returns None.

    Args:
      oneof_group (str): the name of the oneof group to check.

    Returns:
      str or None: The name of the group that is set, or None.

    Raises:
      ValueError: no group with the given name exists
    """
    raise NotImplementedError

  def HasExtension(self, extension_handle):
    """Checks if a certain extension is present for this message.

    Extensions are retrieved using the :attr:`Extensions` mapping (if present).

    Args:
      extension_handle: The handle for the extension to check.

    Returns:
      bool: Whether the extension is present for this message.

    Raises:
      KeyError: if the extension is repeated. Similar to repeated fields,
        there is no separate notion of presence: a "not present" repeated
        extension is an empty list.
    """
    raise NotImplementedError

  def ClearExtension(self, extension_handle):
    """Clears the contents of a given extension.

    Args:
      extension_handle: The handle for the extension to clear.
    """
    raise NotImplementedError

  def UnknownFields(self):
    """Returns the UnknownFieldSet.

    Returns:
      UnknownFieldSet: The unknown fields stored in this message.
    """
    raise NotImplementedError

  def DiscardUnknownFields(self):
    """Clears all fields in the :class:`UnknownFieldSet`.

    This operation is recursive for nested message.
    """
    raise NotImplementedError

  def ByteSize(self):
    """Returns the serialized size of this message.

    Recursively calls ByteSize() on all contained messages.

    Returns:
      int: The number of bytes required to serialize this message.
    """
    raise NotImplementedError

  @classmethod
  def FromString(cls, s):
    raise NotImplementedError

  @staticmethod
  def RegisterExtension(extension_handle):
    raise NotImplementedError

  def _SetListener(self, message_listener):
    """Internal method used by the protocol message implementation.
    Clients should not call this directly.

    Sets a listener that this message will call on certain state transitions.

    The purpose of this method is to register back-edges from children to
    parents at runtime, for the purpose of setting "has" bits and
    byte-size-dirty bits in the parent and ancestor objects whenever a child or
    descendant object is modified.

    If the client wants to disconnect this Message from the object tree, she
    explicitly sets callback to None.

    If message_listener is None, unregisters any existing listener.  Otherwise,
    message_listener must implement the MessageListener interface in
    internal/message_listener.py, and we discard any listener registered
    via a previous _SetListener() call.
    """
    raise NotImplementedError

  def __getstate__(self):
    """Support the pickle protocol."""
    return dict(serialized=self.SerializePartialToString())

  def __setstate__(self, state):
    """Support the pickle protocol."""
    self.__init__()
    serialized = state['serialized']
    # On Python 3, using encoding='latin1' is required for unpickling
    # protos pickled by Python 2.
    if not isinstance(serialized, bytes):
      serialized = serialized.encode('latin1')
    self.ParseFromString(serialized)

  def __reduce__(self):
    message_descriptor = self.DESCRIPTOR
    if message_descriptor.containing_type is None:
      return type(self), (), self.__getstate__()
    # the message type must be nested.
    # Python does not pickle nested classes; use the symbol_database on the
    # receiving end.
    container = message_descriptor
    return (_InternalConstructMessage, (container.full_name,),
            self.__getstate__())


def _InternalConstructMessage(full_name):
  """Constructs a nested message."""
  from google.protobuf import symbol_database  # pylint:disable=g-import-not-at-top

  return symbol_database.Default().GetSymbol(full_name)()

social_vdgnn/lib/python3.6/site-packages/google/protobuf/message.py
eliatrevisan/IA-MPPI-LBM
"""
Python Markdown

A Python implementation of John Gruber's Markdown.

Documentation: https://python-markdown.github.io/
GitHub: https://github.com/Python-Markdown/markdown/
PyPI: https://pypi.org/project/Markdown/

Started by Manfred Stienstra (http://www.dwerg.net/).
Maintained for a few years by Yuri Takhteyev (http://www.freewisdom.org).
Currently maintained by Waylan Limberg (https://github.com/waylan),
Dmitry Shachnev (https://github.com/mitya57) and Isaac Muse (https://github.com/facelessuser).

Copyright 2007-2018 The Python Markdown Project (v. 1.7 and later)
Copyright 2004, 2005, 2006 Yuri Takhteyev (v. 0.2-1.6b)
Copyright 2004 Manfred Stienstra (the original version)

License: BSD (see LICENSE.md for details).
"""

import codecs
import sys
import logging
import importlib
from . import util
from .preprocessors import build_preprocessors
from .blockprocessors import build_block_parser
from .treeprocessors import build_treeprocessors
from .inlinepatterns import build_inlinepatterns
from .postprocessors import build_postprocessors
from .extensions import Extension
from .serializers import to_html_string, to_xhtml_string

__all__ = ['Markdown', 'markdown', 'markdownFromFile']


logger = logging.getLogger('MARKDOWN')


class Markdown:
    """Convert Markdown to HTML."""

    doc_tag = "div"     # Element used to wrap document - later removed

    output_formats = {
        'html':   to_html_string,
        'xhtml':  to_xhtml_string,
    }

    def __init__(self, **kwargs):
        """
        Creates a new Markdown instance.

        Keyword arguments:

        * extensions: A list of extensions.
            If an item is an instance of a subclass of `markdown.extension.Extension`, the  instance will be used
            as-is. If an item is of type string, first an entry point will be loaded. If that fails, the string is
            assumed to use Python dot notation (`path.to.module:ClassName`) to load a markdown.Extension subclass. If
            no class is specified, then a `makeExtension` function is called within the specified module.
        * extension_configs: Configuration settings for extensions.
        * output_format: Format of output. Supported formats are:
            * "xhtml": Outputs XHTML style tags. Default.
            * "html": Outputs HTML style tags.
        * tab_length: Length of tabs in the source. Default: 4

        """

        self.tab_length = kwargs.get('tab_length', 4)

        self.ESCAPED_CHARS = ['\\', '`', '*', '_', '{', '}', '[', ']',
                              '(', ')', '>', '#', '+', '-', '.', '!']

        self.block_level_elements = [
            # Elements which are invalid to wrap in a `<p>` tag.
            # See https://w3c.github.io/html/grouping-content.html#the-p-element
            'address', 'article', 'aside', 'blockquote', 'details', 'div', 'dl',
            'fieldset', 'figcaption', 'figure', 'footer', 'form', 'h1', 'h2', 'h3',
            'h4', 'h5', 'h6', 'header', 'hgroup', 'hr', 'main', 'menu', 'nav', 'ol',
            'p', 'pre', 'section', 'table', 'ul',
            # Other elements which Markdown should not be mucking up the contents of.
            'canvas', 'colgroup', 'dd', 'body', 'dt', 'group', 'iframe', 'li', 'legend',
            'math', 'map', 'noscript', 'output', 'object', 'option', 'progress', 'script',
            'style', 'summary', 'tbody', 'td', 'textarea', 'tfoot', 'th', 'thead', 'tr', 'video'
        ]

        self.registeredExtensions = []
        self.docType = ""
        self.stripTopLevelTags = True

        self.build_parser()

        self.references = {}
        self.htmlStash = util.HtmlStash()
        self.registerExtensions(extensions=kwargs.get('extensions', []),
                                configs=kwargs.get('extension_configs', {}))
        self.set_output_format(kwargs.get('output_format', 'xhtml'))
        self.reset()

    def build_parser(self):
        """ Build the parser from the various parts. """
        self.preprocessors = build_preprocessors(self)
        self.parser = build_block_parser(self)
        self.inlinePatterns = build_inlinepatterns(self)
        self.treeprocessors = build_treeprocessors(self)
        self.postprocessors = build_postprocessors(self)
        return self

    def registerExtensions(self, extensions, configs):
        """
        Register extensions with this instance of Markdown.

        Keyword arguments:

        * extensions: A list of extensions, which can either
           be strings or objects.
        * configs: A dictionary mapping extension names to config options.

        """
        for ext in extensions:
            if isinstance(ext, str):
                ext = self.build_extension(ext, configs.get(ext, {}))
            if isinstance(ext, Extension):
                ext._extendMarkdown(self)
                logger.debug(
                    'Successfully loaded extension "%s.%s".'
                    % (ext.__class__.__module__, ext.__class__.__name__)
                )
            elif ext is not None:
                raise TypeError(
                    'Extension "{}.{}" must be of type: "{}.{}"'.format(
                        ext.__class__.__module__, ext.__class__.__name__,
                        Extension.__module__, Extension.__name__
                    )
                )
        return self

    def build_extension(self, ext_name, configs):
        """
        Build extension from a string name, then return an instance.

        First attempt to load an entry point. The string name must be registered as an entry point in the
        `markdown.extensions` group which points to a subclass of the `markdown.extensions.Extension` class.
        If multiple distributions have registered the same name, the first one found is returned.

        If no entry point is found, assume dot notation (`path.to.module:ClassName`). Load the specified class and
        return an instance. If no class is specified, import the module and call a `makeExtension` function and return
        the Extension instance returned by that function.
        """
        configs = dict(configs)

        entry_points = [ep for ep in util.INSTALLED_EXTENSIONS if ep.name == ext_name]
        if entry_points:
            ext = entry_points[0].load()
            return ext(**configs)

        # Get class name (if provided): `path.to.module:ClassName`
        ext_name, class_name = ext_name.split(':', 1) if ':' in ext_name else (ext_name, '')

        try:
            module = importlib.import_module(ext_name)
            logger.debug(
                'Successfully imported extension module "%s".' % ext_name
            )
        except ImportError as e:
            message = 'Failed loading extension "%s".' % ext_name
            e.args = (message,) + e.args[1:]
            raise

        if class_name:
            # Load given class name from module.
            return getattr(module, class_name)(**configs)
        else:
            # Expect  makeExtension() function to return a class.
            try:
                return module.makeExtension(**configs)
            except AttributeError as e:
                message = e.args[0]
                message = "Failed to initiate extension " \
                          "'%s': %s" % (ext_name, message)
                e.args = (message,) + e.args[1:]
                raise

    def registerExtension(self, extension):
        """ This gets called by the extension """
        self.registeredExtensions.append(extension)
        return self

    def reset(self):
        """
        Resets all state variables so that we can start with a new text.
        """
        self.htmlStash.reset()
        self.references.clear()

        for extension in self.registeredExtensions:
            if hasattr(extension, 'reset'):
                extension.reset()

        return self

    def set_output_format(self, format):
        """ Set the output format for the class instance. """
        self.output_format = format.lower().rstrip('145')  # ignore num
        try:
            self.serializer = self.output_formats[self.output_format]
        except KeyError as e:
            valid_formats = list(self.output_formats.keys())
            valid_formats.sort()
            message = 'Invalid Output Format: "%s". Use one of %s.' \
                % (self.output_format,
                   '"' + '", "'.join(valid_formats) + '"')
            e.args = (message,) + e.args[1:]
            raise
        return self

    def is_block_level(self, tag):
        """Check if the tag is a block level HTML tag."""
        if isinstance(tag, str):
            return tag.lower().rstrip('/') in self.block_level_elements
        # Some ElementTree tags are not strings, so return False.
        return False

    def convert(self, source):
        """
        Convert markdown to serialized XHTML or HTML.

        Keyword arguments:

        * source: Source text as a Unicode string.

        Markdown processing takes place in five steps:

        1. A bunch of "preprocessors" munge the input text.
        2. BlockParser() parses the high-level structural elements of the
           pre-processed text into an ElementTree.
        3. A bunch of "treeprocessors" are run against the ElementTree. One
           such treeprocessor runs InlinePatterns against the ElementTree,
           detecting inline markup.
        4. Some post-processors are run against the text after the ElementTree
           has been serialized into text.
        5. The output is written to a string.

        """

        # Fixup the source text
        if not source.strip():
            return ''  # a blank unicode string

        try:
            source = str(source)
        except UnicodeDecodeError as e:  # pragma: no cover
            # Customise error message while maintaining original trackback
            e.reason += '. -- Note: Markdown only accepts unicode input!'
            raise

        # Split into lines and run the line preprocessors.
        self.lines = source.split("\n")
        for prep in self.preprocessors:
            self.lines = prep.run(self.lines)

        # Parse the high-level elements.
        root = self.parser.parseDocument(self.lines).getroot()

        # Run the tree-processors
        for treeprocessor in self.treeprocessors:
            newRoot = treeprocessor.run(root)
            if newRoot is not None:
                root = newRoot

        # Serialize _properly_.  Strip top-level tags.
        output = self.serializer(root)
        if self.stripTopLevelTags:
            try:
                start = output.index(
                    '<%s>' % self.doc_tag) + len(self.doc_tag) + 2
                end = output.rindex('</%s>' % self.doc_tag)
                output = output[start:end].strip()
            except ValueError as e:  # pragma: no cover
                if output.strip().endswith('<%s />' % self.doc_tag):
                    # We have an empty document
                    output = ''
                else:
                    # We have a serious problem
                    raise ValueError('Markdown failed to strip top-level '
                                     'tags. Document=%r' % output.strip()) from e

        # Run the text post-processors
        for pp in self.postprocessors:
            output = pp.run(output)

        return output.strip()

    def convertFile(self, input=None, output=None, encoding=None):
        """Converts a markdown file and returns the HTML as a unicode string.

        Decodes the file using the provided encoding (defaults to utf-8),
        passes the file content to markdown, and outputs the html to either
        the provided stream or the file with provided name, using the same
        encoding as the source file. The 'xmlcharrefreplace' error handler is
        used when encoding the output.

        **Note:** This is the only place that decoding and encoding of unicode
        takes place in Python-Markdown.  (All other code is unicode-in /
        unicode-out.)

        Keyword arguments:

        * input: File object or path. Reads from stdin if `None`.
        * output: File object or path. Writes to stdout if `None`.
        * encoding: Encoding of input and output files. Defaults to utf-8.

        """

        encoding = encoding or "utf-8"

        # Read the source
        if input:
            if isinstance(input, str):
                input_file = codecs.open(input, mode="r", encoding=encoding)
            else:
                input_file = codecs.getreader(encoding)(input)
            text = input_file.read()
            input_file.close()
        else:
            text = sys.stdin.read()
            if not isinstance(text, str):  # pragma: no cover
                text = text.decode(encoding)

        text = text.lstrip('\ufeff')  # remove the byte-order mark

        # Convert
        html = self.convert(text)

        # Write to file or stdout
        if output:
            if isinstance(output, str):
                output_file = codecs.open(output, "w",
                                          encoding=encoding,
                                          errors="xmlcharrefreplace")
                output_file.write(html)
                output_file.close()
            else:
                writer = codecs.getwriter(encoding)
                output_file = writer(output, errors="xmlcharrefreplace")
                output_file.write(html)
                # Don't close here. User may want to write more.
        else:
            # Encode manually and write bytes to stdout.
            html = html.encode(encoding, "xmlcharrefreplace")
            try:
                # Write bytes directly to buffer (Python 3).
                sys.stdout.buffer.write(html)
            except AttributeError:  # pragma: no cover
                # Probably Python 2, which works with bytes by default.
                sys.stdout.write(html)

        return self


"""
EXPORTED FUNCTIONS
=============================================================================

Those are the two functions we really mean to export: markdown() and
markdownFromFile().
"""


def markdown(text, **kwargs):
    """Convert a markdown string to HTML and return HTML as a unicode string.

    This is a shortcut function for `Markdown` class to cover the most
    basic use case.  It initializes an instance of Markdown, loads the
    necessary extensions and runs the parser on the given text.

    Keyword arguments:

    * text: Markdown formatted text as Unicode or ASCII string.
    * Any arguments accepted by the Markdown class.

    Returns: An HTML document as a string.

    """
    md = Markdown(**kwargs)
    return md.convert(text)


def markdownFromFile(**kwargs):
    """Read markdown code from a file and write it to a file or a stream.

    This is a shortcut function which initializes an instance of Markdown,
    and calls the convertFile method rather than convert.

    Keyword arguments:

    * input: a file name or readable object.
    * output: a file name or writable object.
    * encoding: Encoding of input and output.
    * Any arguments accepted by the Markdown class.

    """
    md = Markdown(**kwargs)
    md.convertFile(kwargs.get('input', None),
                   kwargs.get('output', None),
                   kwargs.get('encoding', None))

social_vdgnn/lib/python3.6/site-packages/markdown/core.py
eliatrevisan/IA-MPPI-LBM
"""
CodeHilite Extension for Python-Markdown
========================================

Adds code/syntax highlighting to standard Python-Markdown code blocks.

See <https://Python-Markdown.github.io/extensions/code_hilite>
for documentation.

Original code Copyright 2006-2008 [Waylan Limberg](http://achinghead.com/).

All changes Copyright 2008-2014 The Python Markdown Project

License: [BSD](https://opensource.org/licenses/bsd-license.php)

"""

from . import Extension
from ..treeprocessors import Treeprocessor
from ..util import parseBoolValue

try:  # pragma: no cover
    from pygments import highlight
    from pygments.lexers import get_lexer_by_name, guess_lexer
    from pygments.formatters import get_formatter_by_name
    pygments = True
except ImportError:  # pragma: no cover
    pygments = False


def parse_hl_lines(expr):
    """Support our syntax for emphasizing certain lines of code.

    expr should be like '1 2' to emphasize lines 1 and 2 of a code block.
    Returns a list of ints, the line numbers to emphasize.
    """
    if not expr:
        return []

    try:
        return list(map(int, expr.split()))
    except ValueError:  # pragma: no cover
        return []


# ------------------ The Main CodeHilite Class ----------------------
class CodeHilite:
    """
    Determine language of source code, and pass it on to the Pygments highlighter.

    Usage:
        code = CodeHilite(src=some_code, lang='python')
        html = code.hilite()

    Arguments:
    * src: Source string or any object with a .readline attribute.

    * lang: String name of Pygments lexer to use for highlighting. Default: `None`.

    * guess_lang: Auto-detect which lexer to use. Ignored if `lang` is set to a valid
      value. Default: `True`.

    * use_pygments: Pass code to pygments for code highlighting. If `False`, the code is
      instead wrapped for highlighting by a JavaScript library. Default: `True`.

    * linenums: An alias to Pygments `linenos` formatter option. Default: `None`.

    * css_class: An alias to Pygments `cssclass` formatter option. Default: 'codehilite'.

    * lang_prefix: Prefix prepended to the language when `use_pygments` is `False`.
      Default: "language-".

    Other Options:
    Any other options are accepted and passed on to the lexer and formatter. Therefore,
    valid options include any options which are accepted by the `html` formatter or
    whichever lexer the code's language uses. Note that most lexers do not have any
    options. However, a few have very useful options, such as PHP's `startinline` option.
    Any invalid options are ignored without error.

    Formatter options: https://pygments.org/docs/formatters/#HtmlFormatter
    Lexer Options: https://pygments.org/docs/lexers/

    Advanced Usage:
        code = CodeHilite(
            src = some_code,
            lang = 'php',
            startinline = True,      # Lexer option. Snippet does not start with `<?php`.
            linenostart = 42,        # Formatter option. Snippet starts on line 42.
            hl_lines = [45, 49, 50], # Formatter option. Highlight lines 45, 49, and 50.
            linenos = 'inline'       # Formatter option. Avoid alignment problems.
        )
        html = code.hilite()

    """

    def __init__(self, src, **options):
        self.src = src
        self.lang = options.pop('lang', None)
        self.guess_lang = options.pop('guess_lang', True)
        self.use_pygments = options.pop('use_pygments', True)
        self.lang_prefix = options.pop('lang_prefix', 'language-')

        if 'linenos' not in options:
            options['linenos'] = options.pop('linenums', None)
        if 'cssclass' not in options:
            options['cssclass'] = options.pop('css_class', 'codehilite')
        if 'wrapcode' not in options:
            # Override pygments default
            options['wrapcode'] = True
        # Disallow use of `full` option
        options['full'] = False

        self.options = options

    def hilite(self, shebang=True):
        """
        Pass code to the [Pygments](http://pygments.pocoo.org/) highliter with
        optional line numbers. The output should then be styled with css to
        your liking. No styles are applied by default - only styling hooks
        (i.e.: <span class="k">).

        returns : A string of html.

        """

        self.src = self.src.strip('\n')

        if self.lang is None and shebang:
            self._parseHeader()

        if pygments and self.use_pygments:
            try:
                lexer = get_lexer_by_name(self.lang, **self.options)
            except ValueError:
                try:
                    if self.guess_lang:
                        lexer = guess_lexer(self.src, **self.options)
                    else:
                        lexer = get_lexer_by_name('text', **self.options)
                except ValueError:  # pragma: no cover
                    lexer = get_lexer_by_name('text', **self.options)
            formatter = get_formatter_by_name('html', **self.options)
            return highlight(self.src, lexer, formatter)
        else:
            # just escape and build markup usable by JS highlighting libs
            txt = self.src.replace('&', '&amp;')
            txt = txt.replace('<', '&lt;')
            txt = txt.replace('>', '&gt;')
            txt = txt.replace('"', '&quot;')
            classes = []
            if self.lang:
                classes.append('{}{}'.format(self.lang_prefix, self.lang))
            if self.options['linenos']:
                classes.append('linenums')
            class_str = ''
            if classes:
                class_str = ' class="{}"'.format(' '.join(classes))
            return '<pre class="{}"><code{}>{}\n</code></pre>\n'.format(
                self.options['cssclass'],
                class_str,
                txt
            )

    def _parseHeader(self):
        """
        Determines language of a code block from shebang line and whether the
        said line should be removed or left in place. If the sheband line
        contains a path (even a single /) then it is assumed to be a real
        shebang line and left alone. However, if no path is given
        (e.i.: #!python or :::python) then it is assumed to be a mock shebang
        for language identification of a code fragment and removed from the
        code block prior to processing for code highlighting. When a mock
        shebang (e.i: #!python) is found, line numbering is turned on. When
        colons are found in place of a shebang (e.i.: :::python), line
        numbering is left in the current state - off by default.

        Also parses optional list of highlight lines, like:

            :::python hl_lines="1 3"
        """

        import re

        # split text into lines
        lines = self.src.split("\n")
        # pull first line to examine
        fl = lines.pop(0)

        c = re.compile(r'''
            (?:(?:^::+)|(?P<shebang>^[#]!)) # Shebang or 2 or more colons
            (?P<path>(?:/\w+)*[/ ])?        # Zero or 1 path
            (?P<lang>[\w#.+-]*)             # The language
            \s*                             # Arbitrary whitespace
            # Optional highlight lines, single- or double-quote-delimited
            (hl_lines=(?P<quot>"|')(?P<hl_lines>.*?)(?P=quot))?
            ''',  re.VERBOSE)
        # search first line for shebang
        m = c.search(fl)
        if m:
            # we have a match
            try:
                self.lang = m.group('lang').lower()
            except IndexError:  # pragma: no cover
                self.lang = None
            if m.group('path'):
                # path exists - restore first line
                lines.insert(0, fl)
            if self.options['linenos'] is None and m.group('shebang'):
                # Overridable and Shebang exists - use line numbers
                self.options['linenos'] = True

            self.options['hl_lines'] = parse_hl_lines(m.group('hl_lines'))
        else:
            # No match
            lines.insert(0, fl)

        self.src = "\n".join(lines).strip("\n")


# ------------------ The Markdown Extension -------------------------------


class HiliteTreeprocessor(Treeprocessor):
    """ Hilight source code in code blocks. """

    def code_unescape(self, text):
        """Unescape code."""
        text = text.replace("&lt;", "<")
        text = text.replace("&gt;", ">")
        # Escaped '&' should be replaced at the end to avoid
        # conflicting with < and >.
        text = text.replace("&amp;", "&")
        return text

    def run(self, root):
        """ Find code blocks and store in htmlStash. """
        blocks = root.iter('pre')
        for block in blocks:
            if len(block) == 1 and block[0].tag == 'code':
                code = CodeHilite(
                    self.code_unescape(block[0].text),
                    tab_length=self.md.tab_length,
                    style=self.config.pop('pygments_style', 'default'),
                    **self.config
                )
                placeholder = self.md.htmlStash.store(code.hilite())
                # Clear codeblock in etree instance
                block.clear()
                # Change to p element which will later
                # be removed when inserting raw html
                block.tag = 'p'
                block.text = placeholder


class CodeHiliteExtension(Extension):
    """ Add source code hilighting to markdown codeblocks. """

    def __init__(self, **kwargs):
        # define default configs
        self.config = {
            'linenums': [None,
                         "Use lines numbers. True|table|inline=yes, False=no, None=auto"],
            'guess_lang': [True,
                           "Automatic language detection - Default: True"],
            'css_class': ["codehilite",
                          "Set class name for wrapper <div> - "
                          "Default: codehilite"],
            'pygments_style': ['default',
                               'Pygments HTML Formatter Style '
                               '(Colorscheme) - Default: default'],
            'noclasses': [False,
                          'Use inline styles instead of CSS classes - '
                          'Default false'],
            'use_pygments': [True,
                             'Use Pygments to Highlight code blocks. '
                             'Disable if using a JavaScript library. '
                             'Default: True'],
            'lang_prefix': [
                'language-',
                'Prefix prepended to the language when use_pygments is false. Default: "language-"'
            ]
            }

        for key, value in kwargs.items():
            if key in self.config:
                self.setConfig(key, value)
            else:
                # manually set unknown keywords.
                if isinstance(value, str):
                    try:
                        # Attempt to parse str as a bool value
                        value = parseBoolValue(value, preserve_none=True)
                    except ValueError:
                        pass  # Assume it's not a bool value. Use as-is.
                self.config[key] = [value, '']

    def extendMarkdown(self, md):
        """ Add HilitePostprocessor to Markdown instance. """
        hiliter = HiliteTreeprocessor(md)
        hiliter.config = self.getConfigs()
        md.treeprocessors.register(hiliter, 'hilite', 30)

        md.registerExtension(self)


def makeExtension(**kwargs):  # pragma: no cover
    return CodeHiliteExtension(**kwargs)

social_vdgnn/lib/python3.6/site-packages/markdown/extensions/codehilite.py
eliatrevisan/IA-MPPI-LBM
"""
Fenced Code Extension for Python Markdown
=========================================

This extension adds Fenced Code Blocks to Python-Markdown.

See <https://Python-Markdown.github.io/extensions/fenced_code_blocks>
for documentation.

Original code Copyright 2007-2008 [Waylan Limberg](http://achinghead.com/).


All changes Copyright 2008-2014 The Python Markdown Project

License: [BSD](https://opensource.org/licenses/bsd-license.php)
"""


from textwrap import dedent
from . import Extension
from ..preprocessors import Preprocessor
from .codehilite import CodeHilite, CodeHiliteExtension, parse_hl_lines
from .attr_list import get_attrs, AttrListExtension
from ..util import parseBoolValue
import re


class FencedCodeExtension(Extension):
    def __init__(self, **kwargs):
        self.config = {
            'lang_prefix': ['language-', 'Prefix prepended to the language. Default: "language-"']
        }
        super().__init__(**kwargs)

    def extendMarkdown(self, md):
        """ Add FencedBlockPreprocessor to the Markdown instance. """
        md.registerExtension(self)

        md.preprocessors.register(FencedBlockPreprocessor(md, self.getConfigs()), 'fenced_code_block', 25)


class FencedBlockPreprocessor(Preprocessor):
    FENCED_BLOCK_RE = re.compile(
        dedent(r'''
            (?P<fence>^(?:~{3,}|`{3,}))[ ]*                          # opening fence
            ((\{(?P<attrs>[^\}\n]*)\})|                              # (optional {attrs} or
            (\.?(?P<lang>[\w#.+-]*)[ ]*)?                            # optional (.)lang
            (hl_lines=(?P<quot>"|')(?P<hl_lines>.*?)(?P=quot)[ ]*)?) # optional hl_lines)
            \n                                                       # newline (end of opening fence)
            (?P<code>.*?)(?<=\n)                                     # the code block
            (?P=fence)[ ]*$                                          # closing fence
        '''),
        re.MULTILINE | re.DOTALL | re.VERBOSE
    )

    def __init__(self, md, config):
        super().__init__(md)
        self.config = config
        self.checked_for_deps = False
        self.codehilite_conf = {}
        self.use_attr_list = False
        # List of options to convert to bool values
        self.bool_options = [
            'linenums',
            'guess_lang',
            'noclasses',
            'use_pygments'
        ]

    def run(self, lines):
        """ Match and store Fenced Code Blocks in the HtmlStash. """

        # Check for dependent extensions
        if not self.checked_for_deps:
            for ext in self.md.registeredExtensions:
                if isinstance(ext, CodeHiliteExtension):
                    self.codehilite_conf = ext.getConfigs()
                if isinstance(ext, AttrListExtension):
                    self.use_attr_list = True

            self.checked_for_deps = True

        text = "\n".join(lines)
        while 1:
            m = self.FENCED_BLOCK_RE.search(text)
            if m:
                lang, id, classes, config = None, '', [], {}
                if m.group('attrs'):
                    id, classes, config = self.handle_attrs(get_attrs(m.group('attrs')))
                    if len(classes):
                        lang = classes.pop(0)
                else:
                    if m.group('lang'):
                        lang = m.group('lang')
                    if m.group('hl_lines'):
                        # Support hl_lines outside of attrs for backward-compatibility
                        config['hl_lines'] = parse_hl_lines(m.group('hl_lines'))

                # If config is not empty, then the codehighlite extension
                # is enabled, so we call it to highlight the code
                if self.codehilite_conf and self.codehilite_conf['use_pygments'] and config.get('use_pygments', True):
                    local_config = self.codehilite_conf.copy()
                    local_config.update(config)
                    # Combine classes with cssclass. Ensure cssclass is at end
                    # as pygments appends a suffix under certain circumstances.
                    # Ignore ID as Pygments does not offer an option to set it.
                    if classes:
                        local_config['css_class'] = '{} {}'.format(
                            ' '.join(classes),
                            local_config['css_class']
                        )
                    highliter = CodeHilite(
                        m.group('code'),
                        lang=lang,
                        style=local_config.pop('pygments_style', 'default'),
                        **local_config
                    )

                    code = highliter.hilite(shebang=False)
                else:
                    id_attr = lang_attr = class_attr = kv_pairs = ''
                    if lang:
                        lang_attr = ' class="{}{}"'.format(self.config.get('lang_prefix', 'language-'), lang)
                    if classes:
                        class_attr = ' class="{}"'.format(' '.join(classes))
                    if id:
                        id_attr = ' id="{}"'.format(id)
                    if self.use_attr_list and config and not config.get('use_pygments', False):
                        # Only assign key/value pairs to code element if attr_list ext is enabled, key/value pairs
                        # were defined on the code block, and the `use_pygments` key was not set to True. The
                        # `use_pygments` key could be either set to False or not defined. It is omitted from output.
                        kv_pairs = ' ' + ' '.join(
                            '{k}="{v}"'.format(k=k, v=v) for k, v in config.items() if k != 'use_pygments'
                        )
                    code = '<pre{id}{cls}><code{lang}{kv}>{code}</code></pre>'.format(
                        id=id_attr,
                        cls=class_attr,
                        lang=lang_attr,
                        kv=kv_pairs,
                        code=self._escape(m.group('code'))
                    )

                placeholder = self.md.htmlStash.store(code)
                text = '{}\n{}\n{}'.format(text[:m.start()],
                                           placeholder,
                                           text[m.end():])
            else:
                break
        return text.split("\n")

    def handle_attrs(self, attrs):
        """ Return tuple: (id, [list, of, classes], {configs}) """
        id = ''
        classes = []
        configs = {}
        for k, v in attrs:
            if k == 'id':
                id = v
            elif k == '.':
                classes.append(v)
            elif k == 'hl_lines':
                configs[k] = parse_hl_lines(v)
            elif k in self.bool_options:
                configs[k] = parseBoolValue(v, fail_on_errors=False, preserve_none=True)
            else:
                configs[k] = v
        return id, classes, configs

    def _escape(self, txt):
        """ basic html escaping """
        txt = txt.replace('&', '&amp;')
        txt = txt.replace('<', '&lt;')
        txt = txt.replace('>', '&gt;')
        txt = txt.replace('"', '&quot;')
        return txt


def makeExtension(**kwargs):  # pragma: no cover
    return FencedCodeExtension(**kwargs)

social_vdgnn/lib/python3.6/site-packages/markdown/extensions/fenced_code.py
eliatrevisan/IA-MPPI-LBM
"""
Python-Markdown Markdown in HTML Extension
===============================

An implementation of [PHP Markdown Extra](http://michelf.com/projects/php-markdown/extra/)'s
parsing of Markdown syntax in raw HTML.

See <https://Python-Markdown.github.io/extensions/raw_html>
for documentation.

Copyright The Python Markdown Project

License: [BSD](https://opensource.org/licenses/bsd-license.php)

"""

from . import Extension
from ..blockprocessors import BlockProcessor
from ..preprocessors import Preprocessor
from ..postprocessors import RawHtmlPostprocessor
from .. import util
from ..htmlparser import HTMLExtractor, blank_line_re
import xml.etree.ElementTree as etree


class HTMLExtractorExtra(HTMLExtractor):
    """
    Override HTMLExtractor and create etree Elements for any elements which should have content parsed as Markdown.
    """

    def __init__(self, md, *args, **kwargs):
        # All block-level tags.
        self.block_level_tags = set(md.block_level_elements.copy())
        # Block-level tags in which the content only gets span level parsing
        self.span_tags = set(
            ['address', 'dd', 'dt', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'legend', 'li', 'p', 'summary', 'td', 'th']
        )
        # Block-level tags which never get their content parsed.
        self.raw_tags = set(['canvas', 'math', 'option', 'pre', 'script', 'style', 'textarea'])

        super().__init__(md, *args, **kwargs)

        # Block-level tags in which the content gets parsed as blocks
        self.block_tags = set(self.block_level_tags) - (self.span_tags | self.raw_tags | self.empty_tags)
        self.span_and_blocks_tags = self.block_tags | self.span_tags

    def reset(self):
        """Reset this instance.  Loses all unprocessed data."""
        self.mdstack = []  # When markdown=1, stack contains a list of tags
        self.treebuilder = etree.TreeBuilder()
        self.mdstate = []  # one of 'block', 'span', 'off', or None
        super().reset()

    def close(self):
        """Handle any buffered data."""
        super().close()
        # Handle any unclosed tags.
        if self.mdstack:
            # Close the outermost parent. handle_endtag will close all unclosed children.
            self.handle_endtag(self.mdstack[0])

    def get_element(self):
        """ Return element from treebuilder and reset treebuilder for later use. """
        element = self.treebuilder.close()
        self.treebuilder = etree.TreeBuilder()
        return element

    def get_state(self, tag, attrs):
        """ Return state from tag and `markdown` attr. One of 'block', 'span', or 'off'. """
        md_attr = attrs.get('markdown', '0')
        if md_attr == 'markdown':
            # `<tag markdown>` is the same as `<tag markdown='1'>`.
            md_attr = '1'
        parent_state = self.mdstate[-1] if self.mdstate else None
        if parent_state == 'off' or (parent_state == 'span' and md_attr != '0'):
            # Only use the parent state if it is more restrictive than the markdown attribute.
            md_attr = parent_state
        if ((md_attr == '1' and tag in self.block_tags) or
                (md_attr == 'block' and tag in self.span_and_blocks_tags)):
            return 'block'
        elif ((md_attr == '1' and tag in self.span_tags) or
              (md_attr == 'span' and tag in self.span_and_blocks_tags)):
            return 'span'
        elif tag in self.block_level_tags:
            return 'off'
        else:  # pragma: no cover
            return None

    def handle_starttag(self, tag, attrs):
        # Handle tags that should always be empty and do not specify a closing tag
        if tag in self.empty_tags and (self.at_line_start() or self.intail):
            attrs = {key: value if value is not None else key for key, value in attrs}
            if "markdown" in attrs:
                attrs.pop('markdown')
                element = etree.Element(tag, attrs)
                data = etree.tostring(element, encoding='unicode', method='html')
            else:
                data = self.get_starttag_text()
            self.handle_empty_tag(data, True)
            return

        if tag in self.block_level_tags and (self.at_line_start() or self.intail):
            # Valueless attr (ex: `<tag checked>`) results in `[('checked', None)]`.
            # Convert to `{'checked': 'checked'}`.
            attrs = {key: value if value is not None else key for key, value in attrs}
            state = self.get_state(tag, attrs)
            if self.inraw or (state in [None, 'off'] and not self.mdstack):
                # fall back to default behavior
                attrs.pop('markdown', None)
                super().handle_starttag(tag, attrs)
            else:
                if 'p' in self.mdstack and tag in self.block_level_tags:
                    # Close unclosed 'p' tag
                    self.handle_endtag('p')
                self.mdstate.append(state)
                self.mdstack.append(tag)
                attrs['markdown'] = state
                self.treebuilder.start(tag, attrs)
        else:
            # Span level tag
            if self.inraw:
                super().handle_starttag(tag, attrs)
            else:
                text = self.get_starttag_text()
                if self.mdstate and self.mdstate[-1] == "off":
                    self.handle_data(self.md.htmlStash.store(text))
                else:
                    self.handle_data(text)
                if tag in self.CDATA_CONTENT_ELEMENTS:
                    # This is presumably a standalone tag in a code span (see #1036).
                    self.clear_cdata_mode()

    def handle_endtag(self, tag):
        if tag in self.block_level_tags:
            if self.inraw:
                super().handle_endtag(tag)
            elif tag in self.mdstack:
                # Close element and any unclosed children
                while self.mdstack:
                    item = self.mdstack.pop()
                    self.mdstate.pop()
                    self.treebuilder.end(item)
                    if item == tag:
                        break
                if not self.mdstack:
                    # Last item in stack is closed. Stash it
                    element = self.get_element()
                    # Get last entry to see if it ends in newlines
                    # If it is an element, assume there is no newlines
                    item = self.cleandoc[-1] if self.cleandoc else ''
                    # If we only have one newline before block element, add another
                    if not item.endswith('\n\n') and item.endswith('\n'):
                        self.cleandoc.append('\n')
                    self.cleandoc.append(self.md.htmlStash.store(element))
                    self.cleandoc.append('\n\n')
                    self.state = []
                    # Check if element has a tail
                    if not blank_line_re.match(
                            self.rawdata[self.line_offset + self.offset + len(self.get_endtag_text(tag)):]):
                        # More content exists after endtag.
                        self.intail = True
            else:
                # Treat orphan closing tag as a span level tag.
                text = self.get_endtag_text(tag)
                if self.mdstate and self.mdstate[-1] == "off":
                    self.handle_data(self.md.htmlStash.store(text))
                else:
                    self.handle_data(text)
        else:
            # Span level tag
            if self.inraw:
                super().handle_endtag(tag)
            else:
                text = self.get_endtag_text(tag)
                if self.mdstate and self.mdstate[-1] == "off":
                    self.handle_data(self.md.htmlStash.store(text))
                else:
                    self.handle_data(text)

    def handle_startendtag(self, tag, attrs):
        if tag in self.empty_tags:
            attrs = {key: value if value is not None else key for key, value in attrs}
            if "markdown" in attrs:
                attrs.pop('markdown')
                element = etree.Element(tag, attrs)
                data = etree.tostring(element, encoding='unicode', method='html')
            else:
                data = self.get_starttag_text()
        else:
            data = self.get_starttag_text()
        self.handle_empty_tag(data, is_block=self.md.is_block_level(tag))

    def handle_data(self, data):
        if self.intail and '\n' in data:
            self.intail = False
        if self.inraw or not self.mdstack:
            super().handle_data(data)
        else:
            self.treebuilder.data(data)

    def handle_empty_tag(self, data, is_block):
        if self.inraw or not self.mdstack:
            super().handle_empty_tag(data, is_block)
        else:
            if self.at_line_start() and is_block:
                self.handle_data('\n' + self.md.htmlStash.store(data) + '\n\n')
            else:
                self.handle_data(self.md.htmlStash.store(data))

    def parse_pi(self, i):
        if self.at_line_start() or self.intail or self.mdstack:
            # The same override exists in HTMLExtractor without the check
            # for mdstack. Therefore, use HTMLExtractor's parent instead.
            return super(HTMLExtractor, self).parse_pi(i)
        # This is not the beginning of a raw block so treat as plain data
        # and avoid consuming any tags which may follow (see #1066).
        self.handle_data('<?')
        return i + 2

    def parse_html_declaration(self, i):
        if self.at_line_start() or self.intail or self.mdstack:
            # The same override exists in HTMLExtractor without the check
            # for mdstack. Therefore, use HTMLExtractor's parent instead.
            return super(HTMLExtractor, self).parse_html_declaration(i)
        # This is not the beginning of a raw block so treat as plain data
        # and avoid consuming any tags which may follow (see #1066).
        self.handle_data('<!')
        return i + 2


class HtmlBlockPreprocessor(Preprocessor):
    """Remove html blocks from the text and store them for later retrieval."""

    def run(self, lines):
        source = '\n'.join(lines)
        parser = HTMLExtractorExtra(self.md)
        parser.feed(source)
        parser.close()
        return ''.join(parser.cleandoc).split('\n')


class MarkdownInHtmlProcessor(BlockProcessor):
    """Process Markdown Inside HTML Blocks which have been stored in the HtmlStash."""

    def test(self, parent, block):
        # ALways return True. `run` will return `False` it not a valid match.
        return True

    def parse_element_content(self, element):
        """
        Resursively parse the text content of an etree Element as Markdown.

        Any block level elements generated from the Markdown will be inserted as children of the element in place
        of the text content. All `markdown` attributes are removed. For any elements in which Markdown parsing has
        been dissabled, the text content of it and its chidlren are wrapped in an `AtomicString`.
        """

        md_attr = element.attrib.pop('markdown', 'off')

        if md_attr == 'block':
            # Parse content as block level
            # The order in which the different parts are parsed (text, children, tails) is important here as the
            # order of elements needs to be preserved. We can't be inserting items at a later point in the current
            # iteration as we don't want to do raw processing on elements created from parsing Markdown text (for
            # example). Therefore, the order of operations is children, tails, text.

            # Recursively parse existing children from raw HTML
            for child in list(element):
                self.parse_element_content(child)

            # Parse Markdown text in tail of children. Do this seperate to avoid raw HTML parsing.
            # Save the position of each item to be inserted later in reverse.
            tails = []
            for pos, child in enumerate(element):
                if child.tail:
                    block = child.tail.rstrip('\n')
                    child.tail = ''
                    # Use a dummy placeholder element.
                    dummy = etree.Element('div')
                    self.parser.parseBlocks(dummy, block.split('\n\n'))
                    children = list(dummy)
                    children.reverse()
                    tails.append((pos + 1, children))

            # Insert the elements created from the tails in reverse.
            tails.reverse()
            for pos, tail in tails:
                for item in tail:
                    element.insert(pos, item)

            # Parse Markdown text content. Do this last to avoid raw HTML parsing.
            if element.text:
                block = element.text.rstrip('\n')
                element.text = ''
                # Use a dummy placeholder element as the content needs to get inserted before existing children.
                dummy = etree.Element('div')
                self.parser.parseBlocks(dummy, block.split('\n\n'))
                children = list(dummy)
                children.reverse()
                for child in children:
                    element.insert(0, child)

        elif md_attr == 'span':
            # Span level parsing will be handled by inlineprocessors.
            # Walk children here to remove any `markdown` attributes.
            for child in list(element):
                self.parse_element_content(child)

        else:
            # Disable inline parsing for everything else
            if element.text is None:
                element.text = ''
            element.text = util.AtomicString(element.text)
            for child in list(element):
                self.parse_element_content(child)
                if child.tail:
                    child.tail = util.AtomicString(child.tail)

    def run(self, parent, blocks):
        m = util.HTML_PLACEHOLDER_RE.match(blocks[0])
        if m:
            index = int(m.group(1))
            element = self.parser.md.htmlStash.rawHtmlBlocks[index]
            if isinstance(element, etree.Element):
                # We have a matched element. Process it.
                blocks.pop(0)
                self.parse_element_content(element)
                parent.append(element)
                # Cleanup stash. Replace element with empty string to avoid confusing postprocessor.
                self.parser.md.htmlStash.rawHtmlBlocks.pop(index)
                self.parser.md.htmlStash.rawHtmlBlocks.insert(index, '')
                # Comfirm the match to the blockparser.
                return True
        # No match found.
        return False


class MarkdownInHTMLPostprocessor(RawHtmlPostprocessor):
    def stash_to_string(self, text):
        """ Override default to handle any etree elements still in the stash. """
        if isinstance(text, etree.Element):
            return self.md.serializer(text)
        else:
            return str(text)


class MarkdownInHtmlExtension(Extension):
    """Add Markdown parsing in HTML to Markdown class."""

    def extendMarkdown(self, md):
        """ Register extension instances. """

        # Replace raw HTML preprocessor
        md.preprocessors.register(HtmlBlockPreprocessor(md), 'html_block', 20)
        # Add blockprocessor which handles the placeholders for etree elements
        md.parser.blockprocessors.register(
            MarkdownInHtmlProcessor(md.parser), 'markdown_block', 105
        )
        # Replace raw HTML postprocessor
        md.postprocessors.register(MarkdownInHTMLPostprocessor(md), 'raw_html', 30)


def makeExtension(**kwargs):  # pragma: no cover
    return MarkdownInHtmlExtension(**kwargs)

social_vdgnn/lib/python3.6/site-packages/markdown/extensions/md_in_html.py
eliatrevisan/IA-MPPI-LBM
"""
Table of Contents Extension for Python-Markdown
===============================================

See <https://Python-Markdown.github.io/extensions/toc>
for documentation.

Oringinal code Copyright 2008 [Jack Miller](https://codezen.org/)

All changes Copyright 2008-2014 The Python Markdown Project

License: [BSD](https://opensource.org/licenses/bsd-license.php)

"""

from . import Extension
from ..treeprocessors import Treeprocessor
from ..util import code_escape, parseBoolValue, AMP_SUBSTITUTE, HTML_PLACEHOLDER_RE, AtomicString
from ..postprocessors import UnescapePostprocessor
import re
import html
import unicodedata
import xml.etree.ElementTree as etree


def slugify(value, separator, unicode=False):
    """ Slugify a string, to make it URL friendly. """
    if not unicode:
        # Replace Extended Latin characters with ASCII, i.e. žlutý → zluty
        value = unicodedata.normalize('NFKD', value)
        value = value.encode('ascii', 'ignore').decode('ascii')
    value = re.sub(r'[^\w\s-]', '', value).strip().lower()
    return re.sub(r'[{}\s]+'.format(separator), separator, value)


def slugify_unicode(value, separator):
    """ Slugify a string, to make it URL friendly while preserving Unicode characters. """
    return slugify(value, separator, unicode=True)


IDCOUNT_RE = re.compile(r'^(.*)_([0-9]+)$')


def unique(id, ids):
    """ Ensure id is unique in set of ids. Append '_1', '_2'... if not """
    while id in ids or not id:
        m = IDCOUNT_RE.match(id)
        if m:
            id = '%s_%d' % (m.group(1), int(m.group(2))+1)
        else:
            id = '%s_%d' % (id, 1)
    ids.add(id)
    return id


def get_name(el):
    """Get title name."""

    text = []
    for c in el.itertext():
        if isinstance(c, AtomicString):
            text.append(html.unescape(c))
        else:
            text.append(c)
    return ''.join(text).strip()


def stashedHTML2text(text, md, strip_entities=True):
    """ Extract raw HTML from stash, reduce to plain text and swap with placeholder. """
    def _html_sub(m):
        """ Substitute raw html with plain text. """
        try:
            raw = md.htmlStash.rawHtmlBlocks[int(m.group(1))]
        except (IndexError, TypeError):  # pragma: no cover
            return m.group(0)
        # Strip out tags and/or entities - leaving text
        res = re.sub(r'(<[^>]+>)', '', raw)
        if strip_entities:
            res = re.sub(r'(&[\#a-zA-Z0-9]+;)', '', res)
        return res

    return HTML_PLACEHOLDER_RE.sub(_html_sub, text)


def unescape(text):
    """ Unescape escaped text. """
    c = UnescapePostprocessor()
    return c.run(text)


def nest_toc_tokens(toc_list):
    """Given an unsorted list with errors and skips, return a nested one.
    [{'level': 1}, {'level': 2}]
    =>
    [{'level': 1, 'children': [{'level': 2, 'children': []}]}]

    A wrong list is also converted:
    [{'level': 2}, {'level': 1}]
    =>
    [{'level': 2, 'children': []}, {'level': 1, 'children': []}]
    """

    ordered_list = []
    if len(toc_list):
        # Initialize everything by processing the first entry
        last = toc_list.pop(0)
        last['children'] = []
        levels = [last['level']]
        ordered_list.append(last)
        parents = []

        # Walk the rest nesting the entries properly
        while toc_list:
            t = toc_list.pop(0)
            current_level = t['level']
            t['children'] = []

            # Reduce depth if current level < last item's level
            if current_level < levels[-1]:
                # Pop last level since we know we are less than it
                levels.pop()

                # Pop parents and levels we are less than or equal to
                to_pop = 0
                for p in reversed(parents):
                    if current_level <= p['level']:
                        to_pop += 1
                    else:  # pragma: no cover
                        break
                if to_pop:
                    levels = levels[:-to_pop]
                    parents = parents[:-to_pop]

                # Note current level as last
                levels.append(current_level)

            # Level is the same, so append to
            # the current parent (if available)
            if current_level == levels[-1]:
                (parents[-1]['children'] if parents
                 else ordered_list).append(t)

            # Current level is > last item's level,
            # So make last item a parent and append current as child
            else:
                last['children'].append(t)
                parents.append(last)
                levels.append(current_level)
            last = t

    return ordered_list


class TocTreeprocessor(Treeprocessor):
    def __init__(self, md, config):
        super().__init__(md)

        self.marker = config["marker"]
        self.title = config["title"]
        self.base_level = int(config["baselevel"]) - 1
        self.slugify = config["slugify"]
        self.sep = config["separator"]
        self.use_anchors = parseBoolValue(config["anchorlink"])
        self.anchorlink_class = config["anchorlink_class"]
        self.use_permalinks = parseBoolValue(config["permalink"], False)
        if self.use_permalinks is None:
            self.use_permalinks = config["permalink"]
        self.permalink_class = config["permalink_class"]
        self.permalink_title = config["permalink_title"]
        self.header_rgx = re.compile("[Hh][123456]")
        if isinstance(config["toc_depth"], str) and '-' in config["toc_depth"]:
            self.toc_top, self.toc_bottom = [int(x) for x in config["toc_depth"].split('-')]
        else:
            self.toc_top = 1
            self.toc_bottom = int(config["toc_depth"])

    def iterparent(self, node):
        ''' Iterator wrapper to get allowed parent and child all at once. '''

        # We do not allow the marker inside a header as that
        # would causes an enless loop of placing a new TOC
        # inside previously generated TOC.
        for child in node:
            if not self.header_rgx.match(child.tag) and child.tag not in ['pre', 'code']:
                yield node, child
                yield from self.iterparent(child)

    def replace_marker(self, root, elem):
        ''' Replace marker with elem. '''
        for (p, c) in self.iterparent(root):
            text = ''.join(c.itertext()).strip()
            if not text:
                continue

            # To keep the output from screwing up the
            # validation by putting a <div> inside of a <p>
            # we actually replace the <p> in its entirety.

            # The <p> element may contain more than a single text content
            # (nl2br can introduce a <br>). In this situation, c.text returns
            # the very first content, ignore children contents or tail content.
            # len(c) == 0 is here to ensure there is only text in the <p>.
            if c.text and c.text.strip() == self.marker and len(c) == 0:
                for i in range(len(p)):
                    if p[i] == c:
                        p[i] = elem
                        break

    def set_level(self, elem):
        ''' Adjust header level according to base level. '''
        level = int(elem.tag[-1]) + self.base_level
        if level > 6:
            level = 6
        elem.tag = 'h%d' % level

    def add_anchor(self, c, elem_id):  # @ReservedAssignment
        anchor = etree.Element("a")
        anchor.text = c.text
        anchor.attrib["href"] = "#" + elem_id
        anchor.attrib["class"] = self.anchorlink_class
        c.text = ""
        for elem in c:
            anchor.append(elem)
        while len(c):
            c.remove(c[0])
        c.append(anchor)

    def add_permalink(self, c, elem_id):
        permalink = etree.Element("a")
        permalink.text = ("%spara;" % AMP_SUBSTITUTE
                          if self.use_permalinks is True
                          else self.use_permalinks)
        permalink.attrib["href"] = "#" + elem_id
        permalink.attrib["class"] = self.permalink_class
        if self.permalink_title:
            permalink.attrib["title"] = self.permalink_title
        c.append(permalink)

    def build_toc_div(self, toc_list):
        """ Return a string div given a toc list. """
        div = etree.Element("div")
        div.attrib["class"] = "toc"

        # Add title to the div
        if self.title:
            header = etree.SubElement(div, "span")
            header.attrib["class"] = "toctitle"
            header.text = self.title

        def build_etree_ul(toc_list, parent):
            ul = etree.SubElement(parent, "ul")
            for item in toc_list:
                # List item link, to be inserted into the toc div
                li = etree.SubElement(ul, "li")
                link = etree.SubElement(li, "a")
                link.text = item.get('name', '')
                link.attrib["href"] = '#' + item.get('id', '')
                if item['children']:
                    build_etree_ul(item['children'], li)
            return ul

        build_etree_ul(toc_list, div)

        if 'prettify' in self.md.treeprocessors:
            self.md.treeprocessors['prettify'].run(div)

        return div

    def run(self, doc):
        # Get a list of id attributes
        used_ids = set()
        for el in doc.iter():
            if "id" in el.attrib:
                used_ids.add(el.attrib["id"])

        toc_tokens = []
        for el in doc.iter():
            if isinstance(el.tag, str) and self.header_rgx.match(el.tag):
                self.set_level(el)
                text = get_name(el)

                # Do not override pre-existing ids
                if "id" not in el.attrib:
                    innertext = unescape(stashedHTML2text(text, self.md))
                    el.attrib["id"] = unique(self.slugify(innertext, self.sep), used_ids)

                if int(el.tag[-1]) >= self.toc_top and int(el.tag[-1]) <= self.toc_bottom:
                    toc_tokens.append({
                        'level': int(el.tag[-1]),
                        'id': el.attrib["id"],
                        'name': unescape(stashedHTML2text(
                            code_escape(el.attrib.get('data-toc-label', text)),
                            self.md, strip_entities=False
                        ))
                    })

                # Remove the data-toc-label attribute as it is no longer needed
                if 'data-toc-label' in el.attrib:
                    del el.attrib['data-toc-label']

                if self.use_anchors:
                    self.add_anchor(el, el.attrib["id"])
                if self.use_permalinks not in [False, None]:
                    self.add_permalink(el, el.attrib["id"])

        toc_tokens = nest_toc_tokens(toc_tokens)
        div = self.build_toc_div(toc_tokens)
        if self.marker:
            self.replace_marker(doc, div)

        # serialize and attach to markdown instance.
        toc = self.md.serializer(div)
        for pp in self.md.postprocessors:
            toc = pp.run(toc)
        self.md.toc_tokens = toc_tokens
        self.md.toc = toc


class TocExtension(Extension):

    TreeProcessorClass = TocTreeprocessor

    def __init__(self, **kwargs):
        self.config = {
            "marker": ['[TOC]',
                       'Text to find and replace with Table of Contents - '
                       'Set to an empty string to disable. Defaults to "[TOC]"'],
            "title": ["",
                      "Title to insert into TOC <div> - "
                      "Defaults to an empty string"],
            "anchorlink": [False,
                           "True if header should be a self link - "
                           "Defaults to False"],
            "anchorlink_class": ['toclink',
                                 'CSS class(es) used for the link. '
                                 'Defaults to "toclink"'],
            "permalink": [0,
                          "True or link text if a Sphinx-style permalink should "
                          "be added - Defaults to False"],
            "permalink_class": ['headerlink',
                                'CSS class(es) used for the link. '
                                'Defaults to "headerlink"'],
            "permalink_title": ["Permanent link",
                                "Title attribute of the permalink - "
                                "Defaults to 'Permanent link'"],
            "baselevel": ['1', 'Base level for headers.'],
            "slugify": [slugify,
                        "Function to generate anchors based on header text - "
                        "Defaults to the headerid ext's slugify function."],
            'separator': ['-', 'Word separator. Defaults to "-".'],
            "toc_depth": [6,
                          'Define the range of section levels to include in'
                          'the Table of Contents. A single integer (b) defines'
                          'the bottom section level (<h1>..<hb>) only.'
                          'A string consisting of two digits separated by a hyphen'
                          'in between ("2-5"), define the top (t) and the'
                          'bottom (b) (<ht>..<hb>). Defaults to `6` (bottom).'],
        }

        super().__init__(**kwargs)

    def extendMarkdown(self, md):
        md.registerExtension(self)
        self.md = md
        self.reset()
        tocext = self.TreeProcessorClass(md, self.getConfigs())
        # Headerid ext is set to '>prettify'. With this set to '_end',
        # it should always come after headerid ext (and honor ids assinged
        # by the header id extension) if both are used. Same goes for
        # attr_list extension. This must come last because we don't want
        # to redefine ids after toc is created. But we do want toc prettified.
        md.treeprocessors.register(tocext, 'toc', 5)

    def reset(self):
        self.md.toc = ''
        self.md.toc_tokens = []


def makeExtension(**kwargs):  # pragma: no cover
    return TocExtension(**kwargs)

social_vdgnn/lib/python3.6/site-packages/markdown/extensions/toc.py
eliatrevisan/IA-MPPI-LBM
"""
Python Markdown

A Python implementation of John Gruber's Markdown.

Documentation: https://python-markdown.github.io/
GitHub: https://github.com/Python-Markdown/markdown/
PyPI: https://pypi.org/project/Markdown/

Started by Manfred Stienstra (http://www.dwerg.net/).
Maintained for a few years by Yuri Takhteyev (http://www.freewisdom.org).
Currently maintained by Waylan Limberg (https://github.com/waylan),
Dmitry Shachnev (https://github.com/mitya57) and Isaac Muse (https://github.com/facelessuser).

Copyright 2007-2018 The Python Markdown Project (v. 1.7 and later)
Copyright 2004, 2005, 2006 Yuri Takhteyev (v. 0.2-1.6b)
Copyright 2004 Manfred Stienstra (the original version)

License: BSD (see LICENSE.md for details).

POST-PROCESSORS
=============================================================================

Markdown also allows post-processors, which are similar to preprocessors in
that they need to implement a "run" method. However, they are run after core
processing.

"""

from collections import OrderedDict
from . import util
import re


def build_postprocessors(md, **kwargs):
    """ Build the default postprocessors for Markdown. """
    postprocessors = util.Registry()
    postprocessors.register(RawHtmlPostprocessor(md), 'raw_html', 30)
    postprocessors.register(AndSubstitutePostprocessor(), 'amp_substitute', 20)
    postprocessors.register(UnescapePostprocessor(), 'unescape', 10)
    return postprocessors


class Postprocessor(util.Processor):
    """
    Postprocessors are run after the ElementTree it converted back into text.

    Each Postprocessor implements a "run" method that takes a pointer to a
    text string, modifies it as necessary and returns a text string.

    Postprocessors must extend markdown.Postprocessor.

    """

    def run(self, text):
        """
        Subclasses of Postprocessor should implement a `run` method, which
        takes the html document as a single text string and returns a
        (possibly modified) string.

        """
        pass  # pragma: no cover


class RawHtmlPostprocessor(Postprocessor):
    """ Restore raw html to the document. """

    BLOCK_LEVEL_REGEX = re.compile(r'^\<\/?([^ >]+)')

    def run(self, text):
        """ Iterate over html stash and restore html. """
        replacements = OrderedDict()
        for i in range(self.md.htmlStash.html_counter):
            html = self.stash_to_string(self.md.htmlStash.rawHtmlBlocks[i])
            if self.isblocklevel(html):
                replacements["<p>{}</p>".format(
                    self.md.htmlStash.get_placeholder(i))] = html
            replacements[self.md.htmlStash.get_placeholder(i)] = html

        def substitute_match(m):
            key = m.group(0)

            if key not in replacements:
                if key[3:-4] in replacements:
                    return f'<p>{ replacements[key[3:-4]] }</p>'
                else:
                    return key

            return replacements[key]

        if replacements:
            base_placeholder = util.HTML_PLACEHOLDER % r'([0-9]+)'
            pattern = re.compile(f'<p>{ base_placeholder }</p>|{ base_placeholder }')
            processed_text = pattern.sub(substitute_match, text)
        else:
            return text

        if processed_text == text:
            return processed_text
        else:
            return self.run(processed_text)

    def isblocklevel(self, html):
        m = self.BLOCK_LEVEL_REGEX.match(html)
        if m:
            if m.group(1)[0] in ('!', '?', '@', '%'):
                # Comment, php etc...
                return True
            return self.md.is_block_level(m.group(1))
        return False

    def stash_to_string(self, text):
        """ Convert a stashed object to a string. """
        return str(text)


class AndSubstitutePostprocessor(Postprocessor):
    """ Restore valid entities """

    def run(self, text):
        text = text.replace(util.AMP_SUBSTITUTE, "&")
        return text


class UnescapePostprocessor(Postprocessor):
    """ Restore escaped chars """

    RE = re.compile(r'{}(\d+){}'.format(util.STX, util.ETX))

    def unescape(self, m):
        return chr(int(m.group(1)))

    def run(self, text):
        return self.RE.sub(self.unescape, text)

social_vdgnn/lib/python3.6/site-packages/markdown/postprocessors.py
eliatrevisan/IA-MPPI-LBM
"""
Python Markdown

A Python implementation of John Gruber's Markdown.

Documentation: https://python-markdown.github.io/
GitHub: https://github.com/Python-Markdown/markdown/
PyPI: https://pypi.org/project/Markdown/

Started by Manfred Stienstra (http://www.dwerg.net/).
Maintained for a few years by Yuri Takhteyev (http://www.freewisdom.org).
Currently maintained by Waylan Limberg (https://github.com/waylan),
Dmitry Shachnev (https://github.com/mitya57) and Isaac Muse (https://github.com/facelessuser).

Copyright 2007-2018 The Python Markdown Project (v. 1.7 and later)
Copyright 2004, 2005, 2006 Yuri Takhteyev (v. 0.2-1.6b)
Copyright 2004 Manfred Stienstra (the original version)

License: BSD (see LICENSE.md for details).
"""

import re
import sys
import warnings
import xml.etree.ElementTree
from collections import namedtuple
from functools import wraps
from itertools import count

from .pep562 import Pep562

if sys.version_info >= (3, 10):
    from importlib import metadata
else:
    # <PY310 use backport
    import importlib_metadata as metadata

PY37 = (3, 7) <= sys.version_info


# TODO: Remove deprecated variables in a future release.
__deprecated__ = {
    'etree': ('xml.etree.ElementTree', xml.etree.ElementTree),
    'string_type': ('str', str),
    'text_type': ('str', str),
    'int2str': ('chr', chr),
    'iterrange': ('range', range)
}


"""
Constants you might want to modify
-----------------------------------------------------------------------------
"""


BLOCK_LEVEL_ELEMENTS = [
    # Elements which are invalid to wrap in a `<p>` tag.
    # See https://w3c.github.io/html/grouping-content.html#the-p-element
    'address', 'article', 'aside', 'blockquote', 'details', 'div', 'dl',
    'fieldset', 'figcaption', 'figure', 'footer', 'form', 'h1', 'h2', 'h3',
    'h4', 'h5', 'h6', 'header', 'hgroup', 'hr', 'main', 'menu', 'nav', 'ol',
    'p', 'pre', 'section', 'table', 'ul',
    # Other elements which Markdown should not be mucking up the contents of.
    'canvas', 'colgroup', 'dd', 'body', 'dt', 'group', 'iframe', 'li', 'legend',
    'math', 'map', 'noscript', 'output', 'object', 'option', 'progress', 'script',
    'style', 'tbody', 'td', 'textarea', 'tfoot', 'th', 'thead', 'tr', 'video'
]

# Placeholders
STX = '\u0002'  # Use STX ("Start of text") for start-of-placeholder
ETX = '\u0003'  # Use ETX ("End of text") for end-of-placeholder
INLINE_PLACEHOLDER_PREFIX = STX+"klzzwxh:"
INLINE_PLACEHOLDER = INLINE_PLACEHOLDER_PREFIX + "%s" + ETX
INLINE_PLACEHOLDER_RE = re.compile(INLINE_PLACEHOLDER % r'([0-9]+)')
AMP_SUBSTITUTE = STX+"amp"+ETX
HTML_PLACEHOLDER = STX + "wzxhzdk:%s" + ETX
HTML_PLACEHOLDER_RE = re.compile(HTML_PLACEHOLDER % r'([0-9]+)')
TAG_PLACEHOLDER = STX + "hzzhzkh:%s" + ETX


"""
Constants you probably do not need to change
-----------------------------------------------------------------------------
"""

# Only load extension entry_points once.
INSTALLED_EXTENSIONS = metadata.entry_points(group='markdown.extensions')
RTL_BIDI_RANGES = (
    ('\u0590', '\u07FF'),
    # Hebrew (0590-05FF), Arabic (0600-06FF),
    # Syriac (0700-074F), Arabic supplement (0750-077F),
    # Thaana (0780-07BF), Nko (07C0-07FF).
    ('\u2D30', '\u2D7F')  # Tifinagh
)


"""
AUXILIARY GLOBAL FUNCTIONS
=============================================================================
"""


def deprecated(message, stacklevel=2):
    """
    Raise a DeprecationWarning when wrapped function/method is called.

    Usage:
        @deprecated("This method will be removed in version X; use Y instead.")
        def some_method()"
            pass
    """
    def wrapper(func):
        @wraps(func)
        def deprecated_func(*args, **kwargs):
            warnings.warn(
                f"'{func.__name__}' is deprecated. {message}",
                category=DeprecationWarning,
                stacklevel=stacklevel
            )
            return func(*args, **kwargs)
        return deprecated_func
    return wrapper


@deprecated("Use 'Markdown.is_block_level' instead.")
def isBlockLevel(tag):
    """Check if the tag is a block level HTML tag."""
    if isinstance(tag, str):
        return tag.lower().rstrip('/') in BLOCK_LEVEL_ELEMENTS
    # Some ElementTree tags are not strings, so return False.
    return False


def parseBoolValue(value, fail_on_errors=True, preserve_none=False):
    """Parses a string representing bool value. If parsing was successful,
       returns True or False. If preserve_none=True, returns True, False,
       or None. If parsing was not successful, raises  ValueError, or, if
       fail_on_errors=False, returns None."""
    if not isinstance(value, str):
        if preserve_none and value is None:
            return value
        return bool(value)
    elif preserve_none and value.lower() == 'none':
        return None
    elif value.lower() in ('true', 'yes', 'y', 'on', '1'):
        return True
    elif value.lower() in ('false', 'no', 'n', 'off', '0', 'none'):
        return False
    elif fail_on_errors:
        raise ValueError('Cannot parse bool value: %r' % value)


def code_escape(text):
    """Escape code."""
    if "&" in text:
        text = text.replace("&", "&amp;")
    if "<" in text:
        text = text.replace("<", "&lt;")
    if ">" in text:
        text = text.replace(">", "&gt;")
    return text


def _get_stack_depth(size=2):
    """Get current stack depth, performantly.
    """
    frame = sys._getframe(size)

    for size in count(size):
        frame = frame.f_back
        if not frame:
            return size


def nearing_recursion_limit():
    """Return true if current stack depth is withing 100 of maximum limit."""
    return sys.getrecursionlimit() - _get_stack_depth() < 100


"""
MISC AUXILIARY CLASSES
=============================================================================
"""


class AtomicString(str):
    """A string which should not be further processed."""
    pass


class Processor:
    def __init__(self, md=None):
        self.md = md

    @property
    @deprecated("Use 'md' instead.")
    def markdown(self):
        # TODO: remove this later
        return self.md


class HtmlStash:
    """
    This class is used for stashing HTML objects that we extract
    in the beginning and replace with place-holders.
    """

    def __init__(self):
        """ Create a HtmlStash. """
        self.html_counter = 0  # for counting inline html segments
        self.rawHtmlBlocks = []
        self.tag_counter = 0
        self.tag_data = []  # list of dictionaries in the order tags appear

    def store(self, html):
        """
        Saves an HTML segment for later reinsertion.  Returns a
        placeholder string that needs to be inserted into the
        document.

        Keyword arguments:

        * html: an html segment

        Returns : a placeholder string

        """
        self.rawHtmlBlocks.append(html)
        placeholder = self.get_placeholder(self.html_counter)
        self.html_counter += 1
        return placeholder

    def reset(self):
        self.html_counter = 0
        self.rawHtmlBlocks = []

    def get_placeholder(self, key):
        return HTML_PLACEHOLDER % key

    def store_tag(self, tag, attrs, left_index, right_index):
        """Store tag data and return a placeholder."""
        self.tag_data.append({'tag': tag, 'attrs': attrs,
                              'left_index': left_index,
                              'right_index': right_index})
        placeholder = TAG_PLACEHOLDER % str(self.tag_counter)
        self.tag_counter += 1  # equal to the tag's index in self.tag_data
        return placeholder


# Used internally by `Registry` for each item in its sorted list.
# Provides an easier to read API when editing the code later.
# For example, `item.name` is more clear than `item[0]`.
_PriorityItem = namedtuple('PriorityItem', ['name', 'priority'])


class Registry:
    """
    A priority sorted registry.

    A `Registry` instance provides two public methods to alter the data of the
    registry: `register` and `deregister`. Use `register` to add items and
    `deregister` to remove items. See each method for specifics.

    When registering an item, a "name" and a "priority" must be provided. All
    items are automatically sorted by "priority" from highest to lowest. The
    "name" is used to remove ("deregister") and get items.

    A `Registry` instance it like a list (which maintains order) when reading
    data. You may iterate over the items, get an item and get a count (length)
    of all items. You may also check that the registry contains an item.

    When getting an item you may use either the index of the item or the
    string-based "name". For example:

        registry = Registry()
        registry.register(SomeItem(), 'itemname', 20)
        # Get the item by index
        item = registry[0]
        # Get the item by name
        item = registry['itemname']

    When checking that the registry contains an item, you may use either the
    string-based "name", or a reference to the actual item. For example:

        someitem = SomeItem()
        registry.register(someitem, 'itemname', 20)
        # Contains the name
        assert 'itemname' in registry
        # Contains the item instance
        assert someitem in registry

    The method `get_index_for_name` is also available to obtain the index of
    an item using that item's assigned "name".
    """

    def __init__(self):
        self._data = {}
        self._priority = []
        self._is_sorted = False

    def __contains__(self, item):
        if isinstance(item, str):
            # Check if an item exists by this name.
            return item in self._data.keys()
        # Check if this instance exists.
        return item in self._data.values()

    def __iter__(self):
        self._sort()
        return iter([self._data[k] for k, p in self._priority])

    def __getitem__(self, key):
        self._sort()
        if isinstance(key, slice):
            data = Registry()
            for k, p in self._priority[key]:
                data.register(self._data[k], k, p)
            return data
        if isinstance(key, int):
            return self._data[self._priority[key].name]
        return self._data[key]

    def __len__(self):
        return len(self._priority)

    def __repr__(self):
        return '<{}({})>'.format(self.__class__.__name__, list(self))

    def get_index_for_name(self, name):
        """
        Return the index of the given name.
        """
        if name in self:
            self._sort()
            return self._priority.index(
                [x for x in self._priority if x.name == name][0]
            )
        raise ValueError('No item named "{}" exists.'.format(name))

    def register(self, item, name, priority):
        """
        Add an item to the registry with the given name and priority.

        Parameters:

        * `item`: The item being registered.
        * `name`: A string used to reference the item.
        * `priority`: An integer or float used to sort against all items.

        If an item is registered with a "name" which already exists, the
        existing item is replaced with the new item. Tread carefully as the
        old item is lost with no way to recover it. The new item will be
        sorted according to its priority and will **not** retain the position
        of the old item.
        """
        if name in self:
            # Remove existing item of same name first
            self.deregister(name)
        self._is_sorted = False
        self._data[name] = item
        self._priority.append(_PriorityItem(name, priority))

    def deregister(self, name, strict=True):
        """
        Remove an item from the registry.

        Set `strict=False` to fail silently.
        """
        try:
            index = self.get_index_for_name(name)
            del self._priority[index]
            del self._data[name]
        except ValueError:
            if strict:
                raise

    def _sort(self):
        """
        Sort the registry by priority from highest to lowest.

        This method is called internally and should never be explicitly called.
        """
        if not self._is_sorted:
            self._priority.sort(key=lambda item: item.priority, reverse=True)
            self._is_sorted = True

    # Deprecated Methods which provide a smooth transition from OrderedDict

    def __setitem__(self, key, value):
        """ Register item with priorty 5 less than lowest existing priority. """
        if isinstance(key, str):
            warnings.warn(
                'Using setitem to register a processor or pattern is deprecated. '
                'Use the `register` method instead.',
                DeprecationWarning,
                stacklevel=2,
            )
            if key in self:
                # Key already exists, replace without altering priority
                self._data[key] = value
                return
            if len(self) == 0:
                # This is the first item. Set priority to 50.
                priority = 50
            else:
                self._sort()
                priority = self._priority[-1].priority - 5
            self.register(value, key, priority)
        else:
            raise TypeError

    def __delitem__(self, key):
        """ Deregister an item by name. """
        if key in self:
            self.deregister(key)
            warnings.warn(
                'Using del to remove a processor or pattern is deprecated. '
                'Use the `deregister` method instead.',
                DeprecationWarning,
                stacklevel=2,
            )
        else:
            raise KeyError('Cannot delete key {}, not registered.'.format(key))

    def add(self, key, value, location):
        """ Register a key by location. """
        if len(self) == 0:
            # This is the first item. Set priority to 50.
            priority = 50
        elif location == '_begin':
            self._sort()
            # Set priority 5 greater than highest existing priority
            priority = self._priority[0].priority + 5
        elif location == '_end':
            self._sort()
            # Set priority 5 less than lowest existing priority
            priority = self._priority[-1].priority - 5
        elif location.startswith('<') or location.startswith('>'):
            # Set priority halfway between existing priorities.
            i = self.get_index_for_name(location[1:])
            if location.startswith('<'):
                after = self._priority[i].priority
                if i > 0:
                    before = self._priority[i-1].priority
                else:
                    # Location is first item`
                    before = after + 10
            else:
                # location.startswith('>')
                before = self._priority[i].priority
                if i < len(self) - 1:
                    after = self._priority[i+1].priority
                else:
                    # location is last item
                    after = before - 10
            priority = before - ((before - after) / 2)
        else:
            raise ValueError('Not a valid location: "%s". Location key '
                             'must start with a ">" or "<".' % location)
        self.register(value, key, priority)
        warnings.warn(
            'Using the add method to register a processor or pattern is deprecated. '
            'Use the `register` method instead.',
            DeprecationWarning,
            stacklevel=2,
        )


def __getattr__(name):
    """Get attribute."""

    deprecated = __deprecated__.get(name)
    if deprecated:
        warnings.warn(
            "'{}' is deprecated. Use '{}' instead.".format(name, deprecated[0]),
            category=DeprecationWarning,
            stacklevel=(3 if PY37 else 4)
        )
        return deprecated[1]
    raise AttributeError("module '{}' has no attribute '{}'".format(__name__, name))


if not PY37:
    Pep562(__name__)

social_vdgnn/lib/python3.6/site-packages/markdown/util.py
eliatrevisan/IA-MPPI-LBM
from .colors import ListedColormap

_magma_data = [[0.001462, 0.000466, 0.013866],
               [0.002258, 0.001295, 0.018331],
               [0.003279, 0.002305, 0.023708],
               [0.004512, 0.003490, 0.029965],
               [0.005950, 0.004843, 0.037130],
               [0.007588, 0.006356, 0.044973],
               [0.009426, 0.008022, 0.052844],
               [0.011465, 0.009828, 0.060750],
               [0.013708, 0.011771, 0.068667],
               [0.016156, 0.013840, 0.076603],
               [0.018815, 0.016026, 0.084584],
               [0.021692, 0.018320, 0.092610],
               [0.024792, 0.020715, 0.100676],
               [0.028123, 0.023201, 0.108787],
               [0.031696, 0.025765, 0.116965],
               [0.035520, 0.028397, 0.125209],
               [0.039608, 0.031090, 0.133515],
               [0.043830, 0.033830, 0.141886],
               [0.048062, 0.036607, 0.150327],
               [0.052320, 0.039407, 0.158841],
               [0.056615, 0.042160, 0.167446],
               [0.060949, 0.044794, 0.176129],
               [0.065330, 0.047318, 0.184892],
               [0.069764, 0.049726, 0.193735],
               [0.074257, 0.052017, 0.202660],
               [0.078815, 0.054184, 0.211667],
               [0.083446, 0.056225, 0.220755],
               [0.088155, 0.058133, 0.229922],
               [0.092949, 0.059904, 0.239164],
               [0.097833, 0.061531, 0.248477],
               [0.102815, 0.063010, 0.257854],
               [0.107899, 0.064335, 0.267289],
               [0.113094, 0.065492, 0.276784],
               [0.118405, 0.066479, 0.286321],
               [0.123833, 0.067295, 0.295879],
               [0.129380, 0.067935, 0.305443],
               [0.135053, 0.068391, 0.315000],
               [0.140858, 0.068654, 0.324538],
               [0.146785, 0.068738, 0.334011],
               [0.152839, 0.068637, 0.343404],
               [0.159018, 0.068354, 0.352688],
               [0.165308, 0.067911, 0.361816],
               [0.171713, 0.067305, 0.370771],
               [0.178212, 0.066576, 0.379497],
               [0.184801, 0.065732, 0.387973],
               [0.191460, 0.064818, 0.396152],
               [0.198177, 0.063862, 0.404009],
               [0.204935, 0.062907, 0.411514],
               [0.211718, 0.061992, 0.418647],
               [0.218512, 0.061158, 0.425392],
               [0.225302, 0.060445, 0.431742],
               [0.232077, 0.059889, 0.437695],
               [0.238826, 0.059517, 0.443256],
               [0.245543, 0.059352, 0.448436],
               [0.252220, 0.059415, 0.453248],
               [0.258857, 0.059706, 0.457710],
               [0.265447, 0.060237, 0.461840],
               [0.271994, 0.060994, 0.465660],
               [0.278493, 0.061978, 0.469190],
               [0.284951, 0.063168, 0.472451],
               [0.291366, 0.064553, 0.475462],
               [0.297740, 0.066117, 0.478243],
               [0.304081, 0.067835, 0.480812],
               [0.310382, 0.069702, 0.483186],
               [0.316654, 0.071690, 0.485380],
               [0.322899, 0.073782, 0.487408],
               [0.329114, 0.075972, 0.489287],
               [0.335308, 0.078236, 0.491024],
               [0.341482, 0.080564, 0.492631],
               [0.347636, 0.082946, 0.494121],
               [0.353773, 0.085373, 0.495501],
               [0.359898, 0.087831, 0.496778],
               [0.366012, 0.090314, 0.497960],
               [0.372116, 0.092816, 0.499053],
               [0.378211, 0.095332, 0.500067],
               [0.384299, 0.097855, 0.501002],
               [0.390384, 0.100379, 0.501864],
               [0.396467, 0.102902, 0.502658],
               [0.402548, 0.105420, 0.503386],
               [0.408629, 0.107930, 0.504052],
               [0.414709, 0.110431, 0.504662],
               [0.420791, 0.112920, 0.505215],
               [0.426877, 0.115395, 0.505714],
               [0.432967, 0.117855, 0.506160],
               [0.439062, 0.120298, 0.506555],
               [0.445163, 0.122724, 0.506901],
               [0.451271, 0.125132, 0.507198],
               [0.457386, 0.127522, 0.507448],
               [0.463508, 0.129893, 0.507652],
               [0.469640, 0.132245, 0.507809],
               [0.475780, 0.134577, 0.507921],
               [0.481929, 0.136891, 0.507989],
               [0.488088, 0.139186, 0.508011],
               [0.494258, 0.141462, 0.507988],
               [0.500438, 0.143719, 0.507920],
               [0.506629, 0.145958, 0.507806],
               [0.512831, 0.148179, 0.507648],
               [0.519045, 0.150383, 0.507443],
               [0.525270, 0.152569, 0.507192],
               [0.531507, 0.154739, 0.506895],
               [0.537755, 0.156894, 0.506551],
               [0.544015, 0.159033, 0.506159],
               [0.550287, 0.161158, 0.505719],
               [0.556571, 0.163269, 0.505230],
               [0.562866, 0.165368, 0.504692],
               [0.569172, 0.167454, 0.504105],
               [0.575490, 0.169530, 0.503466],
               [0.581819, 0.171596, 0.502777],
               [0.588158, 0.173652, 0.502035],
               [0.594508, 0.175701, 0.501241],
               [0.600868, 0.177743, 0.500394],
               [0.607238, 0.179779, 0.499492],
               [0.613617, 0.181811, 0.498536],
               [0.620005, 0.183840, 0.497524],
               [0.626401, 0.185867, 0.496456],
               [0.632805, 0.187893, 0.495332],
               [0.639216, 0.189921, 0.494150],
               [0.645633, 0.191952, 0.492910],
               [0.652056, 0.193986, 0.491611],
               [0.658483, 0.196027, 0.490253],
               [0.664915, 0.198075, 0.488836],
               [0.671349, 0.200133, 0.487358],
               [0.677786, 0.202203, 0.485819],
               [0.684224, 0.204286, 0.484219],
               [0.690661, 0.206384, 0.482558],
               [0.697098, 0.208501, 0.480835],
               [0.703532, 0.210638, 0.479049],
               [0.709962, 0.212797, 0.477201],
               [0.716387, 0.214982, 0.475290],
               [0.722805, 0.217194, 0.473316],
               [0.729216, 0.219437, 0.471279],
               [0.735616, 0.221713, 0.469180],
               [0.742004, 0.224025, 0.467018],
               [0.748378, 0.226377, 0.464794],
               [0.754737, 0.228772, 0.462509],
               [0.761077, 0.231214, 0.460162],
               [0.767398, 0.233705, 0.457755],
               [0.773695, 0.236249, 0.455289],
               [0.779968, 0.238851, 0.452765],
               [0.786212, 0.241514, 0.450184],
               [0.792427, 0.244242, 0.447543],
               [0.798608, 0.247040, 0.444848],
               [0.804752, 0.249911, 0.442102],
               [0.810855, 0.252861, 0.439305],
               [0.816914, 0.255895, 0.436461],
               [0.822926, 0.259016, 0.433573],
               [0.828886, 0.262229, 0.430644],
               [0.834791, 0.265540, 0.427671],
               [0.840636, 0.268953, 0.424666],
               [0.846416, 0.272473, 0.421631],
               [0.852126, 0.276106, 0.418573],
               [0.857763, 0.279857, 0.415496],
               [0.863320, 0.283729, 0.412403],
               [0.868793, 0.287728, 0.409303],
               [0.874176, 0.291859, 0.406205],
               [0.879464, 0.296125, 0.403118],
               [0.884651, 0.300530, 0.400047],
               [0.889731, 0.305079, 0.397002],
               [0.894700, 0.309773, 0.393995],
               [0.899552, 0.314616, 0.391037],
               [0.904281, 0.319610, 0.388137],
               [0.908884, 0.324755, 0.385308],
               [0.913354, 0.330052, 0.382563],
               [0.917689, 0.335500, 0.379915],
               [0.921884, 0.341098, 0.377376],
               [0.925937, 0.346844, 0.374959],
               [0.929845, 0.352734, 0.372677],
               [0.933606, 0.358764, 0.370541],
               [0.937221, 0.364929, 0.368567],
               [0.940687, 0.371224, 0.366762],
               [0.944006, 0.377643, 0.365136],
               [0.947180, 0.384178, 0.363701],
               [0.950210, 0.390820, 0.362468],
               [0.953099, 0.397563, 0.361438],
               [0.955849, 0.404400, 0.360619],
               [0.958464, 0.411324, 0.360014],
               [0.960949, 0.418323, 0.359630],
               [0.963310, 0.425390, 0.359469],
               [0.965549, 0.432519, 0.359529],
               [0.967671, 0.439703, 0.359810],
               [0.969680, 0.446936, 0.360311],
               [0.971582, 0.454210, 0.361030],
               [0.973381, 0.461520, 0.361965],
               [0.975082, 0.468861, 0.363111],
               [0.976690, 0.476226, 0.364466],
               [0.978210, 0.483612, 0.366025],
               [0.979645, 0.491014, 0.367783],
               [0.981000, 0.498428, 0.369734],
               [0.982279, 0.505851, 0.371874],
               [0.983485, 0.513280, 0.374198],
               [0.984622, 0.520713, 0.376698],
               [0.985693, 0.528148, 0.379371],
               [0.986700, 0.535582, 0.382210],
               [0.987646, 0.543015, 0.385210],
               [0.988533, 0.550446, 0.388365],
               [0.989363, 0.557873, 0.391671],
               [0.990138, 0.565296, 0.395122],
               [0.990871, 0.572706, 0.398714],
               [0.991558, 0.580107, 0.402441],
               [0.992196, 0.587502, 0.406299],
               [0.992785, 0.594891, 0.410283],
               [0.993326, 0.602275, 0.414390],
               [0.993834, 0.609644, 0.418613],
               [0.994309, 0.616999, 0.422950],
               [0.994738, 0.624350, 0.427397],
               [0.995122, 0.631696, 0.431951],
               [0.995480, 0.639027, 0.436607],
               [0.995810, 0.646344, 0.441361],
               [0.996096, 0.653659, 0.446213],
               [0.996341, 0.660969, 0.451160],
               [0.996580, 0.668256, 0.456192],
               [0.996775, 0.675541, 0.461314],
               [0.996925, 0.682828, 0.466526],
               [0.997077, 0.690088, 0.471811],
               [0.997186, 0.697349, 0.477182],
               [0.997254, 0.704611, 0.482635],
               [0.997325, 0.711848, 0.488154],
               [0.997351, 0.719089, 0.493755],
               [0.997351, 0.726324, 0.499428],
               [0.997341, 0.733545, 0.505167],
               [0.997285, 0.740772, 0.510983],
               [0.997228, 0.747981, 0.516859],
               [0.997138, 0.755190, 0.522806],
               [0.997019, 0.762398, 0.528821],
               [0.996898, 0.769591, 0.534892],
               [0.996727, 0.776795, 0.541039],
               [0.996571, 0.783977, 0.547233],
               [0.996369, 0.791167, 0.553499],
               [0.996162, 0.798348, 0.559820],
               [0.995932, 0.805527, 0.566202],
               [0.995680, 0.812706, 0.572645],
               [0.995424, 0.819875, 0.579140],
               [0.995131, 0.827052, 0.585701],
               [0.994851, 0.834213, 0.592307],
               [0.994524, 0.841387, 0.598983],
               [0.994222, 0.848540, 0.605696],
               [0.993866, 0.855711, 0.612482],
               [0.993545, 0.862859, 0.619299],
               [0.993170, 0.870024, 0.626189],
               [0.992831, 0.877168, 0.633109],
               [0.992440, 0.884330, 0.640099],
               [0.992089, 0.891470, 0.647116],
               [0.991688, 0.898627, 0.654202],
               [0.991332, 0.905763, 0.661309],
               [0.990930, 0.912915, 0.668481],
               [0.990570, 0.920049, 0.675675],
               [0.990175, 0.927196, 0.682926],
               [0.989815, 0.934329, 0.690198],
               [0.989434, 0.941470, 0.697519],
               [0.989077, 0.948604, 0.704863],
               [0.988717, 0.955742, 0.712242],
               [0.988367, 0.962878, 0.719649],
               [0.988033, 0.970012, 0.727077],
               [0.987691, 0.977154, 0.734536],
               [0.987387, 0.984288, 0.742002],
               [0.987053, 0.991438, 0.749504]]

_inferno_data = [[0.001462, 0.000466, 0.013866],
                 [0.002267, 0.001270, 0.018570],
                 [0.003299, 0.002249, 0.024239],
                 [0.004547, 0.003392, 0.030909],
                 [0.006006, 0.004692, 0.038558],
                 [0.007676, 0.006136, 0.046836],
                 [0.009561, 0.007713, 0.055143],
                 [0.011663, 0.009417, 0.063460],
                 [0.013995, 0.011225, 0.071862],
                 [0.016561, 0.013136, 0.080282],
                 [0.019373, 0.015133, 0.088767],
                 [0.022447, 0.017199, 0.097327],
                 [0.025793, 0.019331, 0.105930],
                 [0.029432, 0.021503, 0.114621],
                 [0.033385, 0.023702, 0.123397],
                 [0.037668, 0.025921, 0.132232],
                 [0.042253, 0.028139, 0.141141],
                 [0.046915, 0.030324, 0.150164],
                 [0.051644, 0.032474, 0.159254],
                 [0.056449, 0.034569, 0.168414],
                 [0.061340, 0.036590, 0.177642],
                 [0.066331, 0.038504, 0.186962],
                 [0.071429, 0.040294, 0.196354],
                 [0.076637, 0.041905, 0.205799],
                 [0.081962, 0.043328, 0.215289],
                 [0.087411, 0.044556, 0.224813],
                 [0.092990, 0.045583, 0.234358],
                 [0.098702, 0.046402, 0.243904],
                 [0.104551, 0.047008, 0.253430],
                 [0.110536, 0.047399, 0.262912],
                 [0.116656, 0.047574, 0.272321],
                 [0.122908, 0.047536, 0.281624],
                 [0.129285, 0.047293, 0.290788],
                 [0.135778, 0.046856, 0.299776],
                 [0.142378, 0.046242, 0.308553],
                 [0.149073, 0.045468, 0.317085],
                 [0.155850, 0.044559, 0.325338],
                 [0.162689, 0.043554, 0.333277],
                 [0.169575, 0.042489, 0.340874],
                 [0.176493, 0.041402, 0.348111],
                 [0.183429, 0.040329, 0.354971],
                 [0.190367, 0.039309, 0.361447],
                 [0.197297, 0.038400, 0.367535],
                 [0.204209, 0.037632, 0.373238],
                 [0.211095, 0.037030, 0.378563],
                 [0.217949, 0.036615, 0.383522],
                 [0.224763, 0.036405, 0.388129],
                 [0.231538, 0.036405, 0.392400],
                 [0.238273, 0.036621, 0.396353],
                 [0.244967, 0.037055, 0.400007],
                 [0.251620, 0.037705, 0.403378],
                 [0.258234, 0.038571, 0.406485],
                 [0.264810, 0.039647, 0.409345],
                 [0.271347, 0.040922, 0.411976],
                 [0.277850, 0.042353, 0.414392],
                 [0.284321, 0.043933, 0.416608],
                 [0.290763, 0.045644, 0.418637],
                 [0.297178, 0.047470, 0.420491],
                 [0.303568, 0.049396, 0.422182],
                 [0.309935, 0.051407, 0.423721],
                 [0.316282, 0.053490, 0.425116],
                 [0.322610, 0.055634, 0.426377],
                 [0.328921, 0.057827, 0.427511],
                 [0.335217, 0.060060, 0.428524],
                 [0.341500, 0.062325, 0.429425],
                 [0.347771, 0.064616, 0.430217],
                 [0.354032, 0.066925, 0.430906],
                 [0.360284, 0.069247, 0.431497],
                 [0.366529, 0.071579, 0.431994],
                 [0.372768, 0.073915, 0.432400],
                 [0.379001, 0.076253, 0.432719],
                 [0.385228, 0.078591, 0.432955],
                 [0.391453, 0.080927, 0.433109],
                 [0.397674, 0.083257, 0.433183],
                 [0.403894, 0.085580, 0.433179],
                 [0.410113, 0.087896, 0.433098],
                 [0.416331, 0.090203, 0.432943],
                 [0.422549, 0.092501, 0.432714],
                 [0.428768, 0.094790, 0.432412],
                 [0.434987, 0.097069, 0.432039],
                 [0.441207, 0.099338, 0.431594],
                 [0.447428, 0.101597, 0.431080],
                 [0.453651, 0.103848, 0.430498],
                 [0.459875, 0.106089, 0.429846],
                 [0.466100, 0.108322, 0.429125],
                 [0.472328, 0.110547, 0.428334],
                 [0.478558, 0.112764, 0.427475],
                 [0.484789, 0.114974, 0.426548],
                 [0.491022, 0.117179, 0.425552],
                 [0.497257, 0.119379, 0.424488],
                 [0.503493, 0.121575, 0.423356],
                 [0.509730, 0.123769, 0.422156],
                 [0.515967, 0.125960, 0.420887],
                 [0.522206, 0.128150, 0.419549],
                 [0.528444, 0.130341, 0.418142],
                 [0.534683, 0.132534, 0.416667],
                 [0.540920, 0.134729, 0.415123],
                 [0.547157, 0.136929, 0.413511],
                 [0.553392, 0.139134, 0.411829],
                 [0.559624, 0.141346, 0.410078],
                 [0.565854, 0.143567, 0.408258],
                 [0.572081, 0.145797, 0.406369],
                 [0.578304, 0.148039, 0.404411],
                 [0.584521, 0.150294, 0.402385],
                 [0.590734, 0.152563, 0.400290],
                 [0.596940, 0.154848, 0.398125],
                 [0.603139, 0.157151, 0.395891],
                 [0.609330, 0.159474, 0.393589],
                 [0.615513, 0.161817, 0.391219],
                 [0.621685, 0.164184, 0.388781],
                 [0.627847, 0.166575, 0.386276],
                 [0.633998, 0.168992, 0.383704],
                 [0.640135, 0.171438, 0.381065],
                 [0.646260, 0.173914, 0.378359],
                 [0.652369, 0.176421, 0.375586],
                 [0.658463, 0.178962, 0.372748],
                 [0.664540, 0.181539, 0.369846],
                 [0.670599, 0.184153, 0.366879],
                 [0.676638, 0.186807, 0.363849],
                 [0.682656, 0.189501, 0.360757],
                 [0.688653, 0.192239, 0.357603],
                 [0.694627, 0.195021, 0.354388],
                 [0.700576, 0.197851, 0.351113],
                 [0.706500, 0.200728, 0.347777],
                 [0.712396, 0.203656, 0.344383],
                 [0.718264, 0.206636, 0.340931],
                 [0.724103, 0.209670, 0.337424],
                 [0.729909, 0.212759, 0.333861],
                 [0.735683, 0.215906, 0.330245],
                 [0.741423, 0.219112, 0.326576],
                 [0.747127, 0.222378, 0.322856],
                 [0.752794, 0.225706, 0.319085],
                 [0.758422, 0.229097, 0.315266],
                 [0.764010, 0.232554, 0.311399],
                 [0.769556, 0.236077, 0.307485],
                 [0.775059, 0.239667, 0.303526],
                 [0.780517, 0.243327, 0.299523],
                 [0.785929, 0.247056, 0.295477],
                 [0.791293, 0.250856, 0.291390],
                 [0.796607, 0.254728, 0.287264],
                 [0.801871, 0.258674, 0.283099],
                 [0.807082, 0.262692, 0.278898],
                 [0.812239, 0.266786, 0.274661],
                 [0.817341, 0.270954, 0.270390],
                 [0.822386, 0.275197, 0.266085],
                 [0.827372, 0.279517, 0.261750],
                 [0.832299, 0.283913, 0.257383],
                 [0.837165, 0.288385, 0.252988],
                 [0.841969, 0.292933, 0.248564],
                 [0.846709, 0.297559, 0.244113],
                 [0.851384, 0.302260, 0.239636],
                 [0.855992, 0.307038, 0.235133],
                 [0.860533, 0.311892, 0.230606],
                 [0.865006, 0.316822, 0.226055],
                 [0.869409, 0.321827, 0.221482],
                 [0.873741, 0.326906, 0.216886],
                 [0.878001, 0.332060, 0.212268],
                 [0.882188, 0.337287, 0.207628],
                 [0.886302, 0.342586, 0.202968],
                 [0.890341, 0.347957, 0.198286],
                 [0.894305, 0.353399, 0.193584],
                 [0.898192, 0.358911, 0.188860],
                 [0.902003, 0.364492, 0.184116],
                 [0.905735, 0.370140, 0.179350],
                 [0.909390, 0.375856, 0.174563],
                 [0.912966, 0.381636, 0.169755],
                 [0.916462, 0.387481, 0.164924],
                 [0.919879, 0.393389, 0.160070],
                 [0.923215, 0.399359, 0.155193],
                 [0.926470, 0.405389, 0.150292],
                 [0.929644, 0.411479, 0.145367],
                 [0.932737, 0.417627, 0.140417],
                 [0.935747, 0.423831, 0.135440],
                 [0.938675, 0.430091, 0.130438],
                 [0.941521, 0.436405, 0.125409],
                 [0.944285, 0.442772, 0.120354],
                 [0.946965, 0.449191, 0.115272],
                 [0.949562, 0.455660, 0.110164],
                 [0.952075, 0.462178, 0.105031],
                 [0.954506, 0.468744, 0.099874],
                 [0.956852, 0.475356, 0.094695],
                 [0.959114, 0.482014, 0.089499],
                 [0.961293, 0.488716, 0.084289],
                 [0.963387, 0.495462, 0.079073],
                 [0.965397, 0.502249, 0.073859],
                 [0.967322, 0.509078, 0.068659],
                 [0.969163, 0.515946, 0.063488],
                 [0.970919, 0.522853, 0.058367],
                 [0.972590, 0.529798, 0.053324],
                 [0.974176, 0.536780, 0.048392],
                 [0.975677, 0.543798, 0.043618],
                 [0.977092, 0.550850, 0.039050],
                 [0.978422, 0.557937, 0.034931],
                 [0.979666, 0.565057, 0.031409],
                 [0.980824, 0.572209, 0.028508],
                 [0.981895, 0.579392, 0.026250],
                 [0.982881, 0.586606, 0.024661],
                 [0.983779, 0.593849, 0.023770],
                 [0.984591, 0.601122, 0.023606],
                 [0.985315, 0.608422, 0.024202],
                 [0.985952, 0.615750, 0.025592],
                 [0.986502, 0.623105, 0.027814],
                 [0.986964, 0.630485, 0.030908],
                 [0.987337, 0.637890, 0.034916],
                 [0.987622, 0.645320, 0.039886],
                 [0.987819, 0.652773, 0.045581],
                 [0.987926, 0.660250, 0.051750],
                 [0.987945, 0.667748, 0.058329],
                 [0.987874, 0.675267, 0.065257],
                 [0.987714, 0.682807, 0.072489],
                 [0.987464, 0.690366, 0.079990],
                 [0.987124, 0.697944, 0.087731],
                 [0.986694, 0.705540, 0.095694],
                 [0.986175, 0.713153, 0.103863],
                 [0.985566, 0.720782, 0.112229],
                 [0.984865, 0.728427, 0.120785],
                 [0.984075, 0.736087, 0.129527],
                 [0.983196, 0.743758, 0.138453],
                 [0.982228, 0.751442, 0.147565],
                 [0.981173, 0.759135, 0.156863],
                 [0.980032, 0.766837, 0.166353],
                 [0.978806, 0.774545, 0.176037],
                 [0.977497, 0.782258, 0.185923],
                 [0.976108, 0.789974, 0.196018],
                 [0.974638, 0.797692, 0.206332],
                 [0.973088, 0.805409, 0.216877],
                 [0.971468, 0.813122, 0.227658],
                 [0.969783, 0.820825, 0.238686],
                 [0.968041, 0.828515, 0.249972],
                 [0.966243, 0.836191, 0.261534],
                 [0.964394, 0.843848, 0.273391],
                 [0.962517, 0.851476, 0.285546],
                 [0.960626, 0.859069, 0.298010],
                 [0.958720, 0.866624, 0.310820],
                 [0.956834, 0.874129, 0.323974],
                 [0.954997, 0.881569, 0.337475],
                 [0.953215, 0.888942, 0.351369],
                 [0.951546, 0.896226, 0.365627],
                 [0.950018, 0.903409, 0.380271],
                 [0.948683, 0.910473, 0.395289],
                 [0.947594, 0.917399, 0.410665],
                 [0.946809, 0.924168, 0.426373],
                 [0.946392, 0.930761, 0.442367],
                 [0.946403, 0.937159, 0.458592],
                 [0.946903, 0.943348, 0.474970],
                 [0.947937, 0.949318, 0.491426],
                 [0.949545, 0.955063, 0.507860],
                 [0.951740, 0.960587, 0.524203],
                 [0.954529, 0.965896, 0.540361],
                 [0.957896, 0.971003, 0.556275],
                 [0.961812, 0.975924, 0.571925],
                 [0.966249, 0.980678, 0.587206],
                 [0.971162, 0.985282, 0.602154],
                 [0.976511, 0.989753, 0.616760],
                 [0.982257, 0.994109, 0.631017],
                 [0.988362, 0.998364, 0.644924]]

_plasma_data = [[0.050383, 0.029803, 0.527975],
                [0.063536, 0.028426, 0.533124],
                [0.075353, 0.027206, 0.538007],
                [0.086222, 0.026125, 0.542658],
                [0.096379, 0.025165, 0.547103],
                [0.105980, 0.024309, 0.551368],
                [0.115124, 0.023556, 0.555468],
                [0.123903, 0.022878, 0.559423],
                [0.132381, 0.022258, 0.563250],
                [0.140603, 0.021687, 0.566959],
                [0.148607, 0.021154, 0.570562],
                [0.156421, 0.020651, 0.574065],
                [0.164070, 0.020171, 0.577478],
                [0.171574, 0.019706, 0.580806],
                [0.178950, 0.019252, 0.584054],
                [0.186213, 0.018803, 0.587228],
                [0.193374, 0.018354, 0.590330],
                [0.200445, 0.017902, 0.593364],
                [0.207435, 0.017442, 0.596333],
                [0.214350, 0.016973, 0.599239],
                [0.221197, 0.016497, 0.602083],
                [0.227983, 0.016007, 0.604867],
                [0.234715, 0.015502, 0.607592],
                [0.241396, 0.014979, 0.610259],
                [0.248032, 0.014439, 0.612868],
                [0.254627, 0.013882, 0.615419],
                [0.261183, 0.013308, 0.617911],
                [0.267703, 0.012716, 0.620346],
                [0.274191, 0.012109, 0.622722],
                [0.280648, 0.011488, 0.625038],
                [0.287076, 0.010855, 0.627295],
                [0.293478, 0.010213, 0.629490],
                [0.299855, 0.009561, 0.631624],
                [0.306210, 0.008902, 0.633694],
                [0.312543, 0.008239, 0.635700],
                [0.318856, 0.007576, 0.637640],
                [0.325150, 0.006915, 0.639512],
                [0.331426, 0.006261, 0.641316],
                [0.337683, 0.005618, 0.643049],
                [0.343925, 0.004991, 0.644710],
                [0.350150, 0.004382, 0.646298],
                [0.356359, 0.003798, 0.647810],
                [0.362553, 0.003243, 0.649245],
                [0.368733, 0.002724, 0.650601],
                [0.374897, 0.002245, 0.651876],
                [0.381047, 0.001814, 0.653068],
                [0.387183, 0.001434, 0.654177],
                [0.393304, 0.001114, 0.655199],
                [0.399411, 0.000859, 0.656133],
                [0.405503, 0.000678, 0.656977],
                [0.411580, 0.000577, 0.657730],
                [0.417642, 0.000564, 0.658390],
                [0.423689, 0.000646, 0.658956],
                [0.429719, 0.000831, 0.659425],
                [0.435734, 0.001127, 0.659797],
                [0.441732, 0.001540, 0.660069],
                [0.447714, 0.002080, 0.660240],
                [0.453677, 0.002755, 0.660310],
                [0.459623, 0.003574, 0.660277],
                [0.465550, 0.004545, 0.660139],
                [0.471457, 0.005678, 0.659897],
                [0.477344, 0.006980, 0.659549],
                [0.483210, 0.008460, 0.659095],
                [0.489055, 0.010127, 0.658534],
                [0.494877, 0.011990, 0.657865],
                [0.500678, 0.014055, 0.657088],
                [0.506454, 0.016333, 0.656202],
                [0.512206, 0.018833, 0.655209],
                [0.517933, 0.021563, 0.654109],
                [0.523633, 0.024532, 0.652901],
                [0.529306, 0.027747, 0.651586],
                [0.534952, 0.031217, 0.650165],
                [0.540570, 0.034950, 0.648640],
                [0.546157, 0.038954, 0.647010],
                [0.551715, 0.043136, 0.645277],
                [0.557243, 0.047331, 0.643443],
                [0.562738, 0.051545, 0.641509],
                [0.568201, 0.055778, 0.639477],
                [0.573632, 0.060028, 0.637349],
                [0.579029, 0.064296, 0.635126],
                [0.584391, 0.068579, 0.632812],
                [0.589719, 0.072878, 0.630408],
                [0.595011, 0.077190, 0.627917],
                [0.600266, 0.081516, 0.625342],
                [0.605485, 0.085854, 0.622686],
                [0.610667, 0.090204, 0.619951],
                [0.615812, 0.094564, 0.617140],
                [0.620919, 0.098934, 0.614257],
                [0.625987, 0.103312, 0.611305],
                [0.631017, 0.107699, 0.608287],
                [0.636008, 0.112092, 0.605205],
                [0.640959, 0.116492, 0.602065],
                [0.645872, 0.120898, 0.598867],
                [0.650746, 0.125309, 0.595617],
                [0.655580, 0.129725, 0.592317],
                [0.660374, 0.134144, 0.588971],
                [0.665129, 0.138566, 0.585582],
                [0.669845, 0.142992, 0.582154],
                [0.674522, 0.147419, 0.578688],
                [0.679160, 0.151848, 0.575189],
                [0.683758, 0.156278, 0.571660],
                [0.688318, 0.160709, 0.568103],
                [0.692840, 0.165141, 0.564522],
                [0.697324, 0.169573, 0.560919],
                [0.701769, 0.174005, 0.557296],
                [0.706178, 0.178437, 0.553657],
                [0.710549, 0.182868, 0.550004],
                [0.714883, 0.187299, 0.546338],
                [0.719181, 0.191729, 0.542663],
                [0.723444, 0.196158, 0.538981],
                [0.727670, 0.200586, 0.535293],
                [0.731862, 0.205013, 0.531601],
                [0.736019, 0.209439, 0.527908],
                [0.740143, 0.213864, 0.524216],
                [0.744232, 0.218288, 0.520524],
                [0.748289, 0.222711, 0.516834],
                [0.752312, 0.227133, 0.513149],
                [0.756304, 0.231555, 0.509468],
                [0.760264, 0.235976, 0.505794],
                [0.764193, 0.240396, 0.502126],
                [0.768090, 0.244817, 0.498465],
                [0.771958, 0.249237, 0.494813],
                [0.775796, 0.253658, 0.491171],
                [0.779604, 0.258078, 0.487539],
                [0.783383, 0.262500, 0.483918],
                [0.787133, 0.266922, 0.480307],
                [0.790855, 0.271345, 0.476706],
                [0.794549, 0.275770, 0.473117],
                [0.798216, 0.280197, 0.469538],
                [0.801855, 0.284626, 0.465971],
                [0.805467, 0.289057, 0.462415],
                [0.809052, 0.293491, 0.458870],
                [0.812612, 0.297928, 0.455338],
                [0.816144, 0.302368, 0.451816],
                [0.819651, 0.306812, 0.448306],
                [0.823132, 0.311261, 0.444806],
                [0.826588, 0.315714, 0.441316],
                [0.830018, 0.320172, 0.437836],
                [0.833422, 0.324635, 0.434366],
                [0.836801, 0.329105, 0.430905],
                [0.840155, 0.333580, 0.427455],
                [0.843484, 0.338062, 0.424013],
                [0.846788, 0.342551, 0.420579],
                [0.850066, 0.347048, 0.417153],
                [0.853319, 0.351553, 0.413734],
                [0.856547, 0.356066, 0.410322],
                [0.859750, 0.360588, 0.406917],
                [0.862927, 0.365119, 0.403519],
                [0.866078, 0.369660, 0.400126],
                [0.869203, 0.374212, 0.396738],
                [0.872303, 0.378774, 0.393355],
                [0.875376, 0.383347, 0.389976],
                [0.878423, 0.387932, 0.386600],
                [0.881443, 0.392529, 0.383229],
                [0.884436, 0.397139, 0.379860],
                [0.887402, 0.401762, 0.376494],
                [0.890340, 0.406398, 0.373130],
                [0.893250, 0.411048, 0.369768],
                [0.896131, 0.415712, 0.366407],
                [0.898984, 0.420392, 0.363047],
                [0.901807, 0.425087, 0.359688],
                [0.904601, 0.429797, 0.356329],
                [0.907365, 0.434524, 0.352970],
                [0.910098, 0.439268, 0.349610],
                [0.912800, 0.444029, 0.346251],
                [0.915471, 0.448807, 0.342890],
                [0.918109, 0.453603, 0.339529],
                [0.920714, 0.458417, 0.336166],
                [0.923287, 0.463251, 0.332801],
                [0.925825, 0.468103, 0.329435],
                [0.928329, 0.472975, 0.326067],
                [0.930798, 0.477867, 0.322697],
                [0.933232, 0.482780, 0.319325],
                [0.935630, 0.487712, 0.315952],
                [0.937990, 0.492667, 0.312575],
                [0.940313, 0.497642, 0.309197],
                [0.942598, 0.502639, 0.305816],
                [0.944844, 0.507658, 0.302433],
                [0.947051, 0.512699, 0.299049],
                [0.949217, 0.517763, 0.295662],
                [0.951344, 0.522850, 0.292275],
                [0.953428, 0.527960, 0.288883],
                [0.955470, 0.533093, 0.285490],
                [0.957469, 0.538250, 0.282096],
                [0.959424, 0.543431, 0.278701],
                [0.961336, 0.548636, 0.275305],
                [0.963203, 0.553865, 0.271909],
                [0.965024, 0.559118, 0.268513],
                [0.966798, 0.564396, 0.265118],
                [0.968526, 0.569700, 0.261721],
                [0.970205, 0.575028, 0.258325],
                [0.971835, 0.580382, 0.254931],
                [0.973416, 0.585761, 0.251540],
                [0.974947, 0.591165, 0.248151],
                [0.976428, 0.596595, 0.244767],
                [0.977856, 0.602051, 0.241387],
                [0.979233, 0.607532, 0.238013],
                [0.980556, 0.613039, 0.234646],
                [0.981826, 0.618572, 0.231287],
                [0.983041, 0.624131, 0.227937],
                [0.984199, 0.629718, 0.224595],
                [0.985301, 0.635330, 0.221265],
                [0.986345, 0.640969, 0.217948],
                [0.987332, 0.646633, 0.214648],
                [0.988260, 0.652325, 0.211364],
                [0.989128, 0.658043, 0.208100],
                [0.989935, 0.663787, 0.204859],
                [0.990681, 0.669558, 0.201642],
                [0.991365, 0.675355, 0.198453],
                [0.991985, 0.681179, 0.195295],
                [0.992541, 0.687030, 0.192170],
                [0.993032, 0.692907, 0.189084],
                [0.993456, 0.698810, 0.186041],
                [0.993814, 0.704741, 0.183043],
                [0.994103, 0.710698, 0.180097],
                [0.994324, 0.716681, 0.177208],
                [0.994474, 0.722691, 0.174381],
                [0.994553, 0.728728, 0.171622],
                [0.994561, 0.734791, 0.168938],
                [0.994495, 0.740880, 0.166335],
                [0.994355, 0.746995, 0.163821],
                [0.994141, 0.753137, 0.161404],
                [0.993851, 0.759304, 0.159092],
                [0.993482, 0.765499, 0.156891],
                [0.993033, 0.771720, 0.154808],
                [0.992505, 0.777967, 0.152855],
                [0.991897, 0.784239, 0.151042],
                [0.991209, 0.790537, 0.149377],
                [0.990439, 0.796859, 0.147870],
                [0.989587, 0.803205, 0.146529],
                [0.988648, 0.809579, 0.145357],
                [0.987621, 0.815978, 0.144363],
                [0.986509, 0.822401, 0.143557],
                [0.985314, 0.828846, 0.142945],
                [0.984031, 0.835315, 0.142528],
                [0.982653, 0.841812, 0.142303],
                [0.981190, 0.848329, 0.142279],
                [0.979644, 0.854866, 0.142453],
                [0.977995, 0.861432, 0.142808],
                [0.976265, 0.868016, 0.143351],
                [0.974443, 0.874622, 0.144061],
                [0.972530, 0.881250, 0.144923],
                [0.970533, 0.887896, 0.145919],
                [0.968443, 0.894564, 0.147014],
                [0.966271, 0.901249, 0.148180],
                [0.964021, 0.907950, 0.149370],
                [0.961681, 0.914672, 0.150520],
                [0.959276, 0.921407, 0.151566],
                [0.956808, 0.928152, 0.152409],
                [0.954287, 0.934908, 0.152921],
                [0.951726, 0.941671, 0.152925],
                [0.949151, 0.948435, 0.152178],
                [0.946602, 0.955190, 0.150328],
                [0.944152, 0.961916, 0.146861],
                [0.941896, 0.968590, 0.140956],
                [0.940015, 0.975158, 0.131326]]

_viridis_data = [[0.267004, 0.004874, 0.329415],
                 [0.268510, 0.009605, 0.335427],
                 [0.269944, 0.014625, 0.341379],
                 [0.271305, 0.019942, 0.347269],
                 [0.272594, 0.025563, 0.353093],
                 [0.273809, 0.031497, 0.358853],
                 [0.274952, 0.037752, 0.364543],
                 [0.276022, 0.044167, 0.370164],
                 [0.277018, 0.050344, 0.375715],
                 [0.277941, 0.056324, 0.381191],
                 [0.278791, 0.062145, 0.386592],
                 [0.279566, 0.067836, 0.391917],
                 [0.280267, 0.073417, 0.397163],
                 [0.280894, 0.078907, 0.402329],
                 [0.281446, 0.084320, 0.407414],
                 [0.281924, 0.089666, 0.412415],
                 [0.282327, 0.094955, 0.417331],
                 [0.282656, 0.100196, 0.422160],
                 [0.282910, 0.105393, 0.426902],
                 [0.283091, 0.110553, 0.431554],
                 [0.283197, 0.115680, 0.436115],
                 [0.283229, 0.120777, 0.440584],
                 [0.283187, 0.125848, 0.444960],
                 [0.283072, 0.130895, 0.449241],
                 [0.282884, 0.135920, 0.453427],
                 [0.282623, 0.140926, 0.457517],
                 [0.282290, 0.145912, 0.461510],
                 [0.281887, 0.150881, 0.465405],
                 [0.281412, 0.155834, 0.469201],
                 [0.280868, 0.160771, 0.472899],
                 [0.280255, 0.165693, 0.476498],
                 [0.279574, 0.170599, 0.479997],
                 [0.278826, 0.175490, 0.483397],
                 [0.278012, 0.180367, 0.486697],
                 [0.277134, 0.185228, 0.489898],
                 [0.276194, 0.190074, 0.493001],
                 [0.275191, 0.194905, 0.496005],
                 [0.274128, 0.199721, 0.498911],
                 [0.273006, 0.204520, 0.501721],
                 [0.271828, 0.209303, 0.504434],
                 [0.270595, 0.214069, 0.507052],
                 [0.269308, 0.218818, 0.509577],
                 [0.267968, 0.223549, 0.512008],
                 [0.266580, 0.228262, 0.514349],
                 [0.265145, 0.232956, 0.516599],
                 [0.263663, 0.237631, 0.518762],
                 [0.262138, 0.242286, 0.520837],
                 [0.260571, 0.246922, 0.522828],
                 [0.258965, 0.251537, 0.524736],
                 [0.257322, 0.256130, 0.526563],
                 [0.255645, 0.260703, 0.528312],
                 [0.253935, 0.265254, 0.529983],
                 [0.252194, 0.269783, 0.531579],
                 [0.250425, 0.274290, 0.533103],
                 [0.248629, 0.278775, 0.534556],
                 [0.246811, 0.283237, 0.535941],
                 [0.244972, 0.287675, 0.537260],
                 [0.243113, 0.292092, 0.538516],
                 [0.241237, 0.296485, 0.539709],
                 [0.239346, 0.300855, 0.540844],
                 [0.237441, 0.305202, 0.541921],
                 [0.235526, 0.309527, 0.542944],
                 [0.233603, 0.313828, 0.543914],
                 [0.231674, 0.318106, 0.544834],
                 [0.229739, 0.322361, 0.545706],
                 [0.227802, 0.326594, 0.546532],
                 [0.225863, 0.330805, 0.547314],
                 [0.223925, 0.334994, 0.548053],
                 [0.221989, 0.339161, 0.548752],
                 [0.220057, 0.343307, 0.549413],
                 [0.218130, 0.347432, 0.550038],
                 [0.216210, 0.351535, 0.550627],
                 [0.214298, 0.355619, 0.551184],
                 [0.212395, 0.359683, 0.551710],
                 [0.210503, 0.363727, 0.552206],
                 [0.208623, 0.367752, 0.552675],
                 [0.206756, 0.371758, 0.553117],
                 [0.204903, 0.375746, 0.553533],
                 [0.203063, 0.379716, 0.553925],
                 [0.201239, 0.383670, 0.554294],
                 [0.199430, 0.387607, 0.554642],
                 [0.197636, 0.391528, 0.554969],
                 [0.195860, 0.395433, 0.555276],
                 [0.194100, 0.399323, 0.555565],
                 [0.192357, 0.403199, 0.555836],
                 [0.190631, 0.407061, 0.556089],
                 [0.188923, 0.410910, 0.556326],
                 [0.187231, 0.414746, 0.556547],
                 [0.185556, 0.418570, 0.556753],
                 [0.183898, 0.422383, 0.556944],
                 [0.182256, 0.426184, 0.557120],
                 [0.180629, 0.429975, 0.557282],
                 [0.179019, 0.433756, 0.557430],
                 [0.177423, 0.437527, 0.557565],
                 [0.175841, 0.441290, 0.557685],
                 [0.174274, 0.445044, 0.557792],
                 [0.172719, 0.448791, 0.557885],
                 [0.171176, 0.452530, 0.557965],
                 [0.169646, 0.456262, 0.558030],
                 [0.168126, 0.459988, 0.558082],
                 [0.166617, 0.463708, 0.558119],
                 [0.165117, 0.467423, 0.558141],
                 [0.163625, 0.471133, 0.558148],
                 [0.162142, 0.474838, 0.558140],
                 [0.160665, 0.478540, 0.558115],
                 [0.159194, 0.482237, 0.558073],
                 [0.157729, 0.485932, 0.558013],
                 [0.156270, 0.489624, 0.557936],
                 [0.154815, 0.493313, 0.557840],
                 [0.153364, 0.497000, 0.557724],
                 [0.151918, 0.500685, 0.557587],
                 [0.150476, 0.504369, 0.557430],
                 [0.149039, 0.508051, 0.557250],
                 [0.147607, 0.511733, 0.557049],
                 [0.146180, 0.515413, 0.556823],
                 [0.144759, 0.519093, 0.556572],
                 [0.143343, 0.522773, 0.556295],
                 [0.141935, 0.526453, 0.555991],
                 [0.140536, 0.530132, 0.555659],
                 [0.139147, 0.533812, 0.555298],
                 [0.137770, 0.537492, 0.554906],
                 [0.136408, 0.541173, 0.554483],
                 [0.135066, 0.544853, 0.554029],
                 [0.133743, 0.548535, 0.553541],
                 [0.132444, 0.552216, 0.553018],
                 [0.131172, 0.555899, 0.552459],
                 [0.129933, 0.559582, 0.551864],
                 [0.128729, 0.563265, 0.551229],
                 [0.127568, 0.566949, 0.550556],
                 [0.126453, 0.570633, 0.549841],
                 [0.125394, 0.574318, 0.549086],
                 [0.124395, 0.578002, 0.548287],
                 [0.123463, 0.581687, 0.547445],
                 [0.122606, 0.585371, 0.546557],
                 [0.121831, 0.589055, 0.545623],
                 [0.121148, 0.592739, 0.544641],
                 [0.120565, 0.596422, 0.543611],
                 [0.120092, 0.600104, 0.542530],
                 [0.119738, 0.603785, 0.541400],
                 [0.119512, 0.607464, 0.540218],
                 [0.119423, 0.611141, 0.538982],
                 [0.119483, 0.614817, 0.537692],
                 [0.119699, 0.618490, 0.536347],
                 [0.120081, 0.622161, 0.534946],
                 [0.120638, 0.625828, 0.533488],
                 [0.121380, 0.629492, 0.531973],
                 [0.122312, 0.633153, 0.530398],
                 [0.123444, 0.636809, 0.528763],
                 [0.124780, 0.640461, 0.527068],
                 [0.126326, 0.644107, 0.525311],
                 [0.128087, 0.647749, 0.523491],
                 [0.130067, 0.651384, 0.521608],
                 [0.132268, 0.655014, 0.519661],
                 [0.134692, 0.658636, 0.517649],
                 [0.137339, 0.662252, 0.515571],
                 [0.140210, 0.665859, 0.513427],
                 [0.143303, 0.669459, 0.511215],
                 [0.146616, 0.673050, 0.508936],
                 [0.150148, 0.676631, 0.506589],
                 [0.153894, 0.680203, 0.504172],
                 [0.157851, 0.683765, 0.501686],
                 [0.162016, 0.687316, 0.499129],
                 [0.166383, 0.690856, 0.496502],
                 [0.170948, 0.694384, 0.493803],
                 [0.175707, 0.697900, 0.491033],
                 [0.180653, 0.701402, 0.488189],
                 [0.185783, 0.704891, 0.485273],
                 [0.191090, 0.708366, 0.482284],
                 [0.196571, 0.711827, 0.479221],
                 [0.202219, 0.715272, 0.476084],
                 [0.208030, 0.718701, 0.472873],
                 [0.214000, 0.722114, 0.469588],
                 [0.220124, 0.725509, 0.466226],
                 [0.226397, 0.728888, 0.462789],
                 [0.232815, 0.732247, 0.459277],
                 [0.239374, 0.735588, 0.455688],
                 [0.246070, 0.738910, 0.452024],
                 [0.252899, 0.742211, 0.448284],
                 [0.259857, 0.745492, 0.444467],
                 [0.266941, 0.748751, 0.440573],
                 [0.274149, 0.751988, 0.436601],
                 [0.281477, 0.755203, 0.432552],
                 [0.288921, 0.758394, 0.428426],
                 [0.296479, 0.761561, 0.424223],
                 [0.304148, 0.764704, 0.419943],
                 [0.311925, 0.767822, 0.415586],
                 [0.319809, 0.770914, 0.411152],
                 [0.327796, 0.773980, 0.406640],
                 [0.335885, 0.777018, 0.402049],
                 [0.344074, 0.780029, 0.397381],
                 [0.352360, 0.783011, 0.392636],
                 [0.360741, 0.785964, 0.387814],
                 [0.369214, 0.788888, 0.382914],
                 [0.377779, 0.791781, 0.377939],
                 [0.386433, 0.794644, 0.372886],
                 [0.395174, 0.797475, 0.367757],
                 [0.404001, 0.800275, 0.362552],
                 [0.412913, 0.803041, 0.357269],
                 [0.421908, 0.805774, 0.351910],
                 [0.430983, 0.808473, 0.346476],
                 [0.440137, 0.811138, 0.340967],
                 [0.449368, 0.813768, 0.335384],
                 [0.458674, 0.816363, 0.329727],
                 [0.468053, 0.818921, 0.323998],
                 [0.477504, 0.821444, 0.318195],
                 [0.487026, 0.823929, 0.312321],
                 [0.496615, 0.826376, 0.306377],
                 [0.506271, 0.828786, 0.300362],
                 [0.515992, 0.831158, 0.294279],
                 [0.525776, 0.833491, 0.288127],
                 [0.535621, 0.835785, 0.281908],
                 [0.545524, 0.838039, 0.275626],
                 [0.555484, 0.840254, 0.269281],
                 [0.565498, 0.842430, 0.262877],
                 [0.575563, 0.844566, 0.256415],
                 [0.585678, 0.846661, 0.249897],
                 [0.595839, 0.848717, 0.243329],
                 [0.606045, 0.850733, 0.236712],
                 [0.616293, 0.852709, 0.230052],
                 [0.626579, 0.854645, 0.223353],
                 [0.636902, 0.856542, 0.216620],
                 [0.647257, 0.858400, 0.209861],
                 [0.657642, 0.860219, 0.203082],
                 [0.668054, 0.861999, 0.196293],
                 [0.678489, 0.863742, 0.189503],
                 [0.688944, 0.865448, 0.182725],
                 [0.699415, 0.867117, 0.175971],
                 [0.709898, 0.868751, 0.169257],
                 [0.720391, 0.870350, 0.162603],
                 [0.730889, 0.871916, 0.156029],
                 [0.741388, 0.873449, 0.149561],
                 [0.751884, 0.874951, 0.143228],
                 [0.762373, 0.876424, 0.137064],
                 [0.772852, 0.877868, 0.131109],
                 [0.783315, 0.879285, 0.125405],
                 [0.793760, 0.880678, 0.120005],
                 [0.804182, 0.882046, 0.114965],
                 [0.814576, 0.883393, 0.110347],
                 [0.824940, 0.884720, 0.106217],
                 [0.835270, 0.886029, 0.102646],
                 [0.845561, 0.887322, 0.099702],
                 [0.855810, 0.888601, 0.097452],
                 [0.866013, 0.889868, 0.095953],
                 [0.876168, 0.891125, 0.095250],
                 [0.886271, 0.892374, 0.095374],
                 [0.896320, 0.893616, 0.096335],
                 [0.906311, 0.894855, 0.098125],
                 [0.916242, 0.896091, 0.100717],
                 [0.926106, 0.897330, 0.104071],
                 [0.935904, 0.898570, 0.108131],
                 [0.945636, 0.899815, 0.112838],
                 [0.955300, 0.901065, 0.118128],
                 [0.964894, 0.902323, 0.123941],
                 [0.974417, 0.903590, 0.130215],
                 [0.983868, 0.904867, 0.136897],
                 [0.993248, 0.906157, 0.143936]]


cmaps = {}
for (name, data) in (('magma', _magma_data),
                     ('inferno', _inferno_data),
                     ('plasma', _plasma_data),
                     ('viridis', _viridis_data)):

    cmaps[name] = ListedColormap(data, name=name)
    # generate reversed colormap
    name = name + '_r'
    cmaps[name] = ListedColormap(list(reversed(data)), name=name)

social_vdgnn/lib/python3.6/site-packages/matplotlib/_cm_listed.py
eliatrevisan/IA-MPPI-LBM
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six
from six.moves import xrange
from itertools import chain, product
from distutils.version import LooseVersion
import io

import datetime

import pytz

import numpy as np
from numpy import ma
from numpy import arange
from cycler import cycler
import pytest

import warnings

import matplotlib
from matplotlib.testing.decorators import image_comparison
import matplotlib.pyplot as plt
import matplotlib.markers as mmarkers
import matplotlib.patches as mpatches
import matplotlib.colors as mcolors
from numpy.testing import assert_allclose, assert_array_equal
from matplotlib.cbook import IgnoredKeywordWarning
from matplotlib.cbook._backports import broadcast_to

# Note: Some test cases are run twice: once normally and once with labeled data
#       These two must be defined in the same test function or need to have
#       different baseline images to prevent race conditions when pytest runs
#       the tests with multiple threads.


@image_comparison(baseline_images=['formatter_ticker_001',
                                   'formatter_ticker_002',
                                   'formatter_ticker_003',
                                   'formatter_ticker_004',
                                   'formatter_ticker_005',
                                   ])
def test_formatter_ticker():
    import matplotlib.testing.jpl_units as units
    units.register()

    # This should affect the tick size.  (Tests issue #543)
    matplotlib.rcParams['lines.markeredgewidth'] = 30

    # This essentially test to see if user specified labels get overwritten
    # by the auto labeler functionality of the axes.
    xdata = [x*units.sec for x in range(10)]
    ydata1 = [(1.5*y - 0.5)*units.km for y in range(10)]
    ydata2 = [(1.75*y - 1.0)*units.km for y in range(10)]

    fig = plt.figure()
    ax = plt.subplot(111)
    ax.set_xlabel("x-label 001")

    fig = plt.figure()
    ax = plt.subplot(111)
    ax.set_xlabel("x-label 001")
    ax.plot(xdata, ydata1, color='blue', xunits="sec")

    fig = plt.figure()
    ax = plt.subplot(111)
    ax.set_xlabel("x-label 001")
    ax.plot(xdata, ydata1, color='blue', xunits="sec")
    ax.set_xlabel("x-label 003")

    fig = plt.figure()
    ax = plt.subplot(111)
    ax.plot(xdata, ydata1, color='blue', xunits="sec")
    ax.plot(xdata, ydata2, color='green', xunits="hour")
    ax.set_xlabel("x-label 004")

    # See SF bug 2846058
    # https://sourceforge.net/tracker/?func=detail&aid=2846058&group_id=80706&atid=560720
    fig = plt.figure()
    ax = plt.subplot(111)
    ax.plot(xdata, ydata1, color='blue', xunits="sec")
    ax.plot(xdata, ydata2, color='green', xunits="hour")
    ax.set_xlabel("x-label 005")
    ax.autoscale_view()


@image_comparison(baseline_images=["formatter_large_small"])
def test_formatter_large_small():
    # github issue #617, pull #619
    if LooseVersion(np.__version__) >= LooseVersion('1.11.0'):
        pytest.skip("Fall out from a fixed numpy bug")
    fig, ax = plt.subplots(1)
    x = [0.500000001, 0.500000002]
    y = [1e64, 1.1e64]
    ax.plot(x, y)


@image_comparison(baseline_images=["twin_axis_locaters_formatters"])
def test_twin_axis_locaters_formatters():
    vals = np.linspace(0, 1, num=5, endpoint=True)
    locs = np.sin(np.pi * vals / 2.0)

    majl = plt.FixedLocator(locs)
    minl = plt.FixedLocator([0.1, 0.2, 0.3])

    fig = plt.figure()
    ax1 = fig.add_subplot(1, 1, 1)
    ax1.plot([0.1, 100], [0, 1])
    ax1.yaxis.set_major_locator(majl)
    ax1.yaxis.set_minor_locator(minl)
    ax1.yaxis.set_major_formatter(plt.FormatStrFormatter('%08.2lf'))
    ax1.yaxis.set_minor_formatter(plt.FixedFormatter(['tricks', 'mind',
                                                      'jedi']))

    ax1.xaxis.set_major_locator(plt.LinearLocator())
    ax1.xaxis.set_minor_locator(plt.FixedLocator([15, 35, 55, 75]))
    ax1.xaxis.set_major_formatter(plt.FormatStrFormatter('%05.2lf'))
    ax1.xaxis.set_minor_formatter(plt.FixedFormatter(['c', '3', 'p', 'o']))
    ax2 = ax1.twiny()
    ax3 = ax1.twinx()


def test_twinx_cla():
    fig, ax = plt.subplots()
    ax2 = ax.twinx()
    ax3 = ax2.twiny()
    plt.draw()
    assert not ax2.xaxis.get_visible()
    assert not ax2.patch.get_visible()
    ax2.cla()
    ax3.cla()

    assert not ax2.xaxis.get_visible()
    assert not ax2.patch.get_visible()
    assert ax2.yaxis.get_visible()

    assert ax3.xaxis.get_visible()
    assert not ax3.patch.get_visible()
    assert not ax3.yaxis.get_visible()

    assert ax.xaxis.get_visible()
    assert ax.patch.get_visible()
    assert ax.yaxis.get_visible()


@image_comparison(baseline_images=['twin_autoscale'], extensions=['png'])
def test_twinx_axis_scales():
    x = np.array([0, 0.5, 1])
    y = 0.5 * x
    x2 = np.array([0, 1, 2])
    y2 = 2 * x2

    fig = plt.figure()
    ax = fig.add_axes((0, 0, 1, 1), autoscalex_on=False, autoscaley_on=False)
    ax.plot(x, y, color='blue', lw=10)

    ax2 = plt.twinx(ax)
    ax2.plot(x2, y2, 'r--', lw=5)

    ax.margins(0, 0)
    ax2.margins(0, 0)


def test_twin_inherit_autoscale_setting():
    fig, ax = plt.subplots()
    ax_x_on = ax.twinx()
    ax.set_autoscalex_on(False)
    ax_x_off = ax.twinx()

    assert ax_x_on.get_autoscalex_on()
    assert not ax_x_off.get_autoscalex_on()

    ax_y_on = ax.twiny()
    ax.set_autoscaley_on(False)
    ax_y_off = ax.twiny()

    assert ax_y_on.get_autoscaley_on()
    assert not ax_y_off.get_autoscaley_on()


def test_inverted_cla():
    # Github PR #5450. Setting autoscale should reset
    # axes to be non-inverted.
    # plotting an image, then 1d graph, axis is now down
    fig = plt.figure(0)
    ax = fig.gca()
    # 1. test that a new axis is not inverted per default
    assert not(ax.xaxis_inverted())
    assert not(ax.yaxis_inverted())
    img = np.random.random((100, 100))
    ax.imshow(img)
    # 2. test that a image axis is inverted
    assert not(ax.xaxis_inverted())
    assert ax.yaxis_inverted()
    # 3. test that clearing and plotting a line, axes are
    # not inverted
    ax.cla()
    x = np.linspace(0, 2*np.pi, 100)
    ax.plot(x, np.cos(x))
    assert not(ax.xaxis_inverted())
    assert not(ax.yaxis_inverted())

    # 4. autoscaling should not bring back axes to normal
    ax.cla()
    ax.imshow(img)
    plt.autoscale()
    assert not(ax.xaxis_inverted())
    assert ax.yaxis_inverted()

    # 5. two shared axes. Clearing the master axis should bring axes in shared
    # axes back to normal
    ax0 = plt.subplot(211)
    ax1 = plt.subplot(212, sharey=ax0)
    ax0.imshow(img)
    ax1.plot(x, np.cos(x))
    ax0.cla()
    assert not(ax1.yaxis_inverted())
    ax1.cla()
    # 6. clearing the nonmaster should not touch limits
    ax0.imshow(img)
    ax1.plot(x, np.cos(x))
    ax1.cla()
    assert ax.yaxis_inverted()

    # clean up
    plt.close(fig)


@image_comparison(baseline_images=["minorticks_on_rcParams_both"],
                  extensions=['png'])
def test_minorticks_on_rcParams_both():
    fig = plt.figure()
    matplotlib.rcParams['xtick.minor.visible'] = True
    matplotlib.rcParams['ytick.minor.visible'] = True

    plt.plot([0, 1], [0, 1])
    plt.axis([0, 1, 0, 1])


@image_comparison(baseline_images=["autoscale_tiny_range"], remove_text=True)
def test_autoscale_tiny_range():
    # github pull #904
    fig, ax = plt.subplots(2, 2)
    ax = ax.flatten()
    for i in xrange(4):
        y1 = 10**(-11 - i)
        ax[i].plot([0, 1], [1, 1 + y1])


@pytest.mark.style('default')
def test_autoscale_tight():
    fig, ax = plt.subplots(1, 1)
    ax.plot([1, 2, 3, 4])
    ax.autoscale(enable=True, axis='x', tight=False)
    ax.autoscale(enable=True, axis='y', tight=True)
    assert_allclose(ax.get_xlim(), (-0.15, 3.15))
    assert_allclose(ax.get_ylim(), (1.0, 4.0))


@pytest.mark.style('default')
def test_autoscale_log_shared():
    # related to github #7587
    # array starts at zero to trigger _minpos handling
    x = np.arange(100, dtype=float)
    fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True)
    ax1.loglog(x, x)
    ax2.semilogx(x, x)
    ax1.autoscale(tight=True)
    ax2.autoscale(tight=True)
    plt.draw()
    lims = (x[1], x[-1])
    assert_allclose(ax1.get_xlim(), lims)
    assert_allclose(ax1.get_ylim(), lims)
    assert_allclose(ax2.get_xlim(), lims)
    assert_allclose(ax2.get_ylim(), (x[0], x[-1]))


@pytest.mark.style('default')
def test_use_sticky_edges():
    fig, ax = plt.subplots()
    ax.imshow([[0, 1], [2, 3]], origin='lower')
    assert_allclose(ax.get_xlim(), (-0.5, 1.5))
    assert_allclose(ax.get_ylim(), (-0.5, 1.5))
    ax.use_sticky_edges = False
    ax.autoscale()
    xlim = (-0.5 - 2 * ax._xmargin, 1.5 + 2 * ax._xmargin)
    ylim = (-0.5 - 2 * ax._ymargin, 1.5 + 2 * ax._ymargin)
    assert_allclose(ax.get_xlim(), xlim)
    assert_allclose(ax.get_ylim(), ylim)
    # Make sure it is reversible:
    ax.use_sticky_edges = True
    ax.autoscale()
    assert_allclose(ax.get_xlim(), (-0.5, 1.5))
    assert_allclose(ax.get_ylim(), (-0.5, 1.5))


@image_comparison(baseline_images=['offset_points'],
                  remove_text=True)
def test_basic_annotate():
    # Setup some data
    t = np.arange(0.0, 5.0, 0.01)
    s = np.cos(2.0*np.pi * t)

    # Offset Points

    fig = plt.figure()
    ax = fig.add_subplot(111, autoscale_on=False, xlim=(-1, 5), ylim=(-3, 5))
    line, = ax.plot(t, s, lw=3, color='purple')

    ax.annotate('local max', xy=(3, 1), xycoords='data',
                xytext=(3, 3), textcoords='offset points')


@image_comparison(baseline_images=['arrow_simple'],
                  extensions=['png'], remove_text=True)
def test_arrow_simple():
    # Simple image test for ax.arrow
    # kwargs that take discrete values
    length_includes_head = (True, False)
    shape = ('full', 'left', 'right')
    head_starts_at_zero = (True, False)
    # Create outer product of values
    kwargs = list(product(length_includes_head, shape, head_starts_at_zero))

    fig, axs = plt.subplots(3, 4)
    for i, (ax, kwarg) in enumerate(zip(axs.flatten(), kwargs)):
        ax.set_xlim(-2, 2)
        ax.set_ylim(-2, 2)
        # Unpack kwargs
        (length_includes_head, shape, head_starts_at_zero) = kwarg
        theta = 2 * np.pi * i / 12
        # Draw arrow
        ax.arrow(0, 0, np.sin(theta), np.cos(theta),
                 width=theta/100,
                 length_includes_head=length_includes_head,
                 shape=shape,
                 head_starts_at_zero=head_starts_at_zero,
                 head_width=theta / 10,
                 head_length=theta / 10)


def test_annotate_default_arrow():
    # Check that we can make an annotation arrow with only default properties.
    fig, ax = plt.subplots()
    ann = ax.annotate("foo", (0, 1), xytext=(2, 3))
    assert ann.arrow_patch is None
    ann = ax.annotate("foo", (0, 1), xytext=(2, 3), arrowprops={})
    assert ann.arrow_patch is not None


@image_comparison(baseline_images=['polar_axes'], style='default')
def test_polar_annotations():
    # you can specify the xypoint and the xytext in different
    # positions and coordinate systems, and optionally turn on a
    # connecting line and mark the point with a marker.  Annotations
    # work on polar axes too.  In the example below, the xy point is
    # in native coordinates (xycoords defaults to 'data').  For a
    # polar axes, this is in (theta, radius) space.  The text in this
    # example is placed in the fractional figure coordinate system.
    # Text keyword args like horizontal and vertical alignment are
    # respected

    # Setup some data
    r = np.arange(0.0, 1.0, 0.001)
    theta = 2.0 * 2.0 * np.pi * r

    fig = plt.figure()
    ax = fig.add_subplot(111, polar=True)
    line, = ax.plot(theta, r, color='#ee8d18', lw=3)
    line, = ax.plot((0, 0), (0, 1), color="#0000ff", lw=1)

    ind = 800
    thisr, thistheta = r[ind], theta[ind]
    ax.plot([thistheta], [thisr], 'o')
    ax.annotate('a polar annotation',
                xy=(thistheta, thisr),  # theta, radius
                xytext=(0.05, 0.05),    # fraction, fraction
                textcoords='figure fraction',
                arrowprops=dict(facecolor='black', shrink=0.05),
                horizontalalignment='left',
                verticalalignment='baseline',
                )

    ax.tick_params(axis='x', tick1On=True, tick2On=True, direction='out')


@image_comparison(baseline_images=['polar_coords'], style='default',
                  remove_text=True)
def test_polar_coord_annotations():
    # You can also use polar notation on a catesian axes.  Here the
    # native coordinate system ('data') is cartesian, so you need to
    # specify the xycoords and textcoords as 'polar' if you want to
    # use (theta, radius)
    from matplotlib.patches import Ellipse
    el = Ellipse((0, 0), 10, 20, facecolor='r', alpha=0.5)

    fig = plt.figure()
    ax = fig.add_subplot(111, aspect='equal')

    ax.add_artist(el)
    el.set_clip_box(ax.bbox)

    ax.annotate('the top',
                xy=(np.pi/2., 10.),      # theta, radius
                xytext=(np.pi/3, 20.),   # theta, radius
                xycoords='polar',
                textcoords='polar',
                arrowprops=dict(facecolor='black', shrink=0.05),
                horizontalalignment='left',
                verticalalignment='baseline',
                clip_on=True,  # clip to the axes bounding box
                )

    ax.set_xlim(-20, 20)
    ax.set_ylim(-20, 20)


@image_comparison(baseline_images=['fill_units'], extensions=['png'],
                  savefig_kwarg={'dpi': 60})
def test_fill_units():
    from datetime import datetime
    import matplotlib.testing.jpl_units as units
    units.register()

    # generate some data
    t = units.Epoch("ET", dt=datetime(2009, 4, 27))
    value = 10.0 * units.deg
    day = units.Duration("ET", 24.0 * 60.0 * 60.0)

    fig = plt.figure()

    # Top-Left
    ax1 = fig.add_subplot(221)
    ax1.plot([t], [value], yunits='deg', color='red')
    ax1.fill([733525.0, 733525.0, 733526.0, 733526.0],
             [0.0, 0.0, 90.0, 0.0], 'b')

    # Top-Right
    ax2 = fig.add_subplot(222)
    ax2.plot([t], [value], yunits='deg', color='red')
    ax2.fill([t, t, t + day, t + day],
             [0.0, 0.0, 90.0, 0.0], 'b')

    # Bottom-Left
    ax3 = fig.add_subplot(223)
    ax3.plot([t], [value], yunits='deg', color='red')
    ax3.fill([733525.0, 733525.0, 733526.0, 733526.0],
             [0 * units.deg, 0 * units.deg, 90 * units.deg, 0 * units.deg],
             'b')

    # Bottom-Right
    ax4 = fig.add_subplot(224)
    ax4.plot([t], [value], yunits='deg', color='red')
    ax4.fill([t, t, t + day, t + day],
             [0 * units.deg, 0 * units.deg, 90 * units.deg, 0 * units.deg],
             facecolor="blue")

    fig.autofmt_xdate()


@image_comparison(baseline_images=['single_point', 'single_point'])
def test_single_point():
    # Issue #1796: don't let lines.marker affect the grid
    matplotlib.rcParams['lines.marker'] = 'o'
    matplotlib.rcParams['axes.grid'] = True

    fig = plt.figure()
    plt.subplot(211)
    plt.plot([0], [0], 'o')

    plt.subplot(212)
    plt.plot([1], [1], 'o')

    # Reuse testcase from above for a labeled data test
    data = {'a': [0], 'b': [1]}

    fig = plt.figure()
    plt.subplot(211)
    plt.plot('a', 'a', 'o', data=data)

    plt.subplot(212)
    plt.plot('b', 'b', 'o', data=data)


@image_comparison(baseline_images=['single_date'])
def test_single_date():
    time1 = [721964.0]
    data1 = [-65.54]

    fig = plt.figure()
    plt.subplot(211)
    plt.plot_date(time1, data1, 'o', color='r')

    plt.subplot(212)
    plt.plot(time1, data1, 'o', color='r')


@image_comparison(baseline_images=['shaped_data'])
def test_shaped_data():
    xdata = np.array([[0.53295185, 0.23052951, 0.19057629, 0.66724975,
                       0.96577916, 0.73136095, 0.60823287, 0.01792100,
                       0.29744742, 0.27164665],
                      [0.27980120, 0.25814229, 0.02818193, 0.12966456,
                       0.57446277, 0.58167607, 0.71028245, 0.69112737,
                       0.89923072, 0.99072476],
                      [0.81218578, 0.80464528, 0.76071809, 0.85616314,
                       0.12757994, 0.94324936, 0.73078663, 0.09658102,
                       0.60703967, 0.77664978],
                      [0.28332265, 0.81479711, 0.86985333, 0.43797066,
                       0.32540082, 0.43819229, 0.92230363, 0.49414252,
                       0.68168256, 0.05922372],
                      [0.10721335, 0.93904142, 0.79163075, 0.73232848,
                       0.90283839, 0.68408046, 0.25502302, 0.95976614,
                       0.59214115, 0.13663711],
                      [0.28087456, 0.33127607, 0.15530412, 0.76558121,
                       0.83389773, 0.03735974, 0.98717738, 0.71432229,
                       0.54881366, 0.86893953],
                      [0.77995937, 0.99555600, 0.29688434, 0.15646162,
                       0.05184800, 0.37161935, 0.12998491, 0.09377296,
                       0.36882507, 0.36583435],
                      [0.37851836, 0.05315792, 0.63144617, 0.25003433,
                       0.69586032, 0.11393988, 0.92362096, 0.88045438,
                       0.93530252, 0.68275072],
                      [0.86486596, 0.83236675, 0.82960664, 0.57796630,
                       0.25724233, 0.84841095, 0.90862812, 0.64414887,
                       0.35652720, 0.71026066],
                      [0.01383268, 0.34060930, 0.76084285, 0.70800694,
                       0.87634056, 0.08213693, 0.54655021, 0.98123181,
                       0.44080053, 0.86815815]])

    y1 = np.arange(10).reshape((1, -1))
    y2 = np.arange(10).reshape((-1, 1))

    fig = plt.figure()
    plt.subplot(411)
    plt.plot(y1)
    plt.subplot(412)
    plt.plot(y2)

    plt.subplot(413)
    with pytest.raises(ValueError):
        plt.plot((y1, y2))

    plt.subplot(414)
    plt.plot(xdata[:, 1], xdata[1, :], 'o')


@image_comparison(baseline_images=['const_xy'])
def test_const_xy():
    fig = plt.figure()

    plt.subplot(311)
    plt.plot(np.arange(10), np.ones((10,)))

    plt.subplot(312)
    plt.plot(np.ones((10,)), np.arange(10))

    plt.subplot(313)
    plt.plot(np.ones((10,)), np.ones((10,)), 'o')


@image_comparison(baseline_images=['polar_wrap_180', 'polar_wrap_360'],
                  style='default')
def test_polar_wrap():
    D2R = np.pi / 180.0

    fig = plt.figure()

    plt.subplot(111, polar=True)

    plt.polar([179*D2R, -179*D2R], [0.2, 0.1], "b.-")
    plt.polar([179*D2R,  181*D2R], [0.2, 0.1], "g.-")
    plt.rgrids([0.05, 0.1, 0.15, 0.2, 0.25, 0.3])
    assert len(fig.axes) == 1, 'More than one polar axes created.'
    fig = plt.figure()

    plt.subplot(111, polar=True)
    plt.polar([2*D2R, -2*D2R], [0.2, 0.1], "b.-")
    plt.polar([2*D2R,  358*D2R], [0.2, 0.1], "g.-")
    plt.polar([358*D2R,  2*D2R], [0.2, 0.1], "r.-")
    plt.rgrids([0.05, 0.1, 0.15, 0.2, 0.25, 0.3])


@image_comparison(baseline_images=['polar_units', 'polar_units_2'],
                  style='default')
def test_polar_units():
    import matplotlib.testing.jpl_units as units
    units.register()

    pi = np.pi
    deg = units.deg
    km = units.km

    x1 = [pi/6.0, pi/4.0, pi/3.0, pi/2.0]
    x2 = [30.0*deg, 45.0*deg, 60.0*deg, 90.0*deg]

    y1 = [1.0, 2.0, 3.0, 4.0]
    y2 = [4.0, 3.0, 2.0, 1.0]

    fig = plt.figure()

    plt.polar(x2, y1, color="blue")

    # polar(x2, y1, color = "red", xunits="rad")
    # polar(x2, y2, color = "green")

    fig = plt.figure()

    # make sure runits and theta units work
    y1 = [y*km for y in y1]
    plt.polar(x2, y1, color="blue", thetaunits="rad", runits="km")
    assert isinstance(plt.gca().get_xaxis().get_major_formatter(),
                      units.UnitDblFormatter)


@image_comparison(baseline_images=['polar_rmin'], style='default')
def test_polar_rmin():
    r = np.arange(0, 3.0, 0.01)
    theta = 2*np.pi*r

    fig = plt.figure()
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], polar=True)
    ax.plot(theta, r)
    ax.set_rmax(2.0)
    ax.set_rmin(0.5)


@image_comparison(baseline_images=['polar_negative_rmin'], style='default')
def test_polar_negative_rmin():
    r = np.arange(-3.0, 0.0, 0.01)
    theta = 2*np.pi*r

    fig = plt.figure()
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], polar=True)
    ax.plot(theta, r)
    ax.set_rmax(0.0)
    ax.set_rmin(-3.0)


@image_comparison(baseline_images=['polar_rorigin'], style='default')
def test_polar_rorigin():
    r = np.arange(0, 3.0, 0.01)
    theta = 2*np.pi*r

    fig = plt.figure()
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], polar=True)
    ax.plot(theta, r)
    ax.set_rmax(2.0)
    ax.set_rmin(0.5)
    ax.set_rorigin(0.0)


@image_comparison(baseline_images=['polar_theta_position'], style='default')
def test_polar_theta_position():
    r = np.arange(0, 3.0, 0.01)
    theta = 2*np.pi*r

    fig = plt.figure()
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], polar=True)
    ax.plot(theta, r)
    ax.set_theta_zero_location("NW", 30)
    ax.set_theta_direction('clockwise')


@image_comparison(baseline_images=['polar_rlabel_position'], style='default')
def test_polar_rlabel_position():
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='polar')
    ax.set_rlabel_position(315)
    ax.tick_params(rotation='auto')


@image_comparison(baseline_images=['polar_theta_wedge'], style='default',
                  tol=0.01 if six.PY2 else 0)
def test_polar_theta_limits():
    r = np.arange(0, 3.0, 0.01)
    theta = 2*np.pi*r

    theta_mins = np.arange(15.0, 361.0, 90.0)
    theta_maxs = np.arange(50.0, 361.0, 90.0)
    DIRECTIONS = ('out', 'in', 'inout')

    fig, axes = plt.subplots(len(theta_mins), len(theta_maxs),
                             subplot_kw={'polar': True},
                             figsize=(8, 6))

    for i, start in enumerate(theta_mins):
        for j, end in enumerate(theta_maxs):
            ax = axes[i, j]
            ax.plot(theta, r)
            if start < end:
                ax.set_thetamin(start)
                ax.set_thetamax(end)
            else:
                # Plot with clockwise orientation instead.
                ax.set_thetamin(end)
                ax.set_thetamax(start)
                ax.set_theta_direction('clockwise')
            ax.tick_params(tick1On=True, tick2On=True,
                           direction=DIRECTIONS[i % len(DIRECTIONS)],
                           rotation='auto')
            ax.yaxis.set_tick_params(label2On=True, rotation='auto')


@image_comparison(baseline_images=['axvspan_epoch'])
def test_axvspan_epoch():
    from datetime import datetime
    import matplotlib.testing.jpl_units as units
    units.register()

    # generate some data
    t0 = units.Epoch("ET", dt=datetime(2009, 1, 20))
    tf = units.Epoch("ET", dt=datetime(2009, 1, 21))

    dt = units.Duration("ET", units.day.convert("sec"))

    fig = plt.figure()

    plt.axvspan(t0, tf, facecolor="blue", alpha=0.25)

    ax = plt.gca()
    ax.set_xlim(t0 - 5.0*dt, tf + 5.0*dt)


@image_comparison(baseline_images=['axhspan_epoch'])
def test_axhspan_epoch():
    from datetime import datetime
    import matplotlib.testing.jpl_units as units
    units.register()

    # generate some data
    t0 = units.Epoch("ET", dt=datetime(2009, 1, 20))
    tf = units.Epoch("ET", dt=datetime(2009, 1, 21))

    dt = units.Duration("ET", units.day.convert("sec"))

    fig = plt.figure()

    plt.axhspan(t0, tf, facecolor="blue", alpha=0.25)

    ax = plt.gca()
    ax.set_ylim(t0 - 5.0*dt, tf + 5.0*dt)


@image_comparison(baseline_images=['hexbin_extent', 'hexbin_extent'],
                  remove_text=True, extensions=['png'])
def test_hexbin_extent():
    # this test exposes sf bug 2856228
    fig = plt.figure()

    ax = fig.add_subplot(111)
    data = (np.arange(2000) / 2000).reshape((2, 1000))
    x, y = data

    ax.hexbin(x, y, extent=[.1, .3, .6, .7])

    # Reuse testcase from above for a labeled data test
    data = {"x": x, "y": y}

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hexbin("x", "y", extent=[.1, .3, .6, .7], data=data)


@image_comparison(baseline_images=['hexbin_empty'], remove_text=True,
                  extensions=['png'])
def test_hexbin_empty():
    # From #3886: creating hexbin from empty dataset raises ValueError
    ax = plt.gca()
    ax.hexbin([], [])


def test_hexbin_pickable():
    # From #1973: Test that picking a hexbin collection works
    class FauxMouseEvent:
        def __init__(self, x, y):
            self.x = x
            self.y = y

    fig = plt.figure()

    ax = fig.add_subplot(111)
    data = (np.arange(200) / 200).reshape((2, 100))
    x, y = data
    hb = ax.hexbin(x, y, extent=[.1, .3, .6, .7], picker=-1)

    assert hb.contains(FauxMouseEvent(400, 300))[0]


@image_comparison(baseline_images=['hexbin_log'],
                  remove_text=True,
                  extensions=['png'])
def test_hexbin_log():
    # Issue #1636
    fig = plt.figure()

    np.random.seed(0)
    n = 100000
    x = np.random.standard_normal(n)
    y = 2.0 + 3.0 * x + 4.0 * np.random.standard_normal(n)
    y = np.power(2, y * 0.5)
    ax = fig.add_subplot(111)
    ax.hexbin(x, y, yscale='log')


def test_inverted_limits():
    # Test gh:1553
    # Calling invert_xaxis prior to plotting should not disable autoscaling
    # while still maintaining the inverted direction
    fig = plt.figure()
    ax = fig.gca()
    ax.invert_xaxis()
    ax.plot([-5, -3, 2, 4], [1, 2, -3, 5])

    assert ax.get_xlim() == (4, -5)
    assert ax.get_ylim() == (-3, 5)
    plt.close()

    fig = plt.figure()
    ax = fig.gca()
    ax.invert_yaxis()
    ax.plot([-5, -3, 2, 4], [1, 2, -3, 5])

    assert ax.get_xlim() == (-5, 4)
    assert ax.get_ylim() == (5, -3)
    plt.close()


@image_comparison(baseline_images=['nonfinite_limits'])
def test_nonfinite_limits():
    x = np.arange(0., np.e, 0.01)
    # silence divide by zero warning from log(0)
    olderr = np.seterr(divide='ignore')
    try:
        y = np.log(x)
    finally:
        np.seterr(**olderr)
    x[len(x)//2] = np.nan
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(x, y)


@image_comparison(baseline_images=['imshow', 'imshow'],
                  remove_text=True, style='mpl20')
def test_imshow():
    # Create a NxN image
    N = 100
    (x, y) = np.indices((N, N))
    x -= N//2
    y -= N//2
    r = np.sqrt(x**2+y**2-x*y)

    # Create a contour plot at N/4 and extract both the clip path and transform
    fig = plt.figure()
    ax = fig.add_subplot(111)

    ax.imshow(r)

    # Reuse testcase from above for a labeled data test
    data = {"r": r}
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.imshow("r", data=data)


@image_comparison(baseline_images=['imshow_clip'], style='mpl20')
def test_imshow_clip():
    # As originally reported by Gellule Xg <gellule.xg@free.fr>

    # Create a NxN image
    N = 100
    (x, y) = np.indices((N, N))
    x -= N//2
    y -= N//2
    r = np.sqrt(x**2+y**2-x*y)

    # Create a contour plot at N/4 and extract both the clip path and transform
    fig = plt.figure()
    ax = fig.add_subplot(111)

    c = ax.contour(r, [N/4])
    x = c.collections[0]
    clipPath = x.get_paths()[0]
    clipTransform = x.get_transform()

    from matplotlib.transforms import TransformedPath
    clip_path = TransformedPath(clipPath, clipTransform)

    # Plot the image clipped by the contour
    ax.imshow(r, clip_path=clip_path)


@image_comparison(baseline_images=['polycollection_joinstyle'],
                  remove_text=True)
def test_polycollection_joinstyle():
    # Bug #2890979 reported by Matthew West

    from matplotlib import collections as mcoll

    fig = plt.figure()
    ax = fig.add_subplot(111)
    verts = np.array([[1, 1], [1, 2], [2, 2], [2, 1]])
    c = mcoll.PolyCollection([verts], linewidths=40)
    ax.add_collection(c)
    ax.set_xbound(0, 3)
    ax.set_ybound(0, 3)


@pytest.mark.parametrize(
    'x, y1, y2', [
        (np.zeros((2, 2)), 3, 3),
        (np.arange(0.0, 2, 0.02), np.zeros((2, 2)), 3),
        (np.arange(0.0, 2, 0.02), 3, np.zeros((2, 2)))
    ], ids=[
        '2d_x_input',
        '2d_y1_input',
        '2d_y2_input'
    ]
)
def test_fill_between_input(x, y1, y2):
    fig = plt.figure()
    ax = fig.add_subplot(211)
    with pytest.raises(ValueError):
        ax.fill_between(x, y1, y2)


@pytest.mark.parametrize(
    'y, x1, x2', [
        (np.zeros((2, 2)), 3, 3),
        (np.arange(0.0, 2, 0.02), np.zeros((2, 2)), 3),
        (np.arange(0.0, 2, 0.02), 3, np.zeros((2, 2)))
    ], ids=[
        '2d_y_input',
        '2d_x1_input',
        '2d_x2_input'
    ]
)
def test_fill_betweenx_input(y, x1, x2):
    fig = plt.figure()
    ax = fig.add_subplot(211)
    with pytest.raises(ValueError):
        ax.fill_betweenx(y, x1, x2)


@image_comparison(baseline_images=['fill_between_interpolate'],
                  remove_text=True)
def test_fill_between_interpolate():
    x = np.arange(0.0, 2, 0.02)
    y1 = np.sin(2*np.pi*x)
    y2 = 1.2*np.sin(4*np.pi*x)

    fig = plt.figure()
    ax = fig.add_subplot(211)
    ax.plot(x, y1, x, y2, color='black')
    ax.fill_between(x, y1, y2, where=y2 >= y1, facecolor='white', hatch='/',
                    interpolate=True)
    ax.fill_between(x, y1, y2, where=y2 <= y1, facecolor='red',
                    interpolate=True)

    # Test support for masked arrays.
    y2 = np.ma.masked_greater(y2, 1.0)
    # Test that plotting works for masked arrays with the first element masked
    y2[0] = np.ma.masked
    ax1 = fig.add_subplot(212, sharex=ax)
    ax1.plot(x, y1, x, y2, color='black')
    ax1.fill_between(x, y1, y2, where=y2 >= y1, facecolor='green',
                     interpolate=True)
    ax1.fill_between(x, y1, y2, where=y2 <= y1, facecolor='red',
                     interpolate=True)


@image_comparison(baseline_images=['fill_between_interpolate_decreasing'],
                  style='mpl20', remove_text=True)
def test_fill_between_interpolate_decreasing():
    p = np.array([724.3, 700, 655])
    t = np.array([9.4, 7, 2.2])
    prof = np.array([7.9, 6.6, 3.8])

    fig = plt.figure(figsize=(9, 9))
    ax = fig.add_subplot(1, 1, 1)

    ax.plot(t, p, 'tab:red')
    ax.plot(prof, p, 'k')

    ax.fill_betweenx(p, t, prof, where=prof < t,
                     facecolor='blue', interpolate=True, alpha=0.4)
    ax.fill_betweenx(p, t, prof, where=prof > t,
                     facecolor='red', interpolate=True, alpha=0.4)

    ax.set_xlim(0, 30)
    ax.set_ylim(800, 600)


@image_comparison(baseline_images=['symlog'])
def test_symlog():
    x = np.array([0, 1, 2, 4, 6, 9, 12, 24])
    y = np.array([1000000, 500000, 100000, 100, 5, 0, 0, 0])

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(x, y)
    ax.set_yscale('symlog')
    ax.set_xscale('linear')
    ax.set_ylim(-1, 10000000)


@image_comparison(baseline_images=['symlog2'],
                  remove_text=True)
def test_symlog2():
    # Numbers from -50 to 50, with 0.1 as step
    x = np.arange(-50, 50, 0.001)

    fig = plt.figure()
    ax = fig.add_subplot(511)
    # Plots a simple linear function 'f(x) = x'
    ax.plot(x, x)
    ax.set_xscale('symlog', linthreshx=20.0)
    ax.grid(True)

    ax = fig.add_subplot(512)
    # Plots a simple linear function 'f(x) = x'
    ax.plot(x, x)
    ax.set_xscale('symlog', linthreshx=2.0)
    ax.grid(True)

    ax = fig.add_subplot(513)
    # Plots a simple linear function 'f(x) = x'
    ax.plot(x, x)
    ax.set_xscale('symlog', linthreshx=1.0)
    ax.grid(True)

    ax = fig.add_subplot(514)
    # Plots a simple linear function 'f(x) = x'
    ax.plot(x, x)
    ax.set_xscale('symlog', linthreshx=0.1)
    ax.grid(True)

    ax = fig.add_subplot(515)
    # Plots a simple linear function 'f(x) = x'
    ax.plot(x, x)
    ax.set_xscale('symlog', linthreshx=0.01)
    ax.grid(True)
    ax.set_ylim(-0.1, 0.1)


def test_pcolorargs_5205():
    # Smoketest to catch issue found in gh:5205
    x = [-1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5]
    y = [-1.5, -1.25, -1.0, -0.75, -0.5, -0.25, 0,
         0.25, 0.5, 0.75, 1.0, 1.25, 1.5]
    X, Y = np.meshgrid(x, y)
    Z = np.hypot(X, Y)

    plt.pcolor(Z)
    plt.pcolor(list(Z))
    plt.pcolor(x, y, Z)
    plt.pcolor(X, Y, list(Z))


@image_comparison(baseline_images=['pcolormesh'], remove_text=True)
def test_pcolormesh():
    n = 12
    x = np.linspace(-1.5, 1.5, n)
    y = np.linspace(-1.5, 1.5, n*2)
    X, Y = np.meshgrid(x, y)
    Qx = np.cos(Y) - np.cos(X)
    Qz = np.sin(Y) + np.sin(X)
    Qx = (Qx + 1.1)
    Z = np.hypot(X, Y) / 5
    Z = (Z - Z.min()) / Z.ptp()

    # The color array can include masked values:
    Zm = ma.masked_where(np.abs(Qz) < 0.5 * np.max(Qz), Z)

    fig = plt.figure()
    ax = fig.add_subplot(131)
    ax.pcolormesh(Qx, Qz, Z, lw=0.5, edgecolors='k')

    ax = fig.add_subplot(132)
    ax.pcolormesh(Qx, Qz, Z, lw=2, edgecolors=['b', 'w'])

    ax = fig.add_subplot(133)
    ax.pcolormesh(Qx, Qz, Z, shading="gouraud")


@image_comparison(baseline_images=['pcolormesh_datetime_axis'],
                  extensions=['png'], remove_text=False)
def test_pcolormesh_datetime_axis():
    fig = plt.figure()
    fig.subplots_adjust(hspace=0.4, top=0.98, bottom=.15)
    base = datetime.datetime(2013, 1, 1)
    x = np.array([base + datetime.timedelta(days=d) for d in range(21)])
    y = np.arange(21)
    z1, z2 = np.meshgrid(np.arange(20), np.arange(20))
    z = z1 * z2
    plt.subplot(221)
    plt.pcolormesh(x[:-1], y[:-1], z)
    plt.subplot(222)
    plt.pcolormesh(x, y, z)
    x = np.repeat(x[np.newaxis], 21, axis=0)
    y = np.repeat(y[:, np.newaxis], 21, axis=1)
    plt.subplot(223)
    plt.pcolormesh(x[:-1, :-1], y[:-1, :-1], z)
    plt.subplot(224)
    plt.pcolormesh(x, y, z)
    for ax in fig.get_axes():
        for label in ax.get_xticklabels():
            label.set_ha('right')
            label.set_rotation(30)


@image_comparison(baseline_images=['pcolor_datetime_axis'],
                  extensions=['png'], remove_text=False)
def test_pcolor_datetime_axis():
    fig = plt.figure()
    fig.subplots_adjust(hspace=0.4, top=0.98, bottom=.15)
    base = datetime.datetime(2013, 1, 1)
    x = np.array([base + datetime.timedelta(days=d) for d in range(21)])
    y = np.arange(21)
    z1, z2 = np.meshgrid(np.arange(20), np.arange(20))
    z = z1 * z2
    plt.subplot(221)
    plt.pcolor(x[:-1], y[:-1], z)
    plt.subplot(222)
    plt.pcolor(x, y, z)
    x = np.repeat(x[np.newaxis], 21, axis=0)
    y = np.repeat(y[:, np.newaxis], 21, axis=1)
    plt.subplot(223)
    plt.pcolor(x[:-1, :-1], y[:-1, :-1], z)
    plt.subplot(224)
    plt.pcolor(x, y, z)
    for ax in fig.get_axes():
        for label in ax.get_xticklabels():
            label.set_ha('right')
            label.set_rotation(30)


def test_pcolorargs():
    n = 12
    x = np.linspace(-1.5, 1.5, n)
    y = np.linspace(-1.5, 1.5, n*2)
    X, Y = np.meshgrid(x, y)
    Z = np.sqrt(X**2 + Y**2)/5

    _, ax = plt.subplots()
    with pytest.raises(TypeError):
        ax.pcolormesh(y, x, Z)
    with pytest.raises(TypeError):
        ax.pcolormesh(X, Y, Z.T)
    with pytest.raises(TypeError):
        ax.pcolormesh(x, y, Z[:-1, :-1], shading="gouraud")
    with pytest.raises(TypeError):
        ax.pcolormesh(X, Y, Z[:-1, :-1], shading="gouraud")
    x[0] = np.NaN
    with pytest.raises(ValueError):
        ax.pcolormesh(x, y, Z[:-1, :-1])
    x = np.ma.array(x, mask=(x < 0))
    with pytest.raises(ValueError):
        ax.pcolormesh(x, y, Z[:-1, :-1])


@image_comparison(baseline_images=['canonical'])
def test_canonical():
    fig, ax = plt.subplots()
    ax.plot([1, 2, 3])


@image_comparison(baseline_images=['arc_angles'], remove_text=True,
                  style='default', extensions=['png'])
def test_arc_angles():
    from matplotlib import patches
    # Ellipse parameters
    w = 2
    h = 1
    centre = (0.2, 0.5)
    scale = 2

    fig, axs = plt.subplots(3, 3)
    for i, ax in enumerate(axs.flat):
        theta2 = i * 360 / 9
        theta1 = theta2 - 45

        ax.add_patch(patches.Ellipse(centre, w, h, alpha=0.3))
        ax.add_patch(patches.Arc(centre, w, h, theta1=theta1, theta2=theta2))
        # Straight lines intersecting start and end of arc
        ax.plot([scale * np.cos(np.deg2rad(theta1)) + centre[0],
                 centre[0],
                 scale * np.cos(np.deg2rad(theta2)) + centre[0]],
                [scale * np.sin(np.deg2rad(theta1)) + centre[1],
                 centre[1],
                 scale * np.sin(np.deg2rad(theta2)) + centre[1]])

        ax.set_xlim(-scale, scale)
        ax.set_ylim(-scale, scale)

        # This looks the same, but it triggers a different code path when it
        # gets large enough.
        w *= 10
        h *= 10
        centre = (centre[0] * 10, centre[1] * 10)
        scale *= 10


@image_comparison(baseline_images=['arc_ellipse'],
                  remove_text=True)
def test_arc_ellipse():
    from matplotlib import patches
    xcenter, ycenter = 0.38, 0.52
    width, height = 1e-1, 3e-1
    angle = -30

    theta = np.arange(0.0, 360.0, 1.0) * np.pi / 180.0
    x = width / 2. * np.cos(theta)
    y = height / 2. * np.sin(theta)

    rtheta = angle * np.pi / 180.
    R = np.array([
        [np.cos(rtheta), -np.sin(rtheta)],
        [np.sin(rtheta), np.cos(rtheta)]])

    x, y = np.dot(R, np.array([x, y]))
    x += xcenter
    y += ycenter

    fig = plt.figure()
    ax = fig.add_subplot(211, aspect='auto')
    ax.fill(x, y, alpha=0.2, facecolor='yellow', edgecolor='yellow',
            linewidth=1, zorder=1)

    e1 = patches.Arc((xcenter, ycenter), width, height,
                     angle=angle, linewidth=2, fill=False, zorder=2)

    ax.add_patch(e1)

    ax = fig.add_subplot(212, aspect='equal')
    ax.fill(x, y, alpha=0.2, facecolor='green', edgecolor='green', zorder=1)
    e2 = patches.Arc((xcenter, ycenter), width, height,
                     angle=angle, linewidth=2, fill=False, zorder=2)

    ax.add_patch(e2)


@image_comparison(baseline_images=['markevery'],
                  remove_text=True)
def test_markevery():
    x = np.linspace(0, 10, 100)
    y = np.sin(x) * np.sqrt(x/10 + 0.5)

    # check marker only plot
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(x, y, 'o', label='default')
    ax.plot(x, y, 'd', markevery=None, label='mark all')
    ax.plot(x, y, 's', markevery=10, label='mark every 10')
    ax.plot(x, y, '+', markevery=(5, 20), label='mark every 5 starting at 10')
    ax.legend()


@image_comparison(baseline_images=['markevery_line'],
                  remove_text=True)
def test_markevery_line():
    x = np.linspace(0, 10, 100)
    y = np.sin(x) * np.sqrt(x/10 + 0.5)

    # check line/marker combos
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(x, y, '-o', label='default')
    ax.plot(x, y, '-d', markevery=None, label='mark all')
    ax.plot(x, y, '-s', markevery=10, label='mark every 10')
    ax.plot(x, y, '-+', markevery=(5, 20), label='mark every 5 starting at 10')
    ax.legend()


@image_comparison(baseline_images=['markevery_linear_scales'],
                  remove_text=True)
def test_markevery_linear_scales():
    cases = [None,
             8,
             (30, 8),
             [16, 24, 30], [0, -1],
             slice(100, 200, 3),
             0.1, 0.3, 1.5,
             (0.0, 0.1), (0.45, 0.1)]

    cols = 3
    gs = matplotlib.gridspec.GridSpec(len(cases) // cols + 1, cols)

    delta = 0.11
    x = np.linspace(0, 10 - 2 * delta, 200) + delta
    y = np.sin(x) + 1.0 + delta

    for i, case in enumerate(cases):
        row = (i // cols)
        col = i % cols
        plt.subplot(gs[row, col])
        plt.title('markevery=%s' % str(case))
        plt.plot(x, y, 'o', ls='-', ms=4,  markevery=case)


@image_comparison(baseline_images=['markevery_linear_scales_zoomed'],
                  remove_text=True)
def test_markevery_linear_scales_zoomed():
    cases = [None,
             8,
             (30, 8),
             [16, 24, 30], [0, -1],
             slice(100, 200, 3),
             0.1, 0.3, 1.5,
             (0.0, 0.1), (0.45, 0.1)]

    cols = 3
    gs = matplotlib.gridspec.GridSpec(len(cases) // cols + 1, cols)

    delta = 0.11
    x = np.linspace(0, 10 - 2 * delta, 200) + delta
    y = np.sin(x) + 1.0 + delta

    for i, case in enumerate(cases):
        row = (i // cols)
        col = i % cols
        plt.subplot(gs[row, col])
        plt.title('markevery=%s' % str(case))
        plt.plot(x, y, 'o', ls='-', ms=4,  markevery=case)
        plt.xlim((6, 6.7))
        plt.ylim((1.1, 1.7))


@image_comparison(baseline_images=['markevery_log_scales'],
                  remove_text=True)
def test_markevery_log_scales():
    cases = [None,
             8,
             (30, 8),
             [16, 24, 30], [0, -1],
             slice(100, 200, 3),
             0.1, 0.3, 1.5,
             (0.0, 0.1), (0.45, 0.1)]

    cols = 3
    gs = matplotlib.gridspec.GridSpec(len(cases) // cols + 1, cols)

    delta = 0.11
    x = np.linspace(0, 10 - 2 * delta, 200) + delta
    y = np.sin(x) + 1.0 + delta

    for i, case in enumerate(cases):
        row = (i // cols)
        col = i % cols
        plt.subplot(gs[row, col])
        plt.title('markevery=%s' % str(case))
        plt.xscale('log')
        plt.yscale('log')
        plt.plot(x, y, 'o', ls='-', ms=4,  markevery=case)


@image_comparison(baseline_images=['markevery_polar'], style='default',
                  remove_text=True)
def test_markevery_polar():
    cases = [None,
             8,
             (30, 8),
             [16, 24, 30], [0, -1],
             slice(100, 200, 3),
             0.1, 0.3, 1.5,
             (0.0, 0.1), (0.45, 0.1)]

    cols = 3
    gs = matplotlib.gridspec.GridSpec(len(cases) // cols + 1, cols)

    r = np.linspace(0, 3.0, 200)
    theta = 2 * np.pi * r

    for i, case in enumerate(cases):
        row = (i // cols)
        col = i % cols
        plt.subplot(gs[row, col], polar=True)
        plt.title('markevery=%s' % str(case))
        plt.plot(theta, r, 'o', ls='-', ms=4,  markevery=case)


@image_comparison(baseline_images=['marker_edges'],
                  remove_text=True)
def test_marker_edges():
    x = np.linspace(0, 1, 10)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(x, np.sin(x), 'y.', ms=30.0, mew=0, mec='r')
    ax.plot(x+0.1, np.sin(x), 'y.', ms=30.0, mew=1, mec='r')
    ax.plot(x+0.2, np.sin(x), 'y.', ms=30.0, mew=2, mec='b')


@image_comparison(baseline_images=['bar_tick_label_single',
                                   'bar_tick_label_single'],
                  extensions=['png'])
def test_bar_tick_label_single():
    # From 2516: plot bar with array of string labels for x axis
    ax = plt.gca()
    ax.bar(0, 1, tick_label='0')

    # Reuse testcase from above for a labeled data test
    data = {"a": 0, "b": 1}
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax = plt.gca()
    ax.bar("a", "b", tick_label='0', data=data)


def test_bar_ticklabel_fail():
    fig, ax = plt.subplots()
    ax.bar([], [])


@image_comparison(baseline_images=['bar_tick_label_multiple'],
                  extensions=['png'])
def test_bar_tick_label_multiple():
    # From 2516: plot bar with array of string labels for x axis
    ax = plt.gca()
    ax.bar([1, 2.5], [1, 2], width=[0.2, 0.5], tick_label=['a', 'b'],
           align='center')


@image_comparison(
    baseline_images=['bar_tick_label_multiple_old_label_alignment'],
    extensions=['png'])
def test_bar_tick_label_multiple_old_alignment():
    # Test that the algnment for class is backward compatible
    matplotlib.rcParams["ytick.alignment"] = "center"
    ax = plt.gca()
    ax.bar([1, 2.5], [1, 2], width=[0.2, 0.5], tick_label=['a', 'b'],
           align='center')


@image_comparison(baseline_images=['barh_tick_label'],
                  extensions=['png'])
def test_barh_tick_label():
    # From 2516: plot barh with array of string labels for y axis
    ax = plt.gca()
    ax.barh([1, 2.5], [1, 2], height=[0.2, 0.5], tick_label=['a', 'b'],
            align='center')


@image_comparison(baseline_images=['hist_log'],
                  remove_text=True)
def test_hist_log():
    data0 = np.linspace(0, 1, 200)**3
    data = np.r_[1-data0, 1+data0]
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist(data, fill=False, log=True)


@image_comparison(baseline_images=['hist_bar_empty'], remove_text=True,
                  extensions=['png'])
def test_hist_bar_empty():
    # From #3886: creating hist from empty dataset raises ValueError
    ax = plt.gca()
    ax.hist([], histtype='bar')


@image_comparison(baseline_images=['hist_step_empty'], remove_text=True,
                  extensions=['png'])
def test_hist_step_empty():
    # From #3886: creating hist from empty dataset raises ValueError
    ax = plt.gca()
    ax.hist([], histtype='step')


@image_comparison(baseline_images=['hist_steplog'], remove_text=True, tol=0.1)
def test_hist_steplog():
    np.random.seed(0)
    data = np.random.standard_normal(2000)
    data += -2.0 - np.min(data)
    data_pos = data + 2.1
    data_big = data_pos + 30
    weights = np.ones_like(data) * 1.e-5

    ax = plt.subplot(4, 1, 1)
    plt.hist(data, 100, histtype='stepfilled', log=True)

    ax = plt.subplot(4, 1, 2)
    plt.hist(data_pos, 100, histtype='stepfilled', log=True)

    ax = plt.subplot(4, 1, 3)
    plt.hist(data, 100, weights=weights, histtype='stepfilled', log=True)

    ax = plt.subplot(4, 1, 4)
    plt.hist(data_big, 100, histtype='stepfilled', log=True,
             orientation='horizontal')


@image_comparison(baseline_images=['hist_step_filled'], remove_text=True,
                  extensions=['png'])
def test_hist_step_filled():
    np.random.seed(0)
    x = np.random.randn(1000, 3)
    n_bins = 10

    kwargs = [{'fill': True}, {'fill': False}, {'fill': None}, {}]*2
    types = ['step']*4+['stepfilled']*4
    fig, axes = plt.subplots(nrows=2, ncols=4)
    axes = axes.flatten()

    for kg, _type, ax in zip(kwargs, types, axes):
        ax.hist(x, n_bins, histtype=_type, stacked=True, **kg)
        ax.set_title('%s/%s' % (kg, _type))
        ax.set_ylim(ymin=-50)

    patches = axes[0].patches
    assert all([p.get_facecolor() == p.get_edgecolor() for p in patches])


@image_comparison(baseline_images=['hist_density'], extensions=['png'])
def test_hist_density():
    np.random.seed(19680801)
    data = np.random.standard_normal(2000)
    fig, ax = plt.subplots()
    ax.hist(data, density=True)


@image_comparison(baseline_images=['hist_step_log_bottom'],
                  remove_text=True, extensions=['png'])
def test_hist_step_log_bottom():
    # check that bottom doesn't get overwritten by the 'minimum' on a
    # log scale histogram (https://github.com/matplotlib/matplotlib/pull/4608)
    np.random.seed(0)
    data = np.random.standard_normal(2000)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    # normal hist (should clip minimum to 1/base)
    ax.hist(data, bins=10, log=True, histtype='stepfilled',
            alpha=0.5, color='b')
    # manual bottom < 1/base (previously buggy, see #4608)
    ax.hist(data, bins=10, log=True, histtype='stepfilled',
            alpha=0.5, color='g', bottom=1e-2)
    # manual bottom > 1/base
    ax.hist(data, bins=10, log=True, histtype='stepfilled',
            alpha=0.5, color='r', bottom=0.5)
    # array bottom with some less than 1/base (should clip to 1/base)
    ax.hist(data, bins=10, log=True, histtype='stepfilled',
            alpha=0.5, color='y', bottom=np.arange(10))
    ax.set_ylim(9e-3, 1e3)


def test_hist_unequal_bins_density():
    # Test correct behavior of normalized histogram with unequal bins
    # https://github.com/matplotlib/matplotlib/issues/9557
    rng = np.random.RandomState(57483)
    t = rng.randn(100)
    bins = [-3, -1, -0.5, 0, 1, 5]
    mpl_heights, _, _ = plt.hist(t, bins=bins, density=True)
    np_heights, _ = np.histogram(t, bins=bins, density=True)
    assert_allclose(mpl_heights, np_heights)


def contour_dat():
    x = np.linspace(-3, 5, 150)
    y = np.linspace(-3, 5, 120)
    z = np.cos(x) + np.sin(y[:, np.newaxis])
    return x, y, z


@image_comparison(baseline_images=['contour_hatching'])
def test_contour_hatching():
    x, y, z = contour_dat()

    fig = plt.figure()
    ax = fig.add_subplot(111)
    cs = ax.contourf(x, y, z, hatches=['-', '/', '\\', '//'],
                     cmap=plt.get_cmap('gray'),
                     extend='both', alpha=0.5)


@image_comparison(baseline_images=['contour_colorbar'])
def test_contour_colorbar():
    x, y, z = contour_dat()

    fig = plt.figure()
    ax = fig.add_subplot(111)
    cs = ax.contourf(x, y, z, levels=np.arange(-1.8, 1.801, 0.2),
                     cmap=plt.get_cmap('RdBu'),
                     vmin=-0.6,
                     vmax=0.6,
                     extend='both')
    cs1 = ax.contour(x, y, z, levels=np.arange(-2.2, -0.599, 0.2),
                     colors=['y'],
                     linestyles='solid',
                     linewidths=2)
    cs2 = ax.contour(x, y, z, levels=np.arange(0.6, 2.2, 0.2),
                     colors=['c'],
                     linewidths=2)
    cbar = fig.colorbar(cs, ax=ax)
    cbar.add_lines(cs1)
    cbar.add_lines(cs2, erase=False)


@image_comparison(baseline_images=['hist2d', 'hist2d'])
def test_hist2d():
    np.random.seed(0)
    # make it not symmetric in case we switch x and y axis
    x = np.random.randn(100)*2+5
    y = np.random.randn(100)-2
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist2d(x, y, bins=10)

    # Reuse testcase from above for a labeled data test
    data = {"x": x, "y": y}
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist2d("x", "y", bins=10, data=data)


@image_comparison(baseline_images=['hist2d_transpose'])
def test_hist2d_transpose():
    np.random.seed(0)
    # make sure the output from np.histogram is transposed before
    # passing to pcolorfast
    x = np.array([5]*100)
    y = np.random.randn(100)-2
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist2d(x, y, bins=10)


@image_comparison(baseline_images=['scatter', 'scatter'])
def test_scatter_plot():
    fig, ax = plt.subplots()
    data = {"x": [3, 4, 2, 6], "y": [2, 5, 2, 3], "c": ['r', 'y', 'b', 'lime'],
            "s": [24, 15, 19, 29]}

    ax.scatter(data["x"], data["y"], c=data["c"], s=data["s"])

    # Reuse testcase from above for a labeled data test
    fig, ax = plt.subplots()
    ax.scatter("x", "y", c="c", s="s", data=data)


@image_comparison(baseline_images=['scatter_marker'], remove_text=True,
                  extensions=['png'])
def test_scatter_marker():
    fig, (ax0, ax1, ax2) = plt.subplots(ncols=3)
    ax0.scatter([3, 4, 2, 6], [2, 5, 2, 3],
                c=[(1, 0, 0), 'y', 'b', 'lime'],
                s=[60, 50, 40, 30],
                edgecolors=['k', 'r', 'g', 'b'],
                marker='s')
    ax1.scatter([3, 4, 2, 6], [2, 5, 2, 3],
                c=[(1, 0, 0), 'y', 'b', 'lime'],
                s=[60, 50, 40, 30],
                edgecolors=['k', 'r', 'g', 'b'],
                marker=mmarkers.MarkerStyle('o', fillstyle='top'))
    # unit area ellipse
    rx, ry = 3, 1
    area = rx * ry * np.pi
    theta = np.linspace(0, 2 * np.pi, 21)
    verts = list(zip(np.cos(theta) * rx / area, np.sin(theta) * ry / area))
    ax2.scatter([3, 4, 2, 6], [2, 5, 2, 3],
                c=[(1, 0, 0), 'y', 'b', 'lime'],
                s=[60, 50, 40, 30],
                edgecolors=['k', 'r', 'g', 'b'],
                verts=verts)


@image_comparison(baseline_images=['scatter_2D'], remove_text=True,
                  extensions=['png'])
def test_scatter_2D():
    x = np.arange(3)
    y = np.arange(2)
    x, y = np.meshgrid(x, y)
    z = x + y
    fig, ax = plt.subplots()
    ax.scatter(x, y, c=z, s=200, edgecolors='face')


def test_scatter_color():
    # Try to catch cases where 'c' kwarg should have been used.
    with pytest.raises(ValueError):
        plt.scatter([1, 2], [1, 2], color=[0.1, 0.2])
    with pytest.raises(ValueError):
        plt.scatter([1, 2, 3], [1, 2, 3], color=[1, 2, 3])


def test_as_mpl_axes_api():
    # tests the _as_mpl_axes api
    from matplotlib.projections.polar import PolarAxes
    import matplotlib.axes as maxes

    class Polar(object):
        def __init__(self):
            self.theta_offset = 0

        def _as_mpl_axes(self):
            # implement the matplotlib axes interface
            return PolarAxes, {'theta_offset': self.theta_offset}
    prj = Polar()
    prj2 = Polar()
    prj2.theta_offset = np.pi
    prj3 = Polar()

    # testing axes creation with plt.axes
    ax = plt.axes([0, 0, 1, 1], projection=prj)
    assert type(ax) == PolarAxes, \
        'Expected a PolarAxes, got %s' % type(ax)
    ax_via_gca = plt.gca(projection=prj)
    assert ax_via_gca is ax
    plt.close()

    # testing axes creation with gca
    ax = plt.gca(projection=prj)
    assert type(ax) == maxes._subplots._subplot_classes[PolarAxes], \
        'Expected a PolarAxesSubplot, got %s' % type(ax)
    ax_via_gca = plt.gca(projection=prj)
    assert ax_via_gca is ax
    # try getting the axes given a different polar projection
    with pytest.warns(UserWarning) as rec:
        ax_via_gca = plt.gca(projection=prj2)
        assert len(rec) == 1
        assert 'Requested projection is different' in str(rec[0].message)
    assert ax_via_gca is not ax
    assert ax.get_theta_offset() == 0, ax.get_theta_offset()
    assert ax_via_gca.get_theta_offset() == np.pi, \
        ax_via_gca.get_theta_offset()
    # try getting the axes given an == (not is) polar projection
    with pytest.warns(UserWarning):
        ax_via_gca = plt.gca(projection=prj3)
        assert len(rec) == 1
        assert 'Requested projection is different' in str(rec[0].message)
    assert ax_via_gca is ax
    plt.close()

    # testing axes creation with subplot
    ax = plt.subplot(121, projection=prj)
    assert type(ax) == maxes._subplots._subplot_classes[PolarAxes], \
        'Expected a PolarAxesSubplot, got %s' % type(ax)
    plt.close()


def test_pyplot_axes():
    # test focusing of Axes in other Figure
    fig1, ax1 = plt.subplots()
    fig2, ax2 = plt.subplots()
    assert ax1 is plt.axes(ax1)
    assert ax1 is plt.gca()
    assert fig1 is plt.gcf()
    plt.close(fig1)
    plt.close(fig2)


@image_comparison(baseline_images=['log_scales'])
def test_log_scales():
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.plot(np.log(np.linspace(0.1, 100)))
    ax.set_yscale('log', basey=5.5)
    ax.invert_yaxis()
    ax.set_xscale('log', basex=9.0)


@image_comparison(baseline_images=['stackplot_test_image',
                                   'stackplot_test_image'])
def test_stackplot():
    fig = plt.figure()
    x = np.linspace(0, 10, 10)
    y1 = 1.0 * x
    y2 = 2.0 * x + 1
    y3 = 3.0 * x + 2
    ax = fig.add_subplot(1, 1, 1)
    ax.stackplot(x, y1, y2, y3)
    ax.set_xlim((0, 10))
    ax.set_ylim((0, 70))

    # Reuse testcase from above for a labeled data test
    data = {"x": x, "y1": y1, "y2": y2, "y3": y3}
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.stackplot("x", "y1", "y2", "y3", data=data)
    ax.set_xlim((0, 10))
    ax.set_ylim((0, 70))


@image_comparison(baseline_images=['stackplot_test_baseline'],
                  remove_text=True)
def test_stackplot_baseline():
    np.random.seed(0)

    def layers(n, m):
        def bump(a):
            x = 1 / (.1 + np.random.random())
            y = 2 * np.random.random() - .5
            z = 10 / (.1 + np.random.random())
            a += x * np.exp(-((np.arange(m) / m - y) * z) ** 2)
        a = np.zeros((m, n))
        for i in range(n):
            for j in range(5):
                bump(a[:, i])
        return a

    d = layers(3, 100)
    d[50, :] = 0  # test for fixed weighted wiggle (issue #6313)

    fig = plt.figure()

    plt.subplot(2, 2, 1)
    plt.stackplot(list(xrange(100)), d.T, baseline='zero')

    plt.subplot(2, 2, 2)
    plt.stackplot(list(xrange(100)), d.T, baseline='sym')

    plt.subplot(2, 2, 3)
    plt.stackplot(list(xrange(100)), d.T, baseline='wiggle')

    plt.subplot(2, 2, 4)
    plt.stackplot(list(xrange(100)), d.T, baseline='weighted_wiggle')


@image_comparison(baseline_images=['bxp_baseline'],
                  extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_baseline():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats)


@image_comparison(baseline_images=['bxp_rangewhis'],
                  extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_rangewhis():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4)),
        whis='range'
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats)


@image_comparison(baseline_images=['bxp_precentilewhis'],
                  extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_precentilewhis():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4)),
        whis=[5, 95]
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats)


@image_comparison(baseline_images=['bxp_with_xlabels'],
                  extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_with_xlabels():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )
    for stats, label in zip(logstats, list('ABCD')):
        stats['label'] = label

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats)


@image_comparison(baseline_images=['bxp_horizontal'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default',
                  tol=0.1)
def test_bxp_horizontal():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_xscale('log')
    ax.bxp(logstats, vert=False)


@image_comparison(baseline_images=['bxp_with_ylabels'],
                  extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default',
                  tol=0.1,)
def test_bxp_with_ylabels():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )
    for stats, label in zip(logstats, list('ABCD')):
        stats['label'] = label

    fig, ax = plt.subplots()
    ax.set_xscale('log')
    ax.bxp(logstats, vert=False)


@image_comparison(baseline_images=['bxp_patchartist'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_patchartist():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, patch_artist=True)


@image_comparison(baseline_images=['bxp_custompatchartist'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 100},
                  style='default')
def test_bxp_custompatchartist():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    boxprops = dict(facecolor='yellow', edgecolor='green', linestyle='dotted')
    ax.bxp(logstats, patch_artist=True, boxprops=boxprops)


@image_comparison(baseline_images=['bxp_customoutlier'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_customoutlier():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    flierprops = dict(linestyle='none', marker='d', markerfacecolor='g')
    ax.bxp(logstats, flierprops=flierprops)


@image_comparison(baseline_images=['bxp_withmean_custompoint'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_showcustommean():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    meanprops = dict(linestyle='none', marker='d', markerfacecolor='green')
    ax.bxp(logstats, showmeans=True, meanprops=meanprops)


@image_comparison(baseline_images=['bxp_custombox'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_custombox():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    boxprops = dict(linestyle='--', color='b', linewidth=3)
    ax.bxp(logstats, boxprops=boxprops)


@image_comparison(baseline_images=['bxp_custommedian'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_custommedian():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    medianprops = dict(linestyle='--', color='b', linewidth=3)
    ax.bxp(logstats, medianprops=medianprops)


@image_comparison(baseline_images=['bxp_customcap'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_customcap():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    capprops = dict(linestyle='--', color='g', linewidth=3)
    ax.bxp(logstats, capprops=capprops)


@image_comparison(baseline_images=['bxp_customwhisker'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_customwhisker():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    whiskerprops = dict(linestyle='-', color='m', linewidth=3)
    ax.bxp(logstats, whiskerprops=whiskerprops)


@image_comparison(baseline_images=['bxp_withnotch'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_shownotches():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, shownotches=True)


@image_comparison(baseline_images=['bxp_nocaps'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_nocaps():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, showcaps=False)


@image_comparison(baseline_images=['bxp_nobox'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_nobox():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, showbox=False)


@image_comparison(baseline_images=['bxp_no_flier_stats'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_no_flier_stats():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )
    for ls in logstats:
        ls.pop('fliers', None)

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, showfliers=False)


@image_comparison(baseline_images=['bxp_withmean_point'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_showmean():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, showmeans=True, meanline=False)


@image_comparison(baseline_images=['bxp_withmean_line'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_showmeanasline():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, showmeans=True, meanline=True)


@image_comparison(baseline_images=['bxp_scalarwidth'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_scalarwidth():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, widths=0.25)


@image_comparison(baseline_images=['bxp_customwidths'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_customwidths():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, widths=[0.10, 0.25, 0.65, 0.85])


@image_comparison(baseline_images=['bxp_custompositions'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_bxp_custompositions():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    ax.bxp(logstats, positions=[1, 5, 6, 7])


def test_bxp_bad_widths():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    with pytest.raises(ValueError):
        ax.bxp(logstats, widths=[1])


def test_bxp_bad_positions():
    np.random.seed(937)
    logstats = matplotlib.cbook.boxplot_stats(
        np.random.lognormal(mean=1.25, sigma=1., size=(37, 4))
    )

    fig, ax = plt.subplots()
    ax.set_yscale('log')
    with pytest.raises(ValueError):
        ax.bxp(logstats, positions=[2, 3])


@image_comparison(baseline_images=['boxplot', 'boxplot'],
                  tol=1,
                  style='default')
def test_boxplot():
    # Randomness used for bootstrapping.
    np.random.seed(937)

    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    fig, ax = plt.subplots()

    ax.boxplot([x, x], bootstrap=10000, notch=1)
    ax.set_ylim((-30, 30))

    # Reuse testcase from above for a labeled data test
    data = {"x": [x, x]}
    fig, ax = plt.subplots()
    ax.boxplot("x", bootstrap=10000, notch=1, data=data)
    ax.set_ylim((-30, 30))


@image_comparison(baseline_images=['boxplot_sym2'],
                  remove_text=True, extensions=['png'],
                  style='default')
def test_boxplot_sym2():
    # Randomness used for bootstrapping.
    np.random.seed(937)

    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    fig, [ax1, ax2] = plt.subplots(1, 2)

    ax1.boxplot([x, x], bootstrap=10000, sym='^')
    ax1.set_ylim((-30, 30))

    ax2.boxplot([x, x], bootstrap=10000, sym='g')
    ax2.set_ylim((-30, 30))


@image_comparison(baseline_images=['boxplot_sym'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40},
                  style='default')
def test_boxplot_sym():
    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    fig, ax = plt.subplots()

    ax.boxplot([x, x], sym='gs')
    ax.set_ylim((-30, 30))


@image_comparison(
    baseline_images=['boxplot_autorange_false_whiskers',
                     'boxplot_autorange_true_whiskers'],
    extensions=['png'],
    style='default'
)
def test_boxplot_autorange_whiskers():
    # Randomness used for bootstrapping.
    np.random.seed(937)

    x = np.ones(140)
    x = np.hstack([0, x, 2])

    fig1, ax1 = plt.subplots()
    ax1.boxplot([x, x], bootstrap=10000, notch=1)
    ax1.set_ylim((-5, 5))

    fig2, ax2 = plt.subplots()
    ax2.boxplot([x, x], bootstrap=10000, notch=1, autorange=True)
    ax2.set_ylim((-5, 5))


def _rc_test_bxp_helper(ax, rc_dict):
    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    with matplotlib.rc_context(rc_dict):
        ax.boxplot([x, x])
    return ax


@image_comparison(baseline_images=['boxplot_rc_parameters'],
                  savefig_kwarg={'dpi': 100}, remove_text=True,
                  tol=1, style='default')
def test_boxplot_rc_parameters():
    # Randomness used for bootstrapping.
    np.random.seed(937)

    fig, ax = plt.subplots(3)

    rc_axis0 = {
        'boxplot.notch': True,
        'boxplot.whiskers': [5, 95],
        'boxplot.bootstrap': 10000,

        'boxplot.flierprops.color': 'b',
        'boxplot.flierprops.marker': 'o',
        'boxplot.flierprops.markerfacecolor': 'g',
        'boxplot.flierprops.markeredgecolor': 'b',
        'boxplot.flierprops.markersize': 5,
        'boxplot.flierprops.linestyle': '--',
        'boxplot.flierprops.linewidth': 2.0,

        'boxplot.boxprops.color': 'r',
        'boxplot.boxprops.linewidth': 2.0,
        'boxplot.boxprops.linestyle': '--',

        'boxplot.capprops.color': 'c',
        'boxplot.capprops.linewidth': 2.0,
        'boxplot.capprops.linestyle': '--',

        'boxplot.medianprops.color': 'k',
        'boxplot.medianprops.linewidth': 2.0,
        'boxplot.medianprops.linestyle': '--',
    }

    rc_axis1 = {
        'boxplot.vertical': False,
        'boxplot.whiskers': 'range',
        'boxplot.patchartist': True,
    }

    rc_axis2 = {
        'boxplot.whiskers': 2.0,
        'boxplot.showcaps': False,
        'boxplot.showbox': False,
        'boxplot.showfliers': False,
        'boxplot.showmeans': True,
        'boxplot.meanline': True,

        'boxplot.meanprops.color': 'c',
        'boxplot.meanprops.linewidth': 2.0,
        'boxplot.meanprops.linestyle': '--',

        'boxplot.whiskerprops.color': 'r',
        'boxplot.whiskerprops.linewidth': 2.0,
        'boxplot.whiskerprops.linestyle': '-.',
    }
    dict_list = [rc_axis0, rc_axis1, rc_axis2]
    for axis, rc_axis in zip(ax, dict_list):
        _rc_test_bxp_helper(axis, rc_axis)

    assert (matplotlib.patches.PathPatch in
            [type(t) for t in ax[1].get_children()])


@image_comparison(baseline_images=['boxplot_with_CIarray'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40}, style='default')
def test_boxplot_with_CIarray():
    # Randomness used for bootstrapping.
    np.random.seed(937)

    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    fig = plt.figure()
    ax = fig.add_subplot(111)
    CIs = np.array([[-1.5, 3.], [-1., 3.5]])

    # show 1 boxplot with mpl medians/conf. interfals, 1 with manual values
    ax.boxplot([x, x], bootstrap=10000, usermedians=[None, 1.0],
               conf_intervals=CIs, notch=1)
    ax.set_ylim((-30, 30))


@image_comparison(baseline_images=['boxplot_no_inverted_whisker'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40}, style='default')
def test_boxplot_no_weird_whisker():
    x = np.array([3, 9000, 150, 88, 350, 200000, 1400, 960],
                 dtype=np.float64)
    ax1 = plt.axes()
    ax1.boxplot(x)
    ax1.set_yscale('log')
    ax1.yaxis.grid(False, which='minor')
    ax1.xaxis.grid(False)


def test_boxplot_bad_medians_1():
    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    fig, ax = plt.subplots()
    with pytest.raises(ValueError):
        ax.boxplot(x, usermedians=[1, 2])


def test_boxplot_bad_medians_2():
    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    fig, ax = plt.subplots()
    with pytest.raises(ValueError):
        ax.boxplot([x, x], usermedians=[[1, 2], [1, 2]])


def test_boxplot_bad_ci_1():
    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    fig, ax = plt.subplots()
    with pytest.raises(ValueError):
        ax.boxplot([x, x], conf_intervals=[[1, 2]])


def test_boxplot_zorder():
    x = np.arange(10)
    fix, ax = plt.subplots()
    assert ax.boxplot(x)['boxes'][0].get_zorder() == 2
    assert ax.boxplot(x, zorder=10)['boxes'][0].get_zorder() == 10


def test_boxplot_bad_ci_2():
    x = np.linspace(-7, 7, 140)
    x = np.hstack([-25, x, 25])
    fig, ax = plt.subplots()
    with pytest.raises(ValueError):
        ax.boxplot([x, x], conf_intervals=[[1, 2], [1]])


@image_comparison(baseline_images=['boxplot_mod_artists_after_plotting'],
                  remove_text=True, extensions=['png'],
                  savefig_kwarg={'dpi': 40}, style='default')
def test_boxplot_mod_artist_after_plotting():
    x = [0.15, 0.11, 0.06, 0.06, 0.12, 0.56, -0.56]
    fig, ax = plt.subplots()
    bp = ax.boxplot(x, sym="o")
    for key in bp:
        for obj in bp[key]:
            obj.set_color('green')


@image_comparison(baseline_images=['violinplot_vert_baseline',
                                   'violinplot_vert_baseline'],
                  extensions=['png'])
def test_vert_violinplot_baseline():
    # First 9 digits of frac(sqrt(2))
    np.random.seed(414213562)
    data = [np.random.normal(size=100) for i in range(4)]
    ax = plt.axes()
    ax.violinplot(data, positions=range(4), showmeans=0, showextrema=0,
                  showmedians=0)

    # Reuse testcase from above for a labeled data test
    data = {"d": data}
    fig, ax = plt.subplots()
    ax = plt.axes()
    ax.violinplot("d", positions=range(4), showmeans=0, showextrema=0,
                  showmedians=0, data=data)


@image_comparison(baseline_images=['violinplot_vert_showmeans'],
                  extensions=['png'])
def test_vert_violinplot_showmeans():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(3))
    np.random.seed(732050807)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), showmeans=1, showextrema=0,
                  showmedians=0)


@image_comparison(baseline_images=['violinplot_vert_showextrema'],
                  extensions=['png'])
def test_vert_violinplot_showextrema():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(5))
    np.random.seed(236067977)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), showmeans=0, showextrema=1,
                  showmedians=0)


@image_comparison(baseline_images=['violinplot_vert_showmedians'],
                  extensions=['png'])
def test_vert_violinplot_showmedians():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(7))
    np.random.seed(645751311)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), showmeans=0, showextrema=0,
                  showmedians=1)


@image_comparison(baseline_images=['violinplot_vert_showall'],
                  extensions=['png'])
def test_vert_violinplot_showall():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(11))
    np.random.seed(316624790)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), showmeans=1, showextrema=1,
                  showmedians=1)


@image_comparison(baseline_images=['violinplot_vert_custompoints_10'],
                  extensions=['png'])
def test_vert_violinplot_custompoints_10():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(13))
    np.random.seed(605551275)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), showmeans=0, showextrema=0,
                  showmedians=0, points=10)


@image_comparison(baseline_images=['violinplot_vert_custompoints_200'],
                  extensions=['png'])
def test_vert_violinplot_custompoints_200():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(17))
    np.random.seed(123105625)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), showmeans=0, showextrema=0,
                  showmedians=0, points=200)


@image_comparison(baseline_images=['violinplot_horiz_baseline'],
                  extensions=['png'])
def test_horiz_violinplot_baseline():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(19))
    np.random.seed(358898943)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), vert=False, showmeans=0,
                  showextrema=0, showmedians=0)


@image_comparison(baseline_images=['violinplot_horiz_showmedians'],
                  extensions=['png'])
def test_horiz_violinplot_showmedians():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(23))
    np.random.seed(795831523)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), vert=False, showmeans=0,
                  showextrema=0, showmedians=1)


@image_comparison(baseline_images=['violinplot_horiz_showmeans'],
                  extensions=['png'])
def test_horiz_violinplot_showmeans():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(29))
    np.random.seed(385164807)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), vert=False, showmeans=1,
                  showextrema=0, showmedians=0)


@image_comparison(baseline_images=['violinplot_horiz_showextrema'],
                  extensions=['png'])
def test_horiz_violinplot_showextrema():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(31))
    np.random.seed(567764362)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), vert=False, showmeans=0,
                  showextrema=1, showmedians=0)


@image_comparison(baseline_images=['violinplot_horiz_showall'],
                  extensions=['png'])
def test_horiz_violinplot_showall():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(37))
    np.random.seed(82762530)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), vert=False, showmeans=1,
                  showextrema=1, showmedians=1)


@image_comparison(baseline_images=['violinplot_horiz_custompoints_10'],
                  extensions=['png'])
def test_horiz_violinplot_custompoints_10():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(41))
    np.random.seed(403124237)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), vert=False, showmeans=0,
                  showextrema=0, showmedians=0, points=10)


@image_comparison(baseline_images=['violinplot_horiz_custompoints_200'],
                  extensions=['png'])
def test_horiz_violinplot_custompoints_200():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(43))
    np.random.seed(557438524)
    data = [np.random.normal(size=100) for i in range(4)]
    ax.violinplot(data, positions=range(4), vert=False, showmeans=0,
                  showextrema=0, showmedians=0, points=200)


def test_violinplot_bad_positions():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(47))
    np.random.seed(855654600)
    data = [np.random.normal(size=100) for i in range(4)]
    with pytest.raises(ValueError):
        ax.violinplot(data, positions=range(5))


def test_violinplot_bad_widths():
    ax = plt.axes()
    # First 9 digits of frac(sqrt(53))
    np.random.seed(280109889)
    data = [np.random.normal(size=100) for i in range(4)]
    with pytest.raises(ValueError):
        ax.violinplot(data, positions=range(4), widths=[1, 2, 3])


def test_manage_xticks():
    _, ax = plt.subplots()
    ax.set_xlim(0, 4)
    old_xlim = ax.get_xlim()
    np.random.seed(0)
    y1 = np.random.normal(10, 3, 20)
    y2 = np.random.normal(3, 1, 20)
    ax.boxplot([y1, y2], positions=[1, 2],
               manage_xticks=False)
    new_xlim = ax.get_xlim()
    assert_array_equal(old_xlim, new_xlim)


def test_tick_space_size_0():
    # allow font size to be zero, which affects ticks when there is
    # no other text in the figure.
    plt.plot([0, 1], [0, 1])
    matplotlib.rcParams.update({'font.size': 0})
    b = io.BytesIO()
    plt.savefig(b, dpi=80, format='raw')


@image_comparison(baseline_images=['errorbar_basic', 'errorbar_mixed',
                                   'errorbar_basic'])
def test_errorbar():
    x = np.arange(0.1, 4, 0.5)
    y = np.exp(-x)

    yerr = 0.1 + 0.2*np.sqrt(x)
    xerr = 0.1 + yerr

    # First illustrate basic pyplot interface, using defaults where possible.
    fig = plt.figure()
    ax = fig.gca()
    ax.errorbar(x, y, xerr=0.2, yerr=0.4)
    ax.set_title("Simplest errorbars, 0.2 in x, 0.4 in y")

    # Now switch to a more OO interface to exercise more features.
    fig, axs = plt.subplots(nrows=2, ncols=2, sharex=True)
    ax = axs[0, 0]
    # Try a Nx1 shaped error just to check
    ax.errorbar(x, y, yerr=np.reshape(yerr, (len(y), 1)), fmt='o')
    ax.set_title('Vert. symmetric')

    # With 4 subplots, reduce the number of axis ticks to avoid crowding.
    ax.locator_params(nbins=4)

    ax = axs[0, 1]
    ax.errorbar(x, y, xerr=xerr, fmt='o', alpha=0.4)
    ax.set_title('Hor. symmetric w/ alpha')

    ax = axs[1, 0]
    ax.errorbar(x, y, yerr=[yerr, 2*yerr], xerr=[xerr, 2*xerr], fmt='--o')
    ax.set_title('H, V asymmetric')

    ax = axs[1, 1]
    ax.set_yscale('log')
    # Here we have to be careful to keep all y values positive:
    ylower = np.maximum(1e-2, y - yerr)
    yerr_lower = y - ylower

    ax.errorbar(x, y, yerr=[yerr_lower, 2*yerr], xerr=xerr,
                fmt='o', ecolor='g', capthick=2)
    ax.set_title('Mixed sym., log y')

    fig.suptitle('Variable errorbars')

    # Reuse te first testcase from above for a labeled data test
    data = {"x": x, "y": y}
    fig = plt.figure()
    ax = fig.gca()
    ax.errorbar("x", "y", xerr=0.2, yerr=0.4, data=data)
    ax.set_title("Simplest errorbars, 0.2 in x, 0.4 in y")


def test_errorbar_colorcycle():

    f, ax = plt.subplots()
    x = np.arange(10)
    y = 2*x

    e1, _, _ = ax.errorbar(x, y, c=None)
    e2, _, _ = ax.errorbar(x, 2*y, c=None)
    ln1, = ax.plot(x, 4*y)

    assert mcolors.to_rgba(e1.get_color()) == mcolors.to_rgba('C0')
    assert mcolors.to_rgba(e2.get_color()) == mcolors.to_rgba('C1')
    assert mcolors.to_rgba(ln1.get_color()) == mcolors.to_rgba('C2')


def test_errorbar_shape():
    fig = plt.figure()
    ax = fig.gca()

    x = np.arange(0.1, 4, 0.5)
    y = np.exp(-x)
    yerr1 = 0.1 + 0.2*np.sqrt(x)
    yerr = np.vstack((yerr1, 2*yerr1)).T
    xerr = 0.1 + yerr

    with pytest.raises(ValueError):
        ax.errorbar(x, y, yerr=yerr, fmt='o')
    with pytest.raises(ValueError):
        ax.errorbar(x, y, xerr=xerr, fmt='o')
    with pytest.raises(ValueError):
        ax.errorbar(x, y, yerr=yerr, xerr=xerr, fmt='o')


@image_comparison(baseline_images=['errorbar_limits'])
def test_errorbar_limits():
    x = np.arange(0.5, 5.5, 0.5)
    y = np.exp(-x)
    xerr = 0.1
    yerr = 0.2
    ls = 'dotted'

    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)

    # standard error bars
    plt.errorbar(x, y, xerr=xerr, yerr=yerr, ls=ls, color='blue')

    # including upper limits
    uplims = np.zeros_like(x)
    uplims[[1, 5, 9]] = True
    plt.errorbar(x, y+0.5, xerr=xerr, yerr=yerr, uplims=uplims, ls=ls,
                 color='green')

    # including lower limits
    lolims = np.zeros_like(x)
    lolims[[2, 4, 8]] = True
    plt.errorbar(x, y+1.0, xerr=xerr, yerr=yerr, lolims=lolims, ls=ls,
                 color='red')

    # including upper and lower limits
    plt.errorbar(x, y+1.5, marker='o', ms=8, xerr=xerr, yerr=yerr,
                 lolims=lolims, uplims=uplims, ls=ls, color='magenta')

    # including xlower and xupper limits
    xerr = 0.2
    yerr = np.zeros_like(x) + 0.2
    yerr[[3, 6]] = 0.3
    xlolims = lolims
    xuplims = uplims
    lolims = np.zeros_like(x)
    uplims = np.zeros_like(x)
    lolims[[6]] = True
    uplims[[3]] = True
    plt.errorbar(x, y+2.1, marker='o', ms=8, xerr=xerr, yerr=yerr,
                 xlolims=xlolims, xuplims=xuplims, uplims=uplims,
                 lolims=lolims, ls='none', mec='blue', capsize=0,
                 color='cyan')
    ax.set_xlim((0, 5.5))
    ax.set_title('Errorbar upper and lower limits')


def test_errobar_nonefmt():
    # Check that passing 'none' as a format still plots errorbars
    x = np.arange(5)
    y = np.arange(5)

    plotline, _, barlines = plt.errorbar(x, y, xerr=1, yerr=1, fmt='none')
    assert plotline is None
    for errbar in barlines:
        assert np.all(errbar.get_color() == mcolors.to_rgba('C0'))


@image_comparison(baseline_images=['hist_stacked_stepfilled',
                                   'hist_stacked_stepfilled'])
def test_hist_stacked_stepfilled():
    # make some data
    d1 = np.linspace(1, 3, 20)
    d2 = np.linspace(0, 10, 50)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist((d1, d2), histtype="stepfilled", stacked=True)

    # Reuse testcase from above for a labeled data test
    data = {"x": (d1, d2)}
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist("x", histtype="stepfilled", stacked=True, data=data)


@image_comparison(baseline_images=['hist_offset'])
def test_hist_offset():
    # make some data
    d1 = np.linspace(0, 10, 50)
    d2 = np.linspace(1, 3, 20)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist(d1, bottom=5)
    ax.hist(d2, bottom=15)


@image_comparison(baseline_images=['hist_step'], extensions=['png'],
                  remove_text=True)
def test_hist_step():
    # make some data
    d1 = np.linspace(1, 3, 20)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist(d1, histtype="step")
    ax.set_ylim(0, 10)
    ax.set_xlim(-1, 5)


@image_comparison(baseline_images=['hist_step_horiz'], extensions=['png'])
def test_hist_step_horiz():
    # make some data
    d1 = np.linspace(0, 10, 50)
    d2 = np.linspace(1, 3, 20)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist((d1, d2), histtype="step", orientation="horizontal")


@image_comparison(baseline_images=['hist_stacked_weights'])
def test_hist_stacked_weighted():
    # make some data
    d1 = np.linspace(0, 10, 50)
    d2 = np.linspace(1, 3, 20)
    w1 = np.linspace(0.01, 3.5, 50)
    w2 = np.linspace(0.05, 2., 20)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist((d1, d2), weights=(w1, w2), histtype="stepfilled", stacked=True)


def test_stem_args():
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)

    x = list(xrange(10))
    y = list(xrange(10))

    # Test the call signatures
    ax.stem(y)
    ax.stem(x, y)
    ax.stem(x, y, 'r--')
    ax.stem(x, y, 'r--', basefmt='b--')


def test_stem_dates():
    fig, ax = plt.subplots(1, 1)
    from dateutil import parser
    x = parser.parse("2013-9-28 11:00:00")
    y = 100

    x1 = parser.parse("2013-9-28 12:00:00")
    y1 = 200

    ax.stem([x, x1], [y, y1], "*-")


@image_comparison(baseline_images=['hist_stacked_stepfilled_alpha'])
def test_hist_stacked_stepfilled_alpha():
    # make some data
    d1 = np.linspace(1, 3, 20)
    d2 = np.linspace(0, 10, 50)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist((d1, d2), histtype="stepfilled", stacked=True, alpha=0.5)


@image_comparison(baseline_images=['hist_stacked_step'])
def test_hist_stacked_step():
    # make some data
    d1 = np.linspace(1, 3, 20)
    d2 = np.linspace(0, 10, 50)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist((d1, d2), histtype="step", stacked=True)


@image_comparison(baseline_images=['hist_stacked_normed'])
def test_hist_stacked_normed():
    # make some data
    d1 = np.linspace(1, 3, 20)
    d2 = np.linspace(0, 10, 50)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist((d1, d2), stacked=True, normed=True)


@image_comparison(baseline_images=['hist_stacked_normed'], extensions=['png'])
def test_hist_stacked_density():
    # make some data
    d1 = np.linspace(1, 3, 20)
    d2 = np.linspace(0, 10, 50)
    fig, ax = plt.subplots()
    ax.hist((d1, d2), stacked=True, density=True)


@pytest.mark.parametrize('normed', [False, True])
@pytest.mark.parametrize('density', [False, True])
def test_hist_normed_density(normed, density):
    # Normed and density should not be used simultaneously
    d1 = np.linspace(1, 3, 20)
    d2 = np.linspace(0, 10, 50)
    fig, ax = plt.subplots()
    # test that kwargs normed and density cannot be set both.
    with pytest.raises(Exception):
        ax.hist((d1, d2), stacked=True, normed=normed, density=density)


@image_comparison(baseline_images=['hist_step_bottom'], extensions=['png'],
                  remove_text=True)
def test_hist_step_bottom():
    # make some data
    d1 = np.linspace(1, 3, 20)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist(d1, bottom=np.arange(10), histtype="stepfilled")


@image_comparison(baseline_images=['hist_stacked_bar'])
def test_hist_stacked_bar():
    # make some data
    d = [[100, 100, 100, 100, 200, 320, 450, 80, 20, 600, 310, 800],
         [20, 23, 50, 11, 100, 420], [120, 120, 120, 140, 140, 150, 180],
         [60, 60, 60, 60, 300, 300, 5, 5, 5, 5, 10, 300],
         [555, 555, 555, 30, 30, 30, 30, 30, 100, 100, 100, 100, 30, 30],
         [30, 30, 30, 30, 400, 400, 400, 400, 400, 400, 400, 400]]
    colors = [(0.5759849696758961, 1.0, 0.0), (0.0, 1.0, 0.350624650815206),
              (0.0, 1.0, 0.6549834156005998), (0.0, 0.6569064625276622, 1.0),
              (0.28302699607823545, 0.0, 1.0), (0.6849123462299822, 0.0, 1.0)]
    labels = ['green', 'orange', ' yellow', 'magenta', 'black']
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist(d, bins=10, histtype='barstacked', align='mid', color=colors,
            label=labels)
    ax.legend(loc='upper right', bbox_to_anchor=(1.0, 1.0), ncol=1)


def test_hist_emptydata():
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hist([[], range(10), range(10)], histtype="step")


@image_comparison(baseline_images=['transparent_markers'], remove_text=True)
def test_transparent_markers():
    np.random.seed(0)
    data = np.random.random(50)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(data, 'D', mfc='none', markersize=100)


@image_comparison(baseline_images=['rgba_markers'], remove_text=True)
def test_rgba_markers():
    fig, axs = plt.subplots(ncols=2)
    rcolors = [(1, 0, 0, 1), (1, 0, 0, 0.5)]
    bcolors = [(0, 0, 1, 1), (0, 0, 1, 0.5)]
    alphas = [None, 0.2]
    kw = dict(ms=100, mew=20)
    for i, alpha in enumerate(alphas):
        for j, rcolor in enumerate(rcolors):
            for k, bcolor in enumerate(bcolors):
                axs[i].plot(j+1, k+1, 'o', mfc=bcolor, mec=rcolor,
                            alpha=alpha, **kw)
                axs[i].plot(j+1, k+3, 'x', mec=rcolor, alpha=alpha, **kw)
    for ax in axs:
        ax.axis([-1, 4, 0, 5])


@image_comparison(baseline_images=['mollweide_grid'], remove_text=True)
def test_mollweide_grid():
    # test that both horizontal and vertical gridlines appear on the Mollweide
    # projection
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='mollweide')
    ax.grid()


def test_mollweide_forward_inverse_closure():
    # test that the round-trip Mollweide forward->inverse transformation is an
    # approximate identity
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='mollweide')

    # set up 1-degree grid in longitude, latitude
    lon = np.linspace(-np.pi, np.pi, 360)
    lat = np.linspace(-np.pi / 2.0, np.pi / 2.0, 180)
    lon, lat = np.meshgrid(lon, lat)
    ll = np.vstack((lon.flatten(), lat.flatten())).T

    # perform forward transform
    xy = ax.transProjection.transform(ll)

    # perform inverse transform
    ll2 = ax.transProjection.inverted().transform(xy)

    # compare
    np.testing.assert_array_almost_equal(ll, ll2, 3)


def test_mollweide_inverse_forward_closure():
    # test that the round-trip Mollweide inverse->forward transformation is an
    # approximate identity
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='mollweide')

    # set up grid in x, y
    x = np.linspace(0, 1, 500)
    x, y = np.meshgrid(x, x)
    xy = np.vstack((x.flatten(), y.flatten())).T

    # perform inverse transform
    ll = ax.transProjection.inverted().transform(xy)

    # perform forward transform
    xy2 = ax.transProjection.transform(ll)

    # compare
    np.testing.assert_array_almost_equal(xy, xy2, 3)


@image_comparison(baseline_images=['test_alpha'], remove_text=True)
def test_alpha():
    np.random.seed(0)
    data = np.random.random(50)

    fig = plt.figure()
    ax = fig.add_subplot(111)

    # alpha=.5 markers, solid line
    ax.plot(data, '-D', color=[1, 0, 0], mfc=[1, 0, 0, .5],
            markersize=20, lw=10)

    # everything solid by kwarg
    ax.plot(data + 2, '-D', color=[1, 0, 0, .5], mfc=[1, 0, 0, .5],
            markersize=20, lw=10,
            alpha=1)

    # everything alpha=.5 by kwarg
    ax.plot(data + 4, '-D', color=[1, 0, 0], mfc=[1, 0, 0],
            markersize=20, lw=10,
            alpha=.5)

    # everything alpha=.5 by colors
    ax.plot(data + 6, '-D', color=[1, 0, 0, .5], mfc=[1, 0, 0, .5],
            markersize=20, lw=10)

    # alpha=.5 line, solid markers
    ax.plot(data + 8, '-D', color=[1, 0, 0, .5], mfc=[1, 0, 0],
            markersize=20, lw=10)


@image_comparison(baseline_images=['eventplot', 'eventplot'], remove_text=True)
def test_eventplot():
    '''
    test that eventplot produces the correct output
    '''
    np.random.seed(0)

    data1 = np.random.random([32, 20]).tolist()
    data2 = np.random.random([6, 20]).tolist()
    data = data1 + data2
    num_datasets = len(data)

    colors1 = [[0, 1, .7]] * len(data1)
    colors2 = [[1, 0, 0],
               [0, 1, 0],
               [0, 0, 1],
               [1, .75, 0],
               [1, 0, 1],
               [0, 1, 1]]
    colors = colors1 + colors2

    lineoffsets1 = 12 + np.arange(0, len(data1)) * .33
    lineoffsets2 = [-15, -3, 1, 1.5, 6, 10]
    lineoffsets = lineoffsets1.tolist() + lineoffsets2

    linelengths1 = [.33] * len(data1)
    linelengths2 = [5, 2, 1, 1, 3, 1.5]
    linelengths = linelengths1 + linelengths2

    fig = plt.figure()
    axobj = fig.add_subplot(111)
    colls = axobj.eventplot(data, colors=colors, lineoffsets=lineoffsets,
                            linelengths=linelengths)

    num_collections = len(colls)
    assert num_collections == num_datasets

    # Reuse testcase from above for a labeled data test
    data = {"pos": data, "c": colors, "lo": lineoffsets, "ll": linelengths}
    fig = plt.figure()
    axobj = fig.add_subplot(111)
    colls = axobj.eventplot("pos", colors="c", lineoffsets="lo",
                            linelengths="ll", data=data)
    num_collections = len(colls)
    assert num_collections == num_datasets


@image_comparison(baseline_images=['test_eventplot_defaults'],
                  extensions=['png'], remove_text=True)
def test_eventplot_defaults():
    '''
    test that eventplot produces the correct output given the default params
    (see bug #3728)
    '''
    np.random.seed(0)

    data1 = np.random.random([32, 20]).tolist()
    data2 = np.random.random([6, 20]).tolist()
    data = data1 + data2

    fig = plt.figure()
    axobj = fig.add_subplot(111)
    colls = axobj.eventplot(data)


@pytest.mark.parametrize(('colors'), [
    ('0.5',),  # string color with multiple characters: not OK before #8193 fix
    ('tab:orange', 'tab:pink', 'tab:cyan', 'bLacK'),  # case-insensitive
    ('red', (0, 1, 0), None, (1, 0, 1, 0.5)),  # a tricky case mixing types
    ('rgbk',)  # len('rgbk') == len(data) and each character is a valid color
])
def test_eventplot_colors(colors):
    '''Test the *colors* parameter of eventplot. Inspired by the issue #8193.
    '''
    data = [[i] for i in range(4)]  # 4 successive events of different nature

    # Build the list of the expected colors
    expected = [c if c is not None else 'C0' for c in colors]
    # Convert the list into an array of RGBA values
    # NB: ['rgbk'] is not a valid argument for to_rgba_array, while 'rgbk' is.
    if len(expected) == 1:
        expected = expected[0]
    expected = broadcast_to(mcolors.to_rgba_array(expected), (len(data), 4))

    fig, ax = plt.subplots()
    if len(colors) == 1:  # tuple with a single string (like '0.5' or 'rgbk')
        colors = colors[0]
    collections = ax.eventplot(data, colors=colors)

    for coll, color in zip(collections, expected):
        assert_allclose(coll.get_color(), color)


@image_comparison(baseline_images=['test_eventplot_problem_kwargs'],
                  extensions=['png'], remove_text=True)
def test_eventplot_problem_kwargs():
    '''
    test that 'singular' versions of LineCollection props raise an
    IgnoredKeywordWarning rather than overriding the 'plural' versions (e.g.
    to prevent 'color' from overriding 'colors', see issue #4297)
    '''
    np.random.seed(0)

    data1 = np.random.random([20]).tolist()
    data2 = np.random.random([10]).tolist()
    data = [data1, data2]

    fig = plt.figure()
    axobj = fig.add_subplot(111)

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        colls = axobj.eventplot(data,
                                colors=['r', 'b'],
                                color=['c', 'm'],
                                linewidths=[2, 1],
                                linewidth=[1, 2],
                                linestyles=['solid', 'dashed'],
                                linestyle=['dashdot', 'dotted'])

        # check that three IgnoredKeywordWarnings were raised
        assert len(w) == 3
        assert all(issubclass(wi.category, IgnoredKeywordWarning) for wi in w)


def test_empty_eventplot():
    fig, ax = plt.subplots(1, 1)
    ax.eventplot([[]], colors=[(0.0, 0.0, 0.0, 0.0)])
    plt.draw()


@image_comparison(baseline_images=['marker_styles'], extensions=['png'],
                  remove_text=True)
def test_marker_styles():
    fig = plt.figure()
    ax = fig.add_subplot(111)
    for y, marker in enumerate(sorted(matplotlib.markers.MarkerStyle.markers,
                                      key=lambda x: str(type(x))+str(x))):
        ax.plot((y % 2)*5 + np.arange(10)*10, np.ones(10)*10*y, linestyle='',
                marker=marker, markersize=10+y/5, label=marker)


@image_comparison(baseline_images=['rc_markerfill'], extensions=['png'])
def test_markers_fillstyle_rcparams():
    fig, ax = plt.subplots()
    x = np.arange(7)
    for idx, (style, marker) in enumerate(
            [('top', 's'), ('bottom', 'o'), ('none', '^')]):
        matplotlib.rcParams['markers.fillstyle'] = style
        ax.plot(x+idx, marker=marker)


@image_comparison(baseline_images=['vertex_markers'], extensions=['png'],
                  remove_text=True)
def test_vertex_markers():
    data = list(xrange(10))
    marker_as_tuple = ((-1, -1), (1, -1), (1, 1), (-1, 1))
    marker_as_list = [(-1, -1), (1, -1), (1, 1), (-1, 1)]
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(data, linestyle='', marker=marker_as_tuple, mfc='k')
    ax.plot(data[::-1], linestyle='', marker=marker_as_list, mfc='b')
    ax.set_xlim([-1, 10])
    ax.set_ylim([-1, 10])


@image_comparison(baseline_images=['vline_hline_zorder',
                                   'errorbar_zorder'])
def test_eb_line_zorder():
    x = list(xrange(10))

    # First illustrate basic pyplot interface, using defaults where possible.
    fig = plt.figure()
    ax = fig.gca()
    ax.plot(x, lw=10, zorder=5)
    ax.axhline(1, color='red', lw=10, zorder=1)
    ax.axhline(5, color='green', lw=10, zorder=10)
    ax.axvline(7, color='m', lw=10, zorder=7)
    ax.axvline(2, color='k', lw=10, zorder=3)

    ax.set_title("axvline and axhline zorder test")

    # Now switch to a more OO interface to exercise more features.
    fig = plt.figure()
    ax = fig.gca()
    x = list(xrange(10))
    y = np.zeros(10)
    yerr = list(xrange(10))
    ax.errorbar(x, y, yerr=yerr, zorder=5, lw=5, color='r')
    for j in range(10):
        ax.axhline(j, lw=5, color='k', zorder=j)
        ax.axhline(-j, lw=5, color='k', zorder=j)

    ax.set_title("errorbar zorder test")


@image_comparison(
    baseline_images=['vlines_basic', 'vlines_with_nan', 'vlines_masked'],
    extensions=['png']
)
def test_vlines():
    # normal
    x1 = [2, 3, 4, 5, 7]
    y1 = [2, -6, 3, 8, 2]
    fig1, ax1 = plt.subplots()
    ax1.vlines(x1, 0, y1, colors='g', linewidth=5)

    # GH #7406
    x2 = [2, 3, 4, 5, 6, 7]
    y2 = [2, -6, 3, 8, np.nan, 2]
    fig2, (ax2, ax3, ax4) = plt.subplots(nrows=3, figsize=(4, 8))
    ax2.vlines(x2, 0, y2, colors='g', linewidth=5)

    x3 = [2, 3, 4, 5, 6, 7]
    y3 = [np.nan, 2, -6, 3, 8, 2]
    ax3.vlines(x3, 0, y3, colors='r', linewidth=3, linestyle='--')

    x4 = [2, 3, 4, 5, 6, 7]
    y4 = [np.nan, 2, -6, 3, 8, np.nan]
    ax4.vlines(x4, 0, y4, colors='k', linewidth=2)

    # tweak the x-axis so we can see the lines better
    for ax in [ax1, ax2, ax3, ax4]:
        ax.set_xlim(0, 10)

    # check that the y-lims are all automatically the same
    assert ax1.get_ylim() == ax2.get_ylim()
    assert ax1.get_ylim() == ax3.get_ylim()
    assert ax1.get_ylim() == ax4.get_ylim()

    fig3, ax5 = plt.subplots()
    x5 = np.ma.masked_equal([2, 4, 6, 8, 10, 12], 8)
    ymin5 = np.ma.masked_equal([0, 1, -1, 0, 2, 1], 2)
    ymax5 = np.ma.masked_equal([13, 14, 15, 16, 17, 18], 18)
    ax5.vlines(x5, ymin5, ymax5, colors='k', linewidth=2)
    ax5.set_xlim(0, 15)


@image_comparison(
    baseline_images=['hlines_basic', 'hlines_with_nan', 'hlines_masked'],
    extensions=['png']
)
def test_hlines():
    # normal
    y1 = [2, 3, 4, 5, 7]
    x1 = [2, -6, 3, 8, 2]
    fig1, ax1 = plt.subplots()
    ax1.hlines(y1, 0, x1, colors='g', linewidth=5)

    # GH #7406
    y2 = [2, 3, 4, 5, 6, 7]
    x2 = [2, -6, 3, 8, np.nan, 2]
    fig2, (ax2, ax3, ax4) = plt.subplots(nrows=3, figsize=(4, 8))
    ax2.hlines(y2, 0, x2, colors='g', linewidth=5)

    y3 = [2, 3, 4, 5, 6, 7]
    x3 = [np.nan, 2, -6, 3, 8, 2]
    ax3.hlines(y3, 0, x3, colors='r', linewidth=3, linestyle='--')

    y4 = [2, 3, 4, 5, 6, 7]
    x4 = [np.nan, 2, -6, 3, 8, np.nan]
    ax4.hlines(y4, 0, x4, colors='k', linewidth=2)

    # tweak the y-axis so we can see the lines better
    for ax in [ax1, ax2, ax3, ax4]:
        ax.set_ylim(0, 10)

    # check that the x-lims are all automatically the same
    assert ax1.get_xlim() == ax2.get_xlim()
    assert ax1.get_xlim() == ax3.get_xlim()
    assert ax1.get_xlim() == ax4.get_xlim()

    fig3, ax5 = plt.subplots()
    y5 = np.ma.masked_equal([2, 4, 6, 8, 10, 12], 8)
    xmin5 = np.ma.masked_equal([0, 1, -1, 0, 2, 1], 2)
    xmax5 = np.ma.masked_equal([13, 14, 15, 16, 17, 18], 18)
    ax5.hlines(y5, xmin5, xmax5, colors='k', linewidth=2)
    ax5.set_ylim(0, 15)


@image_comparison(baseline_images=['step_linestyle', 'step_linestyle'],
                  remove_text=True)
def test_step_linestyle():
    x = y = np.arange(10)

    # First illustrate basic pyplot interface, using defaults where possible.
    fig, ax_lst = plt.subplots(2, 2)
    ax_lst = ax_lst.flatten()

    ln_styles = ['-', '--', '-.', ':']

    for ax, ls in zip(ax_lst, ln_styles):
        ax.step(x, y, lw=5, linestyle=ls, where='pre')
        ax.step(x, y + 1, lw=5, linestyle=ls, where='mid')
        ax.step(x, y + 2, lw=5, linestyle=ls, where='post')
        ax.set_xlim([-1, 5])
        ax.set_ylim([-1, 7])

    # Reuse testcase from above for a labeled data test
    data = {"x": x, "y": y, "y1": y+1, "y2": y+2}
    fig, ax_lst = plt.subplots(2, 2)
    ax_lst = ax_lst.flatten()
    ln_styles = ['-', '--', '-.', ':']
    for ax, ls in zip(ax_lst, ln_styles):
        ax.step("x", "y", lw=5, linestyle=ls, where='pre', data=data)
        ax.step("x", "y1", lw=5, linestyle=ls, where='mid', data=data)
        ax.step("x", "y2", lw=5, linestyle=ls, where='post', data=data)
        ax.set_xlim([-1, 5])
        ax.set_ylim([-1, 7])


@image_comparison(baseline_images=['mixed_collection'], remove_text=True)
def test_mixed_collection():
    from matplotlib import patches
    from matplotlib import collections

    x = list(xrange(10))

    # First illustrate basic pyplot interface, using defaults where possible.
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)

    c = patches.Circle((8, 8), radius=4, facecolor='none', edgecolor='green')

    # PDF can optimize this one
    p1 = collections.PatchCollection([c], match_original=True)
    p1.set_offsets([[0, 0], [24, 24]])
    p1.set_linewidths([1, 5])

    # PDF can't optimize this one, because the alpha of the edge changes
    p2 = collections.PatchCollection([c], match_original=True)
    p2.set_offsets([[48, 0], [-32, -16]])
    p2.set_linewidths([1, 5])
    p2.set_edgecolors([[0, 0, 0.1, 1.0], [0, 0, 0.1, 0.5]])

    ax.patch.set_color('0.5')
    ax.add_collection(p1)
    ax.add_collection(p2)

    ax.set_xlim(0, 16)
    ax.set_ylim(0, 16)


def test_subplot_key_hash():
    ax = plt.subplot(np.float64(5.5), np.int64(1), np.float64(1.2))
    ax.twinx()
    assert (5, 1, 0, None) == ax.get_subplotspec().get_geometry()


@image_comparison(baseline_images=['specgram_freqs',
                                   'specgram_freqs_linear'],
                  remove_text=True, extensions=['png'], tol=0.07,
                  style='default')
def test_specgram_freqs():
    '''test axes.specgram in default (psd) mode with sinusoidal stimuli'''
    n = 1000
    Fs = 10.

    fstims1 = [Fs/4, Fs/5, Fs/11]
    fstims2 = [Fs/4.7, Fs/5.6, Fs/11.9]

    NFFT = int(10 * Fs / min(fstims1 + fstims2))
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y1 = np.zeros(x.size)
    y2 = np.zeros(x.size)
    for fstim1, fstim2 in zip(fstims1, fstims2):
        y1 += np.sin(fstim1 * x * np.pi * 2)
        y2 += np.sin(fstim2 * x * np.pi * 2)
    y = np.hstack([y1, y2])

    fig1 = plt.figure()
    fig2 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    ax21 = fig2.add_subplot(3, 1, 1)
    ax22 = fig2.add_subplot(3, 1, 2)
    ax23 = fig2.add_subplot(3, 1, 3)

    spec11 = ax11.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default')
    spec12 = ax12.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided')
    spec13 = ax13.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided')

    spec21 = ax21.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default',
                           scale='linear', norm=matplotlib.colors.LogNorm())
    spec22 = ax22.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided',
                           scale='linear', norm=matplotlib.colors.LogNorm())
    spec23 = ax23.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided',
                           scale='linear', norm=matplotlib.colors.LogNorm())


@image_comparison(baseline_images=['specgram_noise',
                                   'specgram_noise_linear'],
                  remove_text=True, extensions=['png'], tol=0.01,
                  style='default')
def test_specgram_noise():
    '''test axes.specgram in default (psd) mode with noise stimuli'''
    np.random.seed(0)

    n = 1000
    Fs = 10.

    NFFT = int(10 * Fs / 11)
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)
    y = np.hstack([y1, y2])

    fig1 = plt.figure()
    fig2 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    ax21 = fig2.add_subplot(3, 1, 1)
    ax22 = fig2.add_subplot(3, 1, 2)
    ax23 = fig2.add_subplot(3, 1, 3)

    spec11 = ax11.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default')
    spec12 = ax12.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided')
    spec13 = ax13.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided')

    spec21 = ax21.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default',
                           scale='linear', norm=matplotlib.colors.LogNorm())
    spec22 = ax22.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided',
                           scale='linear', norm=matplotlib.colors.LogNorm())
    spec23 = ax23.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided',
                           scale='linear', norm=matplotlib.colors.LogNorm())


@image_comparison(baseline_images=['specgram_magnitude_freqs',
                                   'specgram_magnitude_freqs_linear'],
                  remove_text=True, extensions=['png'], tol=0.07,
                  style='default')
def test_specgram_magnitude_freqs():
    '''test axes.specgram in magnitude mode with sinusoidal stimuli'''
    n = 1000
    Fs = 10.

    fstims1 = [Fs/4, Fs/5, Fs/11]
    fstims2 = [Fs/4.7, Fs/5.6, Fs/11.9]

    NFFT = int(100 * Fs / min(fstims1 + fstims2))
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y1 = np.zeros(x.size)
    y2 = np.zeros(x.size)
    for i, (fstim1, fstim2) in enumerate(zip(fstims1, fstims2)):
        y1 += np.sin(fstim1 * x * np.pi * 2)
        y2 += np.sin(fstim2 * x * np.pi * 2)
        y1[-1] = y1[-1]/y1[-1]
        y2[-1] = y2[-1]/y2[-1]
    y = np.hstack([y1, y2])

    fig1 = plt.figure()
    fig2 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    ax21 = fig2.add_subplot(3, 1, 1)
    ax22 = fig2.add_subplot(3, 1, 2)
    ax23 = fig2.add_subplot(3, 1, 3)

    spec11 = ax11.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default', mode='magnitude')
    spec12 = ax12.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided', mode='magnitude')
    spec13 = ax13.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided', mode='magnitude')

    spec21 = ax21.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default', mode='magnitude',
                           scale='linear', norm=matplotlib.colors.LogNorm())
    spec22 = ax22.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided', mode='magnitude',
                           scale='linear', norm=matplotlib.colors.LogNorm())
    spec23 = ax23.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided', mode='magnitude',
                           scale='linear', norm=matplotlib.colors.LogNorm())


@image_comparison(baseline_images=['specgram_magnitude_noise',
                                   'specgram_magnitude_noise_linear'],
                  remove_text=True, extensions=['png'],
                  style='default')
def test_specgram_magnitude_noise():
    '''test axes.specgram in magnitude mode with noise stimuli'''
    np.random.seed(0)

    n = 1000
    Fs = 10.

    NFFT = int(10 * Fs / 11)
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)
    y = np.hstack([y1, y2])

    fig1 = plt.figure()
    fig2 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    ax21 = fig2.add_subplot(3, 1, 1)
    ax22 = fig2.add_subplot(3, 1, 2)
    ax23 = fig2.add_subplot(3, 1, 3)

    spec11 = ax11.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default', mode='magnitude')
    spec12 = ax12.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided', mode='magnitude')
    spec13 = ax13.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided', mode='magnitude')

    spec21 = ax21.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default', mode='magnitude',
                           scale='linear', norm=matplotlib.colors.LogNorm())
    spec22 = ax22.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided', mode='magnitude',
                           scale='linear', norm=matplotlib.colors.LogNorm())
    spec23 = ax23.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided', mode='magnitude',
                           scale='linear', norm=matplotlib.colors.LogNorm())


@image_comparison(baseline_images=['specgram_angle_freqs'],
                  remove_text=True, extensions=['png'], tol=0.007,
                  style='default')
def test_specgram_angle_freqs():
    '''test axes.specgram in angle mode with sinusoidal stimuli'''
    n = 1000
    Fs = 10.

    fstims1 = [Fs/4, Fs/5, Fs/11]
    fstims2 = [Fs/4.7, Fs/5.6, Fs/11.9]

    NFFT = int(10 * Fs / min(fstims1 + fstims2))
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y1 = np.zeros(x.size)
    y2 = np.zeros(x.size)
    for i, (fstim1, fstim2) in enumerate(zip(fstims1, fstims2)):
        y1 += np.sin(fstim1 * x * np.pi * 2)
        y2 += np.sin(fstim2 * x * np.pi * 2)
        y1[-1] = y1[-1]/y1[-1]
        y2[-1] = y2[-1]/y2[-1]
    y = np.hstack([y1, y2])

    fig1 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    spec11 = ax11.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default', mode='angle')
    spec12 = ax12.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided', mode='angle')
    spec13 = ax13.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided', mode='angle')

    with pytest.raises(ValueError):
        ax11.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='default',
                      mode='phase', scale='dB')

    with pytest.raises(ValueError):
        ax12.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='onesided',
                      mode='phase', scale='dB')

    with pytest.raises(ValueError):
        ax13.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='twosided',
                      mode='phase', scale='dB')


@image_comparison(baseline_images=['specgram_angle_noise'],
                  remove_text=True, extensions=['png'],
                  style='default')
def test_specgram_noise_angle():
    '''test axes.specgram in angle mode with noise stimuli'''
    np.random.seed(0)

    n = 1000
    Fs = 10.

    NFFT = int(10 * Fs / 11)
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)
    y = np.hstack([y1, y2])

    fig1 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    spec11 = ax11.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default', mode='angle')
    spec12 = ax12.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided', mode='angle')
    spec13 = ax13.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided', mode='angle')

    with pytest.raises(ValueError):
        ax11.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='default',
                      mode='phase', scale='dB')

    with pytest.raises(ValueError):
        ax12.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='onesided',
                      mode='phase', scale='dB')

    with pytest.raises(ValueError):
        ax13.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='twosided',
                      mode='phase', scale='dB')


@image_comparison(baseline_images=['specgram_phase_freqs'],
                  remove_text=True, extensions=['png'],
                  style='default')
def test_specgram_freqs_phase():
    '''test axes.specgram in phase mode with sinusoidal stimuli'''
    n = 1000
    Fs = 10.

    fstims1 = [Fs/4, Fs/5, Fs/11]
    fstims2 = [Fs/4.7, Fs/5.6, Fs/11.9]

    NFFT = int(10 * Fs / min(fstims1 + fstims2))
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y1 = np.zeros(x.size)
    y2 = np.zeros(x.size)
    for i, (fstim1, fstim2) in enumerate(zip(fstims1, fstims2)):
        y1 += np.sin(fstim1 * x * np.pi * 2)
        y2 += np.sin(fstim2 * x * np.pi * 2)
        y1[-1] = y1[-1]/y1[-1]
        y2[-1] = y2[-1]/y2[-1]
    y = np.hstack([y1, y2])

    fig1 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    spec11 = ax11.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default', mode='phase')
    spec12 = ax12.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided', mode='phase')
    spec13 = ax13.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided', mode='phase')

    with pytest.raises(ValueError):
        ax11.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='default',
                      mode='phase', scale='dB')

    with pytest.raises(ValueError):
        ax12.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='onesided',
                      mode='phase', scale='dB')

    with pytest.raises(ValueError):
        ax13.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='twosided',
                      mode='phase', scale='dB')


@image_comparison(baseline_images=['specgram_phase_noise'],
                  remove_text=True, extensions=['png'],
                  style='default')
def test_specgram_noise_phase():
    '''test axes.specgram in phase mode with noise stimuli'''
    np.random.seed(0)

    n = 1000
    Fs = 10.

    NFFT = int(10 * Fs / 11)
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)
    y = np.hstack([y1, y2])

    fig1 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    spec11 = ax11.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default',
                           mode='phase', )
    spec12 = ax12.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided',
                           mode='phase', )
    spec13 = ax13.specgram(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='twosided',
                           mode='phase', )

    with pytest.raises(ValueError):
        ax11.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='default',
                      mode='phase', scale='dB')

    with pytest.raises(ValueError):
        ax12.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='onesided',
                      mode='phase', scale='dB')

    with pytest.raises(ValueError):
        ax13.specgram(y, NFFT=NFFT, Fs=Fs,
                      noverlap=noverlap, pad_to=pad_to, sides='twosided',
                      mode='phase', scale='dB')


@image_comparison(baseline_images=['psd_freqs'], remove_text=True,
                  extensions=['png'])
def test_psd_freqs():
    '''test axes.psd with sinusoidal stimuli'''
    n = 10000
    Fs = 100.

    fstims1 = [Fs/4, Fs/5, Fs/11]
    fstims2 = [Fs/4.7, Fs/5.6, Fs/11.9]

    NFFT = int(1000 * Fs / min(fstims1 + fstims2))
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y1 = np.zeros(x.size)
    y2 = np.zeros(x.size)
    for fstim1, fstim2 in zip(fstims1, fstims2):
        y1 += np.sin(fstim1 * x * np.pi * 2)
        y2 += np.sin(fstim2 * x * np.pi * 2)
    y = np.hstack([y1, y2])

    fig = plt.figure()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    psd1, freqs1 = ax1.psd(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default')
    psd2, freqs2 = ax2.psd(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided',
                           return_line=False)
    psd3, freqs3, line3 = ax3.psd(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                                  pad_to=pad_to, sides='twosided',
                                  return_line=True)

    ax1.set_xlabel('')
    ax2.set_xlabel('')
    ax3.set_xlabel('')
    ax1.set_ylabel('')
    ax2.set_ylabel('')
    ax3.set_ylabel('')


@image_comparison(baseline_images=['psd_noise'], remove_text=True,
                  extensions=['png'])
def test_psd_noise():
    '''test axes.psd with noise stimuli'''
    np.random.seed(0)

    n = 10000
    Fs = 100.

    NFFT = int(1000 * Fs / 11)
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)
    y = np.hstack([y1, y2])

    fig = plt.figure()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    psd1, freqs1 = ax1.psd(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default')
    psd2, freqs2 = ax2.psd(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided',
                           return_line=False)
    psd3, freqs3, line3 = ax3.psd(y, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                                  pad_to=pad_to, sides='twosided',
                                  return_line=True)

    ax1.set_xlabel('')
    ax2.set_xlabel('')
    ax3.set_xlabel('')
    ax1.set_ylabel('')
    ax2.set_ylabel('')
    ax3.set_ylabel('')


@image_comparison(baseline_images=['csd_freqs'], remove_text=True,
                  extensions=['png'])
def test_csd_freqs():
    '''test axes.csd with sinusoidal stimuli'''
    n = 10000
    Fs = 100.

    fstims1 = [Fs/4, Fs/5, Fs/11]
    fstims2 = [Fs/4.7, Fs/5.6, Fs/11.9]

    NFFT = int(1000 * Fs / min(fstims1 + fstims2))
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y1 = np.zeros(x.size)
    y2 = np.zeros(x.size)
    for fstim1, fstim2 in zip(fstims1, fstims2):
        y1 += np.sin(fstim1 * x * np.pi * 2)
        y2 += np.sin(fstim2 * x * np.pi * 2)

    fig = plt.figure()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    csd1, freqs1 = ax1.csd(y1, y2, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default')
    csd2, freqs2 = ax2.csd(y1, y2, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided',
                           return_line=False)
    csd3, freqs3, line3 = ax3.csd(y1, y2, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                                  pad_to=pad_to, sides='twosided',
                                  return_line=True)

    ax1.set_xlabel('')
    ax2.set_xlabel('')
    ax3.set_xlabel('')
    ax1.set_ylabel('')
    ax2.set_ylabel('')
    ax3.set_ylabel('')


@image_comparison(baseline_images=['csd_noise'], remove_text=True,
                  extensions=['png'])
def test_csd_noise():
    '''test axes.csd with noise stimuli'''
    np.random.seed(0)

    n = 10000
    Fs = 100.

    NFFT = int(1000 * Fs / 11)
    noverlap = int(NFFT / 2)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)

    fig = plt.figure()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    csd1, freqs1 = ax1.csd(y1, y2, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='default')
    csd2, freqs2 = ax2.csd(y1, y2, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                           pad_to=pad_to, sides='onesided',
                           return_line=False)
    csd3, freqs3, line3 = ax3.csd(y1, y2, NFFT=NFFT, Fs=Fs, noverlap=noverlap,
                                  pad_to=pad_to, sides='twosided',
                                  return_line=True)

    ax1.set_xlabel('')
    ax2.set_xlabel('')
    ax3.set_xlabel('')
    ax1.set_ylabel('')
    ax2.set_ylabel('')
    ax3.set_ylabel('')


@image_comparison(baseline_images=['magnitude_spectrum_freqs_linear',
                                   'magnitude_spectrum_freqs_dB'],
                  remove_text=True,
                  extensions=['png'])
def test_magnitude_spectrum_freqs():
    '''test axes.magnitude_spectrum with sinusoidal stimuli'''
    n = 10000
    Fs = 100.

    fstims1 = [Fs/4, Fs/5, Fs/11]

    NFFT = int(1000 * Fs / min(fstims1))
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y = np.zeros(x.size)
    for i, fstim1 in enumerate(fstims1):
        y += np.sin(fstim1 * x * np.pi * 2) * 10**i
    y = y

    fig1 = plt.figure()
    fig2 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    ax21 = fig2.add_subplot(3, 1, 1)
    ax22 = fig2.add_subplot(3, 1, 2)
    ax23 = fig2.add_subplot(3, 1, 3)

    spec11, freqs11, line11 = ax11.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='default')
    spec12, freqs12, line12 = ax12.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='onesided')
    spec13, freqs13, line13 = ax13.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='twosided')

    spec21, freqs21, line21 = ax21.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='default',
                                                      scale='dB')
    spec22, freqs22, line22 = ax22.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='onesided',
                                                      scale='dB')
    spec23, freqs23, line23 = ax23.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='twosided',
                                                      scale='dB')

    ax11.set_xlabel('')
    ax12.set_xlabel('')
    ax13.set_xlabel('')
    ax11.set_ylabel('')
    ax12.set_ylabel('')
    ax13.set_ylabel('')

    ax21.set_xlabel('')
    ax22.set_xlabel('')
    ax23.set_xlabel('')
    ax21.set_ylabel('')
    ax22.set_ylabel('')
    ax23.set_ylabel('')


@image_comparison(baseline_images=['magnitude_spectrum_noise_linear',
                                   'magnitude_spectrum_noise_dB'],
                  remove_text=True,
                  extensions=['png'])
def test_magnitude_spectrum_noise():
    '''test axes.magnitude_spectrum with noise stimuli'''
    np.random.seed(0)

    n = 10000
    Fs = 100.

    NFFT = int(1000 * Fs / 11)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)
    y = np.hstack([y1, y2]) - .5

    fig1 = plt.figure()
    fig2 = plt.figure()

    ax11 = fig1.add_subplot(3, 1, 1)
    ax12 = fig1.add_subplot(3, 1, 2)
    ax13 = fig1.add_subplot(3, 1, 3)

    ax21 = fig2.add_subplot(3, 1, 1)
    ax22 = fig2.add_subplot(3, 1, 2)
    ax23 = fig2.add_subplot(3, 1, 3)

    spec11, freqs11, line11 = ax11.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='default')
    spec12, freqs12, line12 = ax12.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='onesided')
    spec13, freqs13, line13 = ax13.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='twosided')

    spec21, freqs21, line21 = ax21.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='default',
                                                      scale='dB')
    spec22, freqs22, line22 = ax22.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='onesided',
                                                      scale='dB')
    spec23, freqs23, line23 = ax23.magnitude_spectrum(y, Fs=Fs, pad_to=pad_to,
                                                      sides='twosided',
                                                      scale='dB')

    ax11.set_xlabel('')
    ax12.set_xlabel('')
    ax13.set_xlabel('')
    ax11.set_ylabel('')
    ax12.set_ylabel('')
    ax13.set_ylabel('')

    ax21.set_xlabel('')
    ax22.set_xlabel('')
    ax23.set_xlabel('')
    ax21.set_ylabel('')
    ax22.set_ylabel('')
    ax23.set_ylabel('')


@image_comparison(baseline_images=['angle_spectrum_freqs'],
                  remove_text=True,
                  extensions=['png'])
def test_angle_spectrum_freqs():
    '''test axes.angle_spectrum with sinusoidal stimuli'''
    n = 10000
    Fs = 100.

    fstims1 = [Fs/4, Fs/5, Fs/11]

    NFFT = int(1000 * Fs / min(fstims1))
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y = np.zeros(x.size)
    for i, fstim1 in enumerate(fstims1):
        y += np.sin(fstim1 * x * np.pi * 2) * 10**i
    y = y

    fig = plt.figure()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    spec1, freqs1, line1 = ax1.angle_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='default')
    spec2, freqs2, line2 = ax2.angle_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='onesided')
    spec3, freqs3, line3 = ax3.angle_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='twosided')

    ax1.set_xlabel('')
    ax2.set_xlabel('')
    ax3.set_xlabel('')
    ax1.set_ylabel('')
    ax2.set_ylabel('')
    ax3.set_ylabel('')


@image_comparison(baseline_images=['angle_spectrum_noise'],
                  remove_text=True,
                  extensions=['png'])
def test_angle_spectrum_noise():
    '''test axes.angle_spectrum with noise stimuli'''
    np.random.seed(0)

    n = 10000
    Fs = 100.

    NFFT = int(1000 * Fs / 11)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)
    y = np.hstack([y1, y2]) - .5

    fig = plt.figure()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    spec1, freqs1, line1 = ax1.angle_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='default')
    spec2, freqs2, line2 = ax2.angle_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='onesided')
    spec3, freqs3, line3 = ax3.angle_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='twosided')

    ax1.set_xlabel('')
    ax2.set_xlabel('')
    ax3.set_xlabel('')
    ax1.set_ylabel('')
    ax2.set_ylabel('')
    ax3.set_ylabel('')


@image_comparison(baseline_images=['phase_spectrum_freqs'],
                  remove_text=True,
                  extensions=['png'])
def test_phase_spectrum_freqs():
    '''test axes.phase_spectrum with sinusoidal stimuli'''
    n = 10000
    Fs = 100.

    fstims1 = [Fs/4, Fs/5, Fs/11]

    NFFT = int(1000 * Fs / min(fstims1))
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    x = np.arange(0, n, 1/Fs)

    y = np.zeros(x.size)
    for i, fstim1 in enumerate(fstims1):
        y += np.sin(fstim1 * x * np.pi * 2) * 10**i
    y = y

    fig = plt.figure()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    spec1, freqs1, line1 = ax1.phase_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='default')
    spec2, freqs2, line2 = ax2.phase_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='onesided')
    spec3, freqs3, line3 = ax3.phase_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='twosided')

    ax1.set_xlabel('')
    ax2.set_xlabel('')
    ax3.set_xlabel('')
    ax1.set_ylabel('')
    ax2.set_ylabel('')
    ax3.set_ylabel('')


@image_comparison(baseline_images=['phase_spectrum_noise'],
                  remove_text=True,
                  extensions=['png'])
def test_phase_spectrum_noise():
    '''test axes.phase_spectrum with noise stimuli'''
    np.random.seed(0)

    n = 10000
    Fs = 100.

    NFFT = int(1000 * Fs / 11)
    pad_to = int(2 ** np.ceil(np.log2(NFFT)))

    y1 = np.random.standard_normal(n)
    y2 = np.random.rand(n)
    y = np.hstack([y1, y2]) - .5

    fig = plt.figure()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    spec1, freqs1, line1 = ax1.phase_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='default')
    spec2, freqs2, line2 = ax2.phase_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='onesided')
    spec3, freqs3, line3 = ax3.phase_spectrum(y, Fs=Fs, pad_to=pad_to,
                                              sides='twosided')

    ax1.set_xlabel('')
    ax2.set_xlabel('')
    ax3.set_xlabel('')
    ax1.set_ylabel('')
    ax2.set_ylabel('')
    ax3.set_ylabel('')


@image_comparison(baseline_images=['twin_spines'], remove_text=True,
                  extensions=['png'])
def test_twin_spines():

    def make_patch_spines_invisible(ax):
        ax.set_frame_on(True)
        ax.patch.set_visible(False)
        for sp in six.itervalues(ax.spines):
            sp.set_visible(False)

    fig = plt.figure(figsize=(4, 3))
    fig.subplots_adjust(right=0.75)

    host = fig.add_subplot(111)
    par1 = host.twinx()
    par2 = host.twinx()

    # Offset the right spine of par2.  The ticks and label have already been
    # placed on the right by twinx above.
    par2.spines["right"].set_position(("axes", 1.2))
    # Having been created by twinx, par2 has its frame off, so the line of
    # its detached spine is invisible.  First, activate the frame but make
    # the patch and spines invisible.
    make_patch_spines_invisible(par2)
    # Second, show the right spine.
    par2.spines["right"].set_visible(True)

    p1, = host.plot([0, 1, 2], [0, 1, 2], "b-")
    p2, = par1.plot([0, 1, 2], [0, 3, 2], "r-")
    p3, = par2.plot([0, 1, 2], [50, 30, 15], "g-")

    host.set_xlim(0, 2)
    host.set_ylim(0, 2)
    par1.set_ylim(0, 4)
    par2.set_ylim(1, 65)

    host.yaxis.label.set_color(p1.get_color())
    par1.yaxis.label.set_color(p2.get_color())
    par2.yaxis.label.set_color(p3.get_color())

    tkw = dict(size=4, width=1.5)
    host.tick_params(axis='y', colors=p1.get_color(), **tkw)
    par1.tick_params(axis='y', colors=p2.get_color(), **tkw)
    par2.tick_params(axis='y', colors=p3.get_color(), **tkw)
    host.tick_params(axis='x', **tkw)


@image_comparison(baseline_images=['twin_spines_on_top', 'twin_spines_on_top'],
                  extensions=['png'], remove_text=True)
def test_twin_spines_on_top():
    matplotlib.rcParams['axes.linewidth'] = 48.0
    matplotlib.rcParams['lines.linewidth'] = 48.0

    fig = plt.figure()
    ax1 = fig.add_subplot(1, 1, 1)

    data = np.array([[1000, 1100, 1200, 1250],
                     [310, 301, 360, 400]])

    ax2 = ax1.twinx()

    ax1.plot(data[0], data[1]/1E3, color='#BEAED4')
    ax1.fill_between(data[0], data[1]/1E3, color='#BEAED4', alpha=.8)

    ax2.plot(data[0], data[1]/1E3, color='#7FC97F')
    ax2.fill_between(data[0], data[1]/1E3, color='#7FC97F', alpha=.5)

    # Reuse testcase from above for a labeled data test
    data = {"i": data[0], "j": data[1]/1E3}
    fig = plt.figure()
    ax1 = fig.add_subplot(1, 1, 1)
    ax2 = ax1.twinx()
    ax1.plot("i", "j", color='#BEAED4', data=data)
    ax1.fill_between("i", "j", color='#BEAED4', alpha=.8, data=data)
    ax2.plot("i", "j", color='#7FC97F', data=data)
    ax2.fill_between("i", "j", color='#7FC97F', alpha=.5, data=data)


def test_rcparam_grid_minor():
    orig_grid = matplotlib.rcParams['axes.grid']
    orig_locator = matplotlib.rcParams['axes.grid.which']

    matplotlib.rcParams['axes.grid'] = True

    values = (
        (('both'), (True, True)),
        (('major'), (True, False)),
        (('minor'), (False, True))
        )

    for locator, result in values:
        matplotlib.rcParams['axes.grid.which'] = locator
        fig = plt.figure()
        ax = fig.add_subplot(1, 1, 1)
        assert (ax.xaxis._gridOnMajor, ax.xaxis._gridOnMinor) == result

    matplotlib.rcParams['axes.grid'] = orig_grid
    matplotlib.rcParams['axes.grid.which'] = orig_locator


def test_vline_limit():
    fig = plt.figure()
    ax = fig.gca()
    ax.axvline(0.5)
    ax.plot([-0.1, 0, 0.2, 0.1])
    (ymin, ymax) = ax.get_ylim()
    assert_allclose(ax.get_ylim(), (-.1, .2))


def test_empty_shared_subplots():
    # empty plots with shared axes inherit limits from populated plots
    fig, axes = plt.subplots(nrows=1, ncols=2, sharex=True, sharey=True)
    axes[0].plot([1, 2, 3], [2, 4, 6])
    x0, x1 = axes[1].get_xlim()
    y0, y1 = axes[1].get_ylim()
    assert x0 <= 1
    assert x1 >= 3
    assert y0 <= 2
    assert y1 >= 6


def test_relim_visible_only():
    x1 = (0., 10.)
    y1 = (0., 10.)
    x2 = (-10., 20.)
    y2 = (-10., 30.)

    fig = matplotlib.figure.Figure()
    ax = fig.add_subplot(111)
    ax.plot(x1, y1)
    assert ax.get_xlim() == x1
    assert ax.get_ylim() == y1
    l = ax.plot(x2, y2)
    assert ax.get_xlim() == x2
    assert ax.get_ylim() == y2
    l[0].set_visible(False)
    assert ax.get_xlim() == x2
    assert ax.get_ylim() == y2

    ax.relim(visible_only=True)
    ax.autoscale_view()

    assert ax.get_xlim() == x1
    assert ax.get_ylim() == y1


def test_text_labelsize():
    """
    tests for issue #1172
    """
    fig = plt.figure()
    ax = fig.gca()
    ax.tick_params(labelsize='large')
    ax.tick_params(direction='out')


@image_comparison(baseline_images=['pie_linewidth_0', 'pie_linewidth_0',
                                   'pie_linewidth_0'],
                  extensions=['png'])
def test_pie_linewidth_0():
    # The slices will be ordered and plotted counter-clockwise.
    labels = 'Frogs', 'Hogs', 'Dogs', 'Logs'
    sizes = [15, 30, 45, 10]
    colors = ['yellowgreen', 'gold', 'lightskyblue', 'lightcoral']
    explode = (0, 0.1, 0, 0)  # only "explode" the 2nd slice (i.e. 'Hogs')

    plt.pie(sizes, explode=explode, labels=labels, colors=colors,
            autopct='%1.1f%%', shadow=True, startangle=90,
            wedgeprops={'linewidth': 0})
    # Set aspect ratio to be equal so that pie is drawn as a circle.
    plt.axis('equal')

    # Reuse testcase from above for a labeled data test
    data = {"l": labels, "s": sizes, "c": colors, "ex": explode}
    fig = plt.figure()
    ax = fig.gca()
    ax.pie("s", explode="ex", labels="l", colors="c",
           autopct='%1.1f%%', shadow=True, startangle=90,
           wedgeprops={'linewidth': 0}, data=data)
    ax.axis('equal')

    # And again to test the pyplot functions which should also be able to be
    # called with a data kwarg
    plt.figure()
    plt.pie("s", explode="ex", labels="l", colors="c",
            autopct='%1.1f%%', shadow=True, startangle=90,
            wedgeprops={'linewidth': 0}, data=data)
    plt.axis('equal')


@image_comparison(baseline_images=['pie_center_radius'], extensions=['png'])
def test_pie_center_radius():
    # The slices will be ordered and plotted counter-clockwise.
    labels = 'Frogs', 'Hogs', 'Dogs', 'Logs'
    sizes = [15, 30, 45, 10]
    colors = ['yellowgreen', 'gold', 'lightskyblue', 'lightcoral']
    explode = (0, 0.1, 0, 0)  # only "explode" the 2nd slice (i.e. 'Hogs')

    plt.pie(sizes, explode=explode, labels=labels, colors=colors,
            autopct='%1.1f%%', shadow=True, startangle=90,
            wedgeprops={'linewidth': 0}, center=(1, 2), radius=1.5)

    plt.annotate("Center point", xy=(1, 2), xytext=(1, 1.5),
                 arrowprops=dict(arrowstyle="->",
                                 connectionstyle="arc3"))
    # Set aspect ratio to be equal so that pie is drawn as a circle.
    plt.axis('equal')


@image_comparison(baseline_images=['pie_linewidth_2'], extensions=['png'])
def test_pie_linewidth_2():
    # The slices will be ordered and plotted counter-clockwise.
    labels = 'Frogs', 'Hogs', 'Dogs', 'Logs'
    sizes = [15, 30, 45, 10]
    colors = ['yellowgreen', 'gold', 'lightskyblue', 'lightcoral']
    explode = (0, 0.1, 0, 0)  # only "explode" the 2nd slice (i.e. 'Hogs')

    plt.pie(sizes, explode=explode, labels=labels, colors=colors,
            autopct='%1.1f%%', shadow=True, startangle=90,
            wedgeprops={'linewidth': 2})
    # Set aspect ratio to be equal so that pie is drawn as a circle.
    plt.axis('equal')


@image_comparison(baseline_images=['pie_ccw_true'], extensions=['png'])
def test_pie_ccw_true():
    # The slices will be ordered and plotted counter-clockwise.
    labels = 'Frogs', 'Hogs', 'Dogs', 'Logs'
    sizes = [15, 30, 45, 10]
    colors = ['yellowgreen', 'gold', 'lightskyblue', 'lightcoral']
    explode = (0, 0.1, 0, 0)  # only "explode" the 2nd slice (i.e. 'Hogs')

    plt.pie(sizes, explode=explode, labels=labels, colors=colors,
            autopct='%1.1f%%', shadow=True, startangle=90,
            counterclock=True)
    # Set aspect ratio to be equal so that pie is drawn as a circle.
    plt.axis('equal')


@image_comparison(baseline_images=['pie_frame_grid'], extensions=['png'])
def test_pie_frame_grid():
    # The slices will be ordered and plotted counter-clockwise.
    labels = 'Frogs', 'Hogs', 'Dogs', 'Logs'
    sizes = [15, 30, 45, 10]
    colors = ['yellowgreen', 'gold', 'lightskyblue', 'lightcoral']
    # only "explode" the 2nd slice (i.e. 'Hogs')
    explode = (0, 0.1, 0, 0)

    plt.pie(sizes, explode=explode, labels=labels, colors=colors,
            autopct='%1.1f%%', shadow=True, startangle=90,
            wedgeprops={'linewidth': 0},
            frame=True, center=(2, 2))

    plt.pie(sizes[::-1], explode=explode, labels=labels, colors=colors,
            autopct='%1.1f%%', shadow=True, startangle=90,
            wedgeprops={'linewidth': 0},
            frame=True, center=(5, 2))

    plt.pie(sizes, explode=explode[::-1], labels=labels, colors=colors,
            autopct='%1.1f%%', shadow=True, startangle=90,
            wedgeprops={'linewidth': 0},
            frame=True, center=(3, 5))
    # Set aspect ratio to be equal so that pie is drawn as a circle.
    plt.axis('equal')


@image_comparison(baseline_images=['pie_rotatelabels_true'],
                  extensions=['png'])
def test_pie_rotatelabels_true():
    # The slices will be ordered and plotted counter-clockwise.
    labels = 'Hogwarts', 'Frogs', 'Dogs', 'Logs'
    sizes = [15, 30, 45, 10]
    colors = ['yellowgreen', 'gold', 'lightskyblue', 'lightcoral']
    explode = (0, 0.1, 0, 0)  # only "explode" the 2nd slice (i.e. 'Hogs')

    plt.pie(sizes, explode=explode, labels=labels, colors=colors,
            autopct='%1.1f%%', shadow=True, startangle=90,
            rotatelabels=True)
    # Set aspect ratio to be equal so that pie is drawn as a circle.
    plt.axis('equal')


@image_comparison(baseline_images=['set_get_ticklabels'], extensions=['png'])
def test_set_get_ticklabels():
    # test issue 2246
    fig, ax = plt.subplots(2)
    ha = ['normal', 'set_x/yticklabels']

    ax[0].plot(arange(10))
    ax[0].set_title(ha[0])

    ax[1].plot(arange(10))
    ax[1].set_title(ha[1])

    # set ticklabel to 1 plot in normal way
    ax[0].set_xticklabels(('a', 'b', 'c', 'd'))
    ax[0].set_yticklabels(('11', '12', '13', '14'))

    # set ticklabel to the other plot, expect the 2 plots have same label
    # setting pass get_ticklabels return value as ticklabels argument
    ax[1].set_xticklabels(ax[0].get_xticklabels())
    ax[1].set_yticklabels(ax[0].get_yticklabels())


@image_comparison(baseline_images=['o_marker_path_snap'], extensions=['png'],
                  savefig_kwarg={'dpi': 72})
def test_o_marker_path_snap():
    fig, ax = plt.subplots()
    ax.margins(.1)
    for ms in range(1, 15):
        ax.plot([1, 2, ], np.ones(2) + ms, 'o', ms=ms)

    for ms in np.linspace(1, 10, 25):
        ax.plot([3, 4, ], np.ones(2) + ms, 'o', ms=ms)


def test_margins():
    # test all ways margins can be called
    data = [1, 10]

    fig1, ax1 = plt.subplots(1, 1)
    ax1.plot(data)
    ax1.margins(1)
    assert ax1.margins() == (1, 1)

    fig2, ax2 = plt.subplots(1, 1)
    ax2.plot(data)
    ax2.margins(1, 0.5)
    assert ax2.margins() == (1, 0.5)

    fig3, ax3 = plt.subplots(1, 1)
    ax3.plot(data)
    ax3.margins(x=1, y=0.5)
    assert ax3.margins() == (1, 0.5)


def test_length_one_hist():
    fig, ax = plt.subplots()
    ax.hist(1)
    ax.hist([1])


def test_pathological_hexbin():
    # issue #2863
    out = io.BytesIO()

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        mylist = [10] * 100
        fig, ax = plt.subplots(1, 1)
        ax.hexbin(mylist, mylist)
        fig.savefig(out)
        assert len(w) == 0


def test_color_None():
    # issue 3855
    fig, ax = plt.subplots()
    ax.plot([1, 2], [1, 2], color=None)


def test_color_alias():
    # issues 4157 and 4162
    fig, ax = plt.subplots()
    line = ax.plot([0, 1], c='lime')[0]
    assert 'lime' == line.get_color()


def test_numerical_hist_label():
    fig, ax = plt.subplots()
    ax.hist([range(15)] * 5, label=range(5))
    ax.legend()


def test_unicode_hist_label():
    fig, ax = plt.subplots()
    a = (b'\xe5\xbe\x88\xe6\xbc\x82\xe4\xba\xae, ' +
         b'r\xc3\xb6m\xc3\xa4n ch\xc3\xa4r\xc3\xa1ct\xc3\xa8rs')
    b = b'\xd7\xa9\xd7\x9c\xd7\x95\xd7\x9d'
    labels = [a.decode('utf-8'),
              'hi aardvark',
              b.decode('utf-8'),
              ]

    ax.hist([range(15)] * 3, label=labels)
    ax.legend()


def test_move_offsetlabel():
    data = np.random.random(10) * 1e-22
    fig, ax = plt.subplots()
    ax.plot(data)
    ax.yaxis.tick_right()
    assert (1, 0.5) == ax.yaxis.offsetText.get_position()


@image_comparison(baseline_images=['rc_spines'], extensions=['png'],
                  savefig_kwarg={'dpi': 40})
def test_rc_spines():
    rc_dict = {
        'axes.spines.left': False,
        'axes.spines.right': False,
        'axes.spines.top': False,
        'axes.spines.bottom': False}
    with matplotlib.rc_context(rc_dict):
        fig, ax = plt.subplots()


@image_comparison(baseline_images=['rc_grid'], extensions=['png'],
                  savefig_kwarg={'dpi': 40})
def test_rc_grid():
    fig = plt.figure()
    rc_dict0 = {
        'axes.grid': True,
        'axes.grid.axis': 'both'
    }
    rc_dict1 = {
        'axes.grid': True,
        'axes.grid.axis': 'x'
    }
    rc_dict2 = {
        'axes.grid': True,
        'axes.grid.axis': 'y'
    }
    dict_list = [rc_dict0, rc_dict1, rc_dict2]

    i = 1
    for rc_dict in dict_list:
        with matplotlib.rc_context(rc_dict):
            fig.add_subplot(3, 1, i)
            i += 1


def test_rc_tick():
    d = {'xtick.bottom': False, 'xtick.top': True,
         'ytick.left': True, 'ytick.right': False}
    with plt.rc_context(rc=d):
        fig = plt.figure()
        ax1 = fig.add_subplot(1, 1, 1)
        xax = ax1.xaxis
        yax = ax1.yaxis
        # tick1On bottom/left
        assert not xax._major_tick_kw['tick1On']
        assert xax._major_tick_kw['tick2On']
        assert not xax._minor_tick_kw['tick1On']
        assert xax._minor_tick_kw['tick2On']

        assert yax._major_tick_kw['tick1On']
        assert not yax._major_tick_kw['tick2On']
        assert yax._minor_tick_kw['tick1On']
        assert not yax._minor_tick_kw['tick2On']


def test_rc_major_minor_tick():
    d = {'xtick.top': True, 'ytick.right': True,  # Enable all ticks
         'xtick.bottom': True, 'ytick.left': True,
         # Selectively disable
         'xtick.minor.bottom': False, 'xtick.major.bottom': False,
         'ytick.major.left': False, 'ytick.minor.left': False}
    with plt.rc_context(rc=d):
        fig = plt.figure()
        ax1 = fig.add_subplot(1, 1, 1)
        xax = ax1.xaxis
        yax = ax1.yaxis
        # tick1On bottom/left
        assert not xax._major_tick_kw['tick1On']
        assert xax._major_tick_kw['tick2On']
        assert not xax._minor_tick_kw['tick1On']
        assert xax._minor_tick_kw['tick2On']

        assert not yax._major_tick_kw['tick1On']
        assert yax._major_tick_kw['tick2On']
        assert not yax._minor_tick_kw['tick1On']
        assert yax._minor_tick_kw['tick2On']


def test_square_plot():
    x = np.arange(4)
    y = np.array([1., 3., 5., 7.])
    fig, ax = plt.subplots()
    ax.plot(x, y, 'mo')
    ax.axis('square')
    xlim, ylim = ax.get_xlim(), ax.get_ylim()
    assert np.diff(xlim) == np.diff(ylim)
    assert ax.get_aspect() == 'equal'


def test_no_None():
    fig, ax = plt.subplots()
    with pytest.raises(ValueError):
        plt.plot(None)
    with pytest.raises(ValueError):
        plt.plot(None, None)


def test_pcolor_fast_non_uniform():
    Z = np.arange(6).reshape((3, 2))
    X = np.array([0, 1, 2, 10])
    Y = np.array([0, 1, 2])

    plt.figure()
    ax = plt.subplot(111)
    ax.pcolorfast(X, Y, Z.T)


def test_shared_scale():
    fig, axs = plt.subplots(2, 2, sharex=True, sharey=True)

    axs[0, 0].set_xscale("log")
    axs[0, 0].set_yscale("log")

    for ax in axs.flat:
        assert ax.get_yscale() == 'log'
        assert ax.get_xscale() == 'log'

    axs[1, 1].set_xscale("linear")
    axs[1, 1].set_yscale("linear")

    for ax in axs.flat:
        assert ax.get_yscale() == 'linear'
        assert ax.get_xscale() == 'linear'


def test_violin_point_mass():
    """Violin plot should handle point mass pdf gracefully."""
    plt.violinplot(np.array([0, 0]))


def generate_errorbar_inputs():
    base_xy = cycler('x', [np.arange(5)]) + cycler('y', [np.ones((5, ))])
    err_cycler = cycler('err', [1,
                                [1, 1, 1, 1, 1],
                                [[1, 1, 1, 1, 1],
                                 [1, 1, 1, 1, 1]],
                                [[1]] * 5,
                                np.ones(5),
                                np.ones((2, 5)),
                                np.ones((5, 1)),
                                None
                                ])
    xerr_cy = cycler('xerr', err_cycler)
    yerr_cy = cycler('yerr', err_cycler)

    empty = ((cycler('x', [[]]) + cycler('y', [[]])) *
             cycler('xerr', [[], None]) * cycler('yerr', [[], None]))
    xerr_only = base_xy * xerr_cy
    yerr_only = base_xy * yerr_cy
    both_err = base_xy * yerr_cy * xerr_cy

    test_cyclers = chain(xerr_only, yerr_only, both_err, empty)

    return test_cyclers


@pytest.mark.parametrize('kwargs', generate_errorbar_inputs())
def test_errorbar_inputs_shotgun(kwargs):
    ax = plt.gca()
    eb = ax.errorbar(**kwargs)
    eb.remove()


def test_axisbg_warning():
    fig = plt.figure()
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        ax = matplotlib.axes.Axes(fig, [0, 0, 1, 1], axisbg='r')
        assert len(w) == 1
        msg = "The axisbg attribute was deprecated in version 2.0."
        assert str(w[0].message).startswith(msg)


@image_comparison(baseline_images=["dash_offset"], remove_text=True)
def test_dash_offset():
    fig, ax = plt.subplots()
    x = np.linspace(0, 10)
    y = np.ones_like(x)
    for j in range(0, 100, 2):
        ax.plot(x, j*y, ls=(j, (10, 10)), lw=5, color='k')


def test_title_location_roundtrip():
    fig, ax = plt.subplots()
    ax.set_title('aardvark')
    ax.set_title('left', loc='left')
    ax.set_title('right', loc='right')

    assert 'left' == ax.get_title(loc='left')
    assert 'right' == ax.get_title(loc='right')
    assert 'aardvark' == ax.get_title()

    with pytest.raises(ValueError):
        ax.get_title(loc='foo')
    with pytest.raises(ValueError):
        ax.set_title('fail', loc='foo')


@image_comparison(baseline_images=["loglog"], remove_text=True,
                  extensions=['png'])
def test_loglog():
    fig, ax = plt.subplots()
    x = np.arange(1, 11)
    ax.loglog(x, x**3, lw=5)
    ax.tick_params(length=25, width=2)
    ax.tick_params(length=15, width=2, which='minor')


@image_comparison(baseline_images=["test_loglog_nonpos"],
                  remove_text=True, extensions=['png'], style='mpl20')
def test_loglog_nonpos():
    fig, ax = plt.subplots(3, 3)
    x = np.arange(1, 11)
    y = x**3
    y[7] = -3.
    x[4] = -10
    for nn, mcx in enumerate(['mask', 'clip', '']):
        for mm, mcy in enumerate(['mask', 'clip', '']):
            kws = {}
            if mcx:
                kws['nonposx'] = mcx
            if mcy:
                kws['nonposy'] = mcy
            ax[mm, nn].loglog(x, y**3, lw=2, **kws)


@pytest.mark.style('default')
def test_axes_margins():
    fig, ax = plt.subplots()
    ax.plot([0, 1, 2, 3])
    assert ax.get_ybound()[0] != 0

    fig, ax = plt.subplots()
    ax.bar([0, 1, 2, 3], [1, 1, 1, 1])
    assert ax.get_ybound()[0] == 0

    fig, ax = plt.subplots()
    ax.barh([0, 1, 2, 3], [1, 1, 1, 1])
    assert ax.get_xbound()[0] == 0

    fig, ax = plt.subplots()
    ax.pcolor(np.zeros((10, 10)))
    assert ax.get_xbound() == (0, 10)
    assert ax.get_ybound() == (0, 10)

    fig, ax = plt.subplots()
    ax.pcolorfast(np.zeros((10, 10)))
    assert ax.get_xbound() == (0, 10)
    assert ax.get_ybound() == (0, 10)

    fig, ax = plt.subplots()
    ax.hist(np.arange(10))
    assert ax.get_ybound()[0] == 0

    fig, ax = plt.subplots()
    ax.imshow(np.zeros((10, 10)))
    assert ax.get_xbound() == (-0.5, 9.5)
    assert ax.get_ybound() == (-0.5, 9.5)


@pytest.fixture(params=['x', 'y'])
def shared_axis_remover(request):
    def _helper_x(ax):
        ax2 = ax.twinx()
        ax2.remove()
        ax.set_xlim(0, 15)
        r = ax.xaxis.get_major_locator()()
        assert r[-1] > 14

    def _helper_y(ax):
        ax2 = ax.twiny()
        ax2.remove()
        ax.set_ylim(0, 15)
        r = ax.yaxis.get_major_locator()()
        assert r[-1] > 14

    if request.param == 'x':
        return _helper_x
    elif request.param == 'y':
        return _helper_y
    else:
        assert False, 'Request param %s is invalid.' % (request.param, )


@pytest.fixture(params=['gca', 'subplots', 'subplots_shared', 'add_axes'])
def shared_axes_generator(request):
    # test all of the ways to get fig/ax sets
    if request.param == 'gca':
        fig = plt.figure()
        ax = fig.gca()
    elif request.param == 'subplots':
        fig, ax = plt.subplots()
    elif request.param == 'subplots_shared':
        fig, ax_lst = plt.subplots(2, 2, sharex='all', sharey='all')
        ax = ax_lst[0][0]
    elif request.param == 'add_axes':
        fig = plt.figure()
        ax = fig.add_axes([.1, .1, .8, .8])
    else:
        assert False, 'Request param %s is invalid.' % (request.param, )

    return fig, ax


def test_remove_shared_axes(shared_axes_generator, shared_axis_remover):
    # test all of the ways to get fig/ax sets
    fig, ax = shared_axes_generator
    shared_axis_remover(ax)


def test_remove_shared_axes_relim():
    fig, ax_lst = plt.subplots(2, 2, sharex='all', sharey='all')
    ax = ax_lst[0][0]
    orig_xlim = ax_lst[0][1].get_xlim()
    ax.remove()
    ax.set_xlim(0, 5)
    assert_array_equal(ax_lst[0][1].get_xlim(), orig_xlim)


def test_adjust_numtick_aspect():
    fig, ax = plt.subplots()
    ax.yaxis.get_major_locator().set_params(nbins='auto')
    ax.set_xlim(0, 1000)
    ax.set_aspect('equal')
    fig.canvas.draw()
    assert len(ax.yaxis.get_major_locator()()) == 2
    ax.set_ylim(0, 1000)
    fig.canvas.draw()
    assert len(ax.yaxis.get_major_locator()()) > 2


@image_comparison(baseline_images=["auto_numticks"], style='default',
                  extensions=['png'])
def test_auto_numticks():
    # Make tiny, empty subplots, verify that there are only 3 ticks.
    fig, axes = plt.subplots(4, 4)


@image_comparison(baseline_images=["auto_numticks_log"], style='default',
                  extensions=['png'])
def test_auto_numticks_log():
    # Verify that there are not too many ticks with a large log range.
    fig, ax = plt.subplots()
    matplotlib.rcParams['axes.autolimit_mode'] = 'round_numbers'
    ax.loglog([1e-20, 1e5], [1e-16, 10])


def test_broken_barh_empty():
    fig, ax = plt.subplots()
    ax.broken_barh([], (.1, .5))


def test_pandas_pcolormesh():
    pd = pytest.importorskip('pandas')
    from pandas.tseries import converter
    converter.register()

    time = pd.date_range('2000-01-01', periods=10)
    depth = np.arange(20)
    data = np.random.rand(20, 10)

    fig, ax = plt.subplots()
    ax.pcolormesh(time, depth, data)


def test_pandas_indexing_dates():
    pd = pytest.importorskip('pandas')
    from pandas.tseries import converter
    converter.register()

    dates = np.arange('2005-02', '2005-03', dtype='datetime64[D]')
    values = np.sin(np.array(range(len(dates))))
    df = pd.DataFrame({'dates': dates, 'values': values})

    ax = plt.gca()

    without_zero_index = df[np.array(df.index) % 2 == 1].copy()
    ax.plot('dates', 'values', data=without_zero_index)


def test_pandas_errorbar_indexing():
    pd = pytest.importorskip('pandas')
    from pandas.tseries import converter
    converter.register()

    df = pd.DataFrame(np.random.uniform(size=(5, 4)),
                      columns=['x', 'y', 'xe', 'ye'],
                      index=[1, 2, 3, 4, 5])
    fig, ax = plt.subplots()
    ax.errorbar('x', 'y', xerr='xe', yerr='ye', data=df)


def test_pandas_indexing_hist():
    pd = pytest.importorskip('pandas')
    from pandas.tseries import converter
    converter.register()

    ser_1 = pd.Series(data=[1, 2, 2, 3, 3, 4, 4, 4, 4, 5])
    ser_2 = ser_1.iloc[1:]
    fig, axes = plt.subplots()
    axes.hist(ser_2)


def test_pandas_bar_align_center():
    # Tests fix for issue 8767
    pd = pytest.importorskip('pandas')
    from pandas.tseries import converter
    converter.register()

    df = pd.DataFrame({'a': range(2), 'b': range(2)})

    fig, ax = plt.subplots(1)

    rect = ax.bar(df.loc[df['a'] == 1, 'b'],
                  df.loc[df['a'] == 1, 'b'],
                  align='center')

    fig.canvas.draw()


def test_axis_set_tick_params_labelsize_labelcolor():
    # Tests fix for issue 4346
    axis_1 = plt.subplot()
    axis_1.yaxis.set_tick_params(labelsize=30, labelcolor='red',
                                 direction='out')

    # Expected values after setting the ticks
    assert axis_1.yaxis.majorTicks[0]._size == 4.0
    assert axis_1.yaxis.majorTicks[0]._color == 'k'
    assert axis_1.yaxis.majorTicks[0]._labelsize == 30.0
    assert axis_1.yaxis.majorTicks[0]._labelcolor == 'red'


def test_none_kwargs():
    fig, ax = plt.subplots()
    ln, = ax.plot(range(32), linestyle=None)
    assert ln.get_linestyle() == '-'


def test_ls_ds_conflict():
    with pytest.raises(ValueError):
        plt.plot(range(32), linestyle='steps-pre:', drawstyle='steps-post')


def test_bar_uint8():
    xs = [0, 1, 2, 3]
    b = plt.bar(np.array(xs, dtype=np.uint8), [2, 3, 4, 5])
    for (patch, x) in zip(b.patches, xs):
        assert patch.xy[0] == x


@image_comparison(baseline_images=['date_timezone_x'], extensions=['png'])
def test_date_timezone_x():
    # Tests issue 5575
    time_index = [pytz.timezone('Canada/Eastern').localize(datetime.datetime(
        year=2016, month=2, day=22, hour=x)) for x in range(3)]

    # Same Timezone
    fig = plt.figure(figsize=(20, 12))
    plt.subplot(2, 1, 1)
    plt.plot_date(time_index, [3] * 3, tz='Canada/Eastern')

    # Different Timezone
    plt.subplot(2, 1, 2)
    plt.plot_date(time_index, [3] * 3, tz='UTC')


@image_comparison(baseline_images=['date_timezone_y'],
                  extensions=['png'])
def test_date_timezone_y():
    # Tests issue 5575
    time_index = [pytz.timezone('Canada/Eastern').localize(datetime.datetime(
        year=2016, month=2, day=22, hour=x)) for x in range(3)]

    # Same Timezone
    fig = plt.figure(figsize=(20, 12))
    plt.subplot(2, 1, 1)
    plt.plot_date([3] * 3,
                  time_index, tz='Canada/Eastern', xdate=False, ydate=True)

    # Different Timezone
    plt.subplot(2, 1, 2)
    plt.plot_date([3] * 3, time_index, tz='UTC', xdate=False, ydate=True)


@image_comparison(baseline_images=['date_timezone_x_and_y'],
                  extensions=['png'])
def test_date_timezone_x_and_y():
    # Tests issue 5575
    time_index = [pytz.timezone('UTC').localize(datetime.datetime(
        year=2016, month=2, day=22, hour=x)) for x in range(3)]

    # Same Timezone
    fig = plt.figure(figsize=(20, 12))
    plt.subplot(2, 1, 1)
    plt.plot_date(time_index, time_index, tz='UTC', ydate=True)

    # Different Timezone
    plt.subplot(2, 1, 2)
    plt.plot_date(time_index, time_index, tz='US/Eastern', ydate=True)


@image_comparison(baseline_images=['axisbelow'],
                  extensions=['png'], remove_text=True)
def test_axisbelow():
    # Test 'line' setting added in 6287.
    # Show only grids, not frame or ticks, to make this test
    # independent of future change to drawing order of those elements.
    fig, axs = plt.subplots(ncols=3, sharex=True, sharey=True)
    settings = (False, 'line', True)

    for ax, setting in zip(axs, settings):
        ax.plot((0, 10), (0, 10), lw=10, color='m')
        circ = mpatches.Circle((3, 3), color='r')
        ax.add_patch(circ)
        ax.grid(color='c', linestyle='-', linewidth=3)
        ax.tick_params(top=False, bottom=False,
                       left=False, right=False)
        for spine in ax.spines.values():
            spine.set_visible(False)
        ax.set_axisbelow(setting)


def test_offset_label_color():
    # Tests issue 6440
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.plot([1.01e9, 1.02e9, 1.03e9])
    ax.yaxis.set_tick_params(labelcolor='red')
    assert ax.yaxis.get_offset_text().get_color() == 'red'


def test_large_offset():
    fig, ax = plt.subplots()
    ax.plot((1 + np.array([0, 1.e-12])) * 1.e27)
    fig.canvas.draw()


def test_barb_units():
    fig, ax = plt.subplots()
    dates = [datetime.datetime(2017, 7, 15, 18, i) for i in range(0, 60, 10)]
    y = np.linspace(0, 5, len(dates))
    u = v = np.linspace(0, 50, len(dates))
    ax.barbs(dates, y, u, v)


def test_quiver_units():
    fig, ax = plt.subplots()
    dates = [datetime.datetime(2017, 7, 15, 18, i) for i in range(0, 60, 10)]
    y = np.linspace(0, 5, len(dates))
    u = v = np.linspace(0, 50, len(dates))
    ax.quiver(dates, y, u, v)


def test_bar_color_cycle():
    ccov = mcolors.colorConverter.to_rgb
    fig, ax = plt.subplots()
    for j in range(5):
        ln, = ax.plot(range(3))
        brs = ax.bar(range(3), range(3))
        for br in brs:
            assert ccov(ln.get_color()) == ccov(br.get_facecolor())


def test_tick_param_label_rotation():
    fix, (ax, ax2) = plt.subplots(1, 2)
    ax.plot([0, 1], [0, 1])
    ax2.plot([0, 1], [0, 1])
    ax.xaxis.set_tick_params(which='both', rotation=75)
    ax.yaxis.set_tick_params(which='both', rotation=90)
    for text in ax.get_xticklabels(which='both'):
        assert text.get_rotation() == 75
    for text in ax.get_yticklabels(which='both'):
        assert text.get_rotation() == 90

    ax2.tick_params(axis='x', labelrotation=53)
    ax2.tick_params(axis='y', rotation=35)
    for text in ax2.get_xticklabels(which='major'):
        assert text.get_rotation() == 53
    for text in ax2.get_yticklabels(which='major'):
        assert text.get_rotation() == 35


@pytest.mark.style('default')
def test_fillbetween_cycle():
    fig, ax = plt.subplots()

    for j in range(3):
        cc = ax.fill_between(range(3), range(3))
        target = mcolors.to_rgba('C{}'.format(j))
        assert tuple(cc.get_facecolors().squeeze()) == tuple(target)

    for j in range(3, 6):
        cc = ax.fill_betweenx(range(3), range(3))
        target = mcolors.to_rgba('C{}'.format(j))
        assert tuple(cc.get_facecolors().squeeze()) == tuple(target)

    target = mcolors.to_rgba('k')

    for al in ['facecolor', 'facecolors', 'color']:
        cc = ax.fill_between(range(3), range(3), **{al: 'k'})
        assert tuple(cc.get_facecolors().squeeze()) == tuple(target)

    edge_target = mcolors.to_rgba('k')
    for j, el in enumerate(['edgecolor', 'edgecolors'], start=6):
        cc = ax.fill_between(range(3), range(3), **{el: 'k'})
        face_target = mcolors.to_rgba('C{}'.format(j))
        assert tuple(cc.get_facecolors().squeeze()) == tuple(face_target)
        assert tuple(cc.get_edgecolors().squeeze()) == tuple(edge_target)


def test_log_margins():
    plt.rcParams['axes.autolimit_mode'] = 'data'
    fig, ax = plt.subplots()
    margin = 0.05
    ax.set_xmargin(margin)
    ax.semilogx([10, 100], [10, 100])
    xlim0, xlim1 = ax.get_xlim()
    transform = ax.xaxis.get_transform()
    xlim0t, xlim1t = transform.transform([xlim0, xlim1])
    x0t, x1t = transform.transform([10, 100])
    delta = (x1t - x0t) * margin
    assert_allclose([xlim0t + delta, xlim1t - delta], [x0t, x1t])


def test_color_length_mismatch():
    N = 5
    x, y = np.arange(N), np.arange(N)
    colors = np.arange(N+1)
    fig, ax = plt.subplots()
    with pytest.raises(ValueError):
        ax.scatter(x, y, c=colors)
    c_rgb = (0.5, 0.5, 0.5)
    ax.scatter(x, y, c=c_rgb)
    ax.scatter(x, y, c=[c_rgb] * N)


def test_scatter_color_masking():
    x = np.array([1, 2, 3])
    y = np.array([1, np.nan, 3])
    colors = np.array(['k', 'w', 'k'])
    linewidths = np.array([1, 2, 3])
    s = plt.scatter(x, y, color=colors, linewidths=linewidths)

    facecolors = s.get_facecolors()
    linecolors = s.get_edgecolors()
    linewidths = s.get_linewidths()
    assert_array_equal(facecolors[1], np.array([0, 0, 0, 1]))
    assert_array_equal(linecolors[1], np.array([0, 0, 0, 1]))
    assert linewidths[1] == 3


def test_eventplot_legend():
    plt.eventplot([1.0], label='Label')
    plt.legend()


def test_bar_broadcast_args():
    fig, ax = plt.subplots()
    # Check that a bar chart with a single height for all bars works.
    ax.bar(range(4), 1)
    # Check that a horizontal chart with one width works.
    ax.bar(0, 1, bottom=range(4), width=1, orientation='horizontal')
    # Check that edgecolor gets broadcasted.
    rect1, rect2 = ax.bar([0, 1], [0, 1], edgecolor=(.1, .2, .3, .4))
    assert rect1.get_edgecolor() == rect2.get_edgecolor() == (.1, .2, .3, .4)


def test_invalid_axis_limits():
    plt.plot([0, 1], [0, 1])
    with pytest.raises(ValueError):
        plt.xlim(np.nan)
    with pytest.raises(ValueError):
        plt.xlim(np.inf)
    with pytest.raises(ValueError):
        plt.ylim(np.nan)
    with pytest.raises(ValueError):
        plt.ylim(np.inf)


# Test all 4 combinations of logs/symlogs for minorticks_on()
@pytest.mark.parametrize('xscale', ['symlog', 'log'])
@pytest.mark.parametrize('yscale', ['symlog', 'log'])
def test_minorticks_on(xscale, yscale):
    ax = plt.subplot(111)
    ax.plot([1, 2, 3, 4])
    ax.set_xscale(xscale)
    ax.set_yscale(yscale)
    ax.minorticks_on()


def test_twinx_knows_limits():
    fig, ax = plt.subplots()

    ax.axvspan(1, 2)
    xtwin = ax.twinx()
    xtwin.plot([0, 0.5], [1, 2])
    # control axis
    fig2, ax2 = plt.subplots()

    ax2.axvspan(1, 2)
    ax2.plot([0, 0.5], [1, 2])

    assert((xtwin.viewLim.intervalx == ax2.viewLim.intervalx).all())


@pytest.mark.style('mpl20')
@pytest.mark.parametrize('args, kwargs, warning_count',
                         [((1, 1), {'width': 1, 'bottom': 1}, 0),
                          ((1, ), {'height': 1, 'bottom': 1}, 0),
                          ((), {'x': 1, 'height': 1}, 0),
                          ((), {'left': 1, 'height': 1}, 1)])
def test_bar_signature(args, kwargs, warning_count):
    fig, ax = plt.subplots()
    with warnings.catch_warnings(record=True) as w:
        r, = ax.bar(*args, **kwargs)

    assert r.get_width() == kwargs.get('width', 0.8)
    assert r.get_y() == kwargs.get('bottom', 0)
    assert len(w) == warning_count


@pytest.mark.style('mpl20')
@pytest.mark.parametrize('args, kwargs, warning_count',
                         [((1, 1), {'height': 1, 'left': 1}, 0),
                          ((1, ), {'width': 1, 'left': 1}, 0),
                          ((), {'y': 1, 'width': 1}, 0),
                          ((), {'bottom': 1, 'width': 1}, 1)])
def test_barh_signature(args, kwargs, warning_count):
    fig, ax = plt.subplots()
    with warnings.catch_warnings(record=True) as w:
        r, = ax.barh(*args, **kwargs)

    assert r.get_height() == kwargs.get('height', 0.8)
    assert r.get_x() == kwargs.get('left', 0)
    assert len(w) == warning_count


def test_zero_linewidth():
    # Check that setting a zero linewidth doesn't error
    plt.plot([0, 1], [0, 1], ls='--', lw=0)


def test_patch_deprecations():
    fig, ax = plt.subplots()
    with warnings.catch_warnings(record=True) as w:
        assert ax.patch == ax.axesPatch
        assert fig.patch == fig.figurePatch

    assert len(w) == 2


def test_polar_gridlines():
    fig = plt.figure()
    ax = fig.add_subplot(111, polar=True)

    # make all major grid lines lighter, only x grid lines set in 2.1.0
    ax.grid(alpha=0.2)

    # hide y tick labels, no effect in 2.1.0
    plt.setp(ax.yaxis.get_ticklabels(), visible=False)

    fig.canvas.draw()

    assert ax.xaxis.majorTicks[0].gridline.get_alpha() == .2
    assert ax.yaxis.majorTicks[0].gridline.get_alpha() == .2


def test_empty_errorbar_legend():
    fig, ax = plt.subplots()
    ax.errorbar([], [], xerr=[], label='empty y')
    ax.errorbar([], [], yerr=[], label='empty x')
    ax.legend()

social_vdgnn/lib/python3.6/site-packages/matplotlib/tests/test_axes.py
eliatrevisan/IA-MPPI-LBM
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six

from copy import copy
import io
import os
import warnings

import numpy as np
from numpy import ma
from numpy.testing import assert_array_equal

from matplotlib import (
    colors, image as mimage, mlab, patches, pyplot as plt,
    rc_context, rcParams)
from matplotlib.image import (AxesImage, BboxImage, FigureImage,
                              NonUniformImage, PcolorImage)
from matplotlib.testing.decorators import image_comparison
from matplotlib.transforms import Bbox, Affine2D, TransformedBbox

import pytest


try:
    from PIL import Image
    HAS_PIL = True
except ImportError:
    HAS_PIL = False
needs_pillow = pytest.mark.xfail(not HAS_PIL, reason='Test requires Pillow')


@image_comparison(baseline_images=['image_interps'], style='mpl20')
def test_image_interps():
    'make the basic nearest, bilinear and bicubic interps'
    X = np.arange(100)
    X = X.reshape(5, 20)

    fig = plt.figure()
    ax1 = fig.add_subplot(311)
    ax1.imshow(X, interpolation='nearest')
    ax1.set_title('three interpolations')
    ax1.set_ylabel('nearest')

    ax2 = fig.add_subplot(312)
    ax2.imshow(X, interpolation='bilinear')
    ax2.set_ylabel('bilinear')

    ax3 = fig.add_subplot(313)
    ax3.imshow(X, interpolation='bicubic')
    ax3.set_ylabel('bicubic')


@image_comparison(baseline_images=['interp_nearest_vs_none'],
                  extensions=['pdf', 'svg'], remove_text=True)
def test_interp_nearest_vs_none():
    'Test the effect of "nearest" and "none" interpolation'
    # Setting dpi to something really small makes the difference very
    # visible. This works fine with pdf, since the dpi setting doesn't
    # affect anything but images, but the agg output becomes unusably
    # small.
    rcParams['savefig.dpi'] = 3
    X = np.array([[[218, 165, 32], [122, 103, 238]],
                  [[127, 255, 0], [255, 99, 71]]], dtype=np.uint8)
    fig = plt.figure()
    ax1 = fig.add_subplot(121)
    ax1.imshow(X, interpolation='none')
    ax1.set_title('interpolation none')
    ax2 = fig.add_subplot(122)
    ax2.imshow(X, interpolation='nearest')
    ax2.set_title('interpolation nearest')


def do_figimage(suppressComposite):
    """ Helper for the next two tests """
    fig = plt.figure(figsize=(2,2), dpi=100)
    fig.suppressComposite = suppressComposite
    x,y = np.ix_(np.arange(100.0)/100.0, np.arange(100.0)/100.0)
    z = np.sin(x**2 + y**2 - x*y)
    c = np.sin(20*x**2 + 50*y**2)
    img = z + c/5

    fig.figimage(img, xo=0, yo=0, origin='lower')
    fig.figimage(img[::-1,:], xo=0, yo=100, origin='lower')
    fig.figimage(img[:,::-1], xo=100, yo=0, origin='lower')
    fig.figimage(img[::-1,::-1], xo=100, yo=100, origin='lower')

@image_comparison(baseline_images=['figimage-0'],
                  extensions=['png','pdf'])
def test_figimage0():
    'test the figimage method'

    suppressComposite = False
    do_figimage(suppressComposite)


@image_comparison(baseline_images=['figimage-1'],
                  extensions=['png','pdf'])
def test_figimage1():
    'test the figimage method'
    suppressComposite = True
    do_figimage(suppressComposite)


def test_image_python_io():
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot([1,2,3])
    buffer = io.BytesIO()
    fig.savefig(buffer)
    buffer.seek(0)
    plt.imread(buffer)


@needs_pillow
def test_imread_pil_uint16():
    img = plt.imread(os.path.join(os.path.dirname(__file__),
                     'baseline_images', 'test_image', 'uint16.tif'))
    assert (img.dtype == np.uint16)
    assert np.sum(img) == 134184960


def test_imsave():
    # The goal here is that the user can specify an output logical DPI
    # for the image, but this will not actually add any extra pixels
    # to the image, it will merely be used for metadata purposes.

    # So we do the traditional case (dpi == 1), and the new case (dpi
    # == 100) and read the resulting PNG files back in and make sure
    # the data is 100% identical.
    from numpy import random
    random.seed(1)
    data = random.rand(256, 128)

    buff_dpi1 = io.BytesIO()
    plt.imsave(buff_dpi1, data, dpi=1)

    buff_dpi100 = io.BytesIO()
    plt.imsave(buff_dpi100, data, dpi=100)

    buff_dpi1.seek(0)
    arr_dpi1 = plt.imread(buff_dpi1)

    buff_dpi100.seek(0)
    arr_dpi100 = plt.imread(buff_dpi100)

    assert arr_dpi1.shape == (256, 128, 4)
    assert arr_dpi100.shape == (256, 128, 4)

    assert_array_equal(arr_dpi1, arr_dpi100)

def test_imsave_color_alpha():
    # Test that imsave accept arrays with ndim=3 where the third dimension is
    # color and alpha without raising any exceptions, and that the data is
    # acceptably preserved through a save/read roundtrip.
    from numpy import random
    random.seed(1)

    for origin in ['lower', 'upper']:
        data = random.rand(16, 16, 4)
        buff = io.BytesIO()
        plt.imsave(buff, data, origin=origin)

        buff.seek(0)
        arr_buf = plt.imread(buff)

        # Recreate the float -> uint8 conversion of the data
        # We can only expect to be the same with 8 bits of precision,
        # since that's what the PNG file used.
        data = (255*data).astype('uint8')
        if origin == 'lower':
            data = data[::-1]
        arr_buf = (255*arr_buf).astype('uint8')

        assert_array_equal(data, arr_buf)

@image_comparison(baseline_images=['image_alpha'], remove_text=True)
def test_image_alpha():
    plt.figure()

    np.random.seed(0)
    Z = np.random.rand(6, 6)

    plt.subplot(131)
    plt.imshow(Z, alpha=1.0, interpolation='none')

    plt.subplot(132)
    plt.imshow(Z, alpha=0.5, interpolation='none')

    plt.subplot(133)
    plt.imshow(Z, alpha=0.5, interpolation='nearest')

def test_cursor_data():
    from matplotlib.backend_bases import MouseEvent

    fig, ax = plt.subplots()
    im = ax.imshow(np.arange(100).reshape(10, 10), origin='upper')

    x, y = 4, 4
    xdisp, ydisp = ax.transData.transform_point([x, y])

    event = MouseEvent('motion_notify_event', fig.canvas, xdisp, ydisp)
    z = im.get_cursor_data(event)
    assert z == 44, "Did not get 44, got %d" % z

    # Now try for a point outside the image
    # Tests issue #4957
    x, y = 10.1, 4
    xdisp, ydisp = ax.transData.transform_point([x, y])

    event = MouseEvent('motion_notify_event', fig.canvas, xdisp, ydisp)
    z = im.get_cursor_data(event)
    assert z is None, "Did not get None, got %d" % z

    # Hmm, something is wrong here... I get 0, not None...
    # But, this works further down in the tests with extents flipped
    #x, y = 0.1, -0.1
    #xdisp, ydisp = ax.transData.transform_point([x, y])
    #event = MouseEvent('motion_notify_event', fig.canvas, xdisp, ydisp)
    #z = im.get_cursor_data(event)
    #assert z is None, "Did not get None, got %d" % z

    ax.clear()
    # Now try with the extents flipped.
    im = ax.imshow(np.arange(100).reshape(10, 10), origin='lower')

    x, y = 4, 4
    xdisp, ydisp = ax.transData.transform_point([x, y])

    event = MouseEvent('motion_notify_event', fig.canvas, xdisp, ydisp)
    z = im.get_cursor_data(event)
    assert z == 44, "Did not get 44, got %d" % z

    fig, ax = plt.subplots()
    im = ax.imshow(np.arange(100).reshape(10, 10), extent=[0, 0.5, 0, 0.5])

    x, y = 0.25, 0.25
    xdisp, ydisp = ax.transData.transform_point([x, y])

    event = MouseEvent('motion_notify_event', fig.canvas, xdisp, ydisp)
    z = im.get_cursor_data(event)
    assert z == 55, "Did not get 55, got %d" % z

    # Now try for a point outside the image
    # Tests issue #4957
    x, y = 0.75, 0.25
    xdisp, ydisp = ax.transData.transform_point([x, y])

    event = MouseEvent('motion_notify_event', fig.canvas, xdisp, ydisp)
    z = im.get_cursor_data(event)
    assert z is None, "Did not get None, got %d" % z

    x, y = 0.01, -0.01
    xdisp, ydisp = ax.transData.transform_point([x, y])

    event = MouseEvent('motion_notify_event', fig.canvas, xdisp, ydisp)
    z = im.get_cursor_data(event)
    assert z is None, "Did not get None, got %d" % z


@image_comparison(baseline_images=['image_clip'], style='mpl20')
def test_image_clip():
    d = [[1, 2], [3, 4]]

    fig, ax = plt.subplots()
    im = ax.imshow(d)
    patch = patches.Circle((0, 0), radius=1, transform=ax.transData)
    im.set_clip_path(patch)


@image_comparison(baseline_images=['image_cliprect'], style='mpl20')
def test_image_cliprect():
    import matplotlib.patches as patches

    fig = plt.figure()
    ax = fig.add_subplot(111)
    d = [[1,2],[3,4]]

    im = ax.imshow(d, extent=(0,5,0,5))

    rect = patches.Rectangle(xy=(1,1), width=2, height=2, transform=im.axes.transData)
    im.set_clip_path(rect)


@image_comparison(baseline_images=['imshow'], remove_text=True, style='mpl20')
def test_imshow():
    import numpy as np
    import matplotlib.pyplot as plt

    fig = plt.figure()
    arr = np.arange(100).reshape((10, 10))
    ax = fig.add_subplot(111)
    ax.imshow(arr, interpolation="bilinear", extent=(1,2,1,2))
    ax.set_xlim(0,3)
    ax.set_ylim(0,3)

@image_comparison(baseline_images=['no_interpolation_origin'], remove_text=True)
def test_no_interpolation_origin():
    fig = plt.figure()
    ax = fig.add_subplot(211)
    ax.imshow(np.arange(100).reshape((2, 50)), origin="lower", interpolation='none')

    ax = fig.add_subplot(212)
    ax.imshow(np.arange(100).reshape((2, 50)), interpolation='none')

@image_comparison(baseline_images=['image_shift'], remove_text=True,
                  extensions=['pdf', 'svg'])
def test_image_shift():
    from matplotlib.colors import LogNorm

    imgData = [[1.0/(x) + 1.0/(y) for x in range(1,100)] for y in range(1,100)]
    tMin=734717.945208
    tMax=734717.946366

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.imshow(imgData, norm=LogNorm(), interpolation='none',
              extent=(tMin, tMax, 1, 100))
    ax.set_aspect('auto')

def test_image_edges():
    f = plt.figure(figsize=[1, 1])
    ax = f.add_axes([0, 0, 1, 1], frameon=False)

    data = np.tile(np.arange(12), 15).reshape(20, 9)

    im = ax.imshow(data, origin='upper', extent=[-10, 10, -10, 10],
                   interpolation='none', cmap='gray')

    x = y = 2
    ax.set_xlim([-x, x])
    ax.set_ylim([-y, y])

    ax.set_xticks([])
    ax.set_yticks([])

    buf = io.BytesIO()
    f.savefig(buf, facecolor=(0, 1, 0))

    buf.seek(0)

    im = plt.imread(buf)
    r, g, b, a = sum(im[:, 0])
    r, g, b, a = sum(im[:, -1])

    assert g != 100, 'Expected a non-green edge - but sadly, it was.'


@image_comparison(baseline_images=['image_composite_background'],
                  remove_text=True,
                  style='mpl20')
def test_image_composite_background():
    fig = plt.figure()
    ax = fig.add_subplot(111)
    arr = np.arange(12).reshape(4, 3)
    ax.imshow(arr, extent=[0, 2, 15, 0])
    ax.imshow(arr, extent=[4, 6, 15, 0])
    ax.set_facecolor((1, 0, 0, 0.5))
    ax.set_xlim([0, 12])


@image_comparison(baseline_images=['image_composite_alpha'],
                  remove_text=True)
def test_image_composite_alpha():
    """
    Tests that the alpha value is recognized and correctly applied in the
    process of compositing images together.
    """
    fig = plt.figure()
    ax = fig.add_subplot(111)
    arr = np.zeros((11, 21, 4))
    arr[:, :, 0] = 1
    arr[:, :, 3] = np.concatenate((np.arange(0, 1.1, 0.1), np.arange(0, 1, 0.1)[::-1]))
    arr2 = np.zeros((21, 11, 4))
    arr2[:, :, 0] = 1
    arr2[:, :, 1] = 1
    arr2[:, :, 3] = np.concatenate((np.arange(0, 1.1, 0.1), np.arange(0, 1, 0.1)[::-1]))[:, np.newaxis]
    ax.imshow(arr, extent=[1, 2, 5, 0], alpha=0.3)
    ax.imshow(arr, extent=[2, 3, 5, 0], alpha=0.6)
    ax.imshow(arr, extent=[3, 4, 5, 0])
    ax.imshow(arr2, extent=[0, 5, 1, 2])
    ax.imshow(arr2, extent=[0, 5, 2, 3], alpha=0.6)
    ax.imshow(arr2, extent=[0, 5, 3, 4], alpha=0.3)
    ax.set_facecolor((0, 0.5, 0, 1))
    ax.set_xlim([0, 5])
    ax.set_ylim([5, 0])


@image_comparison(baseline_images=['rasterize_10dpi'],
                  extensions=['pdf', 'svg'],
                  remove_text=True, style='mpl20')
def test_rasterize_dpi():
    # This test should check rasterized rendering with high output resolution.
    # It plots a rasterized line and a normal image with implot. So it will catch
    # when images end up in the wrong place in case of non-standard dpi setting.
    # Instead of high-res rasterization i use low-res.  Therefore the fact that the
    # resolution is non-standard is easily checked by image_comparison.
    import numpy as np
    import matplotlib.pyplot as plt

    img = np.asarray([[1, 2], [3, 4]])

    fig, axes = plt.subplots(1, 3, figsize = (3, 1))

    axes[0].imshow(img)

    axes[1].plot([0,1],[0,1], linewidth=20., rasterized=True)
    axes[1].set(xlim = (0,1), ylim = (-1, 2))

    axes[2].plot([0,1],[0,1], linewidth=20.)
    axes[2].set(xlim = (0,1), ylim = (-1, 2))

    # Low-dpi PDF rasterization errors prevent proper image comparison tests.
    # Hide detailed structures like the axes spines.
    for ax in axes:
        ax.set_xticks([])
        ax.set_yticks([])
        for spine in ax.spines.values():
            spine.set_visible(False)

    rcParams['savefig.dpi'] = 10


@image_comparison(baseline_images=['bbox_image_inverted'], remove_text=True,
                  style='mpl20')
def test_bbox_image_inverted():
    # This is just used to produce an image to feed to BboxImage
    image = np.arange(100).reshape((10, 10))

    ax = plt.subplot(111)
    bbox_im = BboxImage(
        TransformedBbox(Bbox([[100, 100], [0, 0]]), ax.transData))
    bbox_im.set_data(image)
    bbox_im.set_clip_on(False)
    ax.set_xlim(0, 100)
    ax.set_ylim(0, 100)
    ax.add_artist(bbox_im)

    image = np.identity(10)

    bbox_im = BboxImage(
        TransformedBbox(Bbox([[0.1, 0.2], [0.3, 0.25]]), ax.figure.transFigure))
    bbox_im.set_data(image)
    bbox_im.set_clip_on(False)
    ax.add_artist(bbox_im)


def test_get_window_extent_for_AxisImage():
    # Create a figure of known size (1000x1000 pixels), place an image
    # object at a given location and check that get_window_extent()
    # returns the correct bounding box values (in pixels).

    im = np.array([[0.25, 0.75, 1.0, 0.75], [0.1, 0.65, 0.5, 0.4], \
        [0.6, 0.3, 0.0, 0.2], [0.7, 0.9, 0.4, 0.6]])
    fig = plt.figure(figsize=(10, 10), dpi=100)
    ax = plt.subplot()
    ax.set_position([0, 0, 1, 1])
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    im_obj = ax.imshow(im, extent=[0.4, 0.7, 0.2, 0.9], interpolation='nearest')

    fig.canvas.draw()
    renderer = fig.canvas.renderer
    im_bbox = im_obj.get_window_extent(renderer)

    assert_array_equal(im_bbox.get_points(), [[400, 200], [700, 900]])


@image_comparison(baseline_images=['zoom_and_clip_upper_origin'],
                  remove_text=True,
                  extensions=['png'],
                  style='mpl20')
def test_zoom_and_clip_upper_origin():
    image = np.arange(100)
    image = image.reshape((10, 10))

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.imshow(image)
    ax.set_ylim(2.0, -0.5)
    ax.set_xlim(-0.5, 2.0)


def test_nonuniformimage_setcmap():
    ax = plt.gca()
    im = NonUniformImage(ax)
    im.set_cmap('Blues')


def test_nonuniformimage_setnorm():
    ax = plt.gca()
    im = NonUniformImage(ax)
    im.set_norm(plt.Normalize())


@needs_pillow
def test_jpeg_alpha():
    plt.figure(figsize=(1, 1), dpi=300)
    # Create an image that is all black, with a gradient from 0-1 in
    # the alpha channel from left to right.
    im = np.zeros((300, 300, 4), dtype=float)
    im[..., 3] = np.linspace(0.0, 1.0, 300)

    plt.figimage(im)

    buff = io.BytesIO()
    with rc_context({'savefig.facecolor': 'red'}):
        plt.savefig(buff, transparent=True, format='jpg', dpi=300)

    buff.seek(0)
    image = Image.open(buff)

    # If this fails, there will be only one color (all black). If this
    # is working, we should have all 256 shades of grey represented.
    num_colors = len(image.getcolors(256))
    assert 175 <= num_colors <= 185, 'num colors: %d' % (num_colors, )
    # The fully transparent part should be red, not white or black
    # or anything else
    corner_pixel = image.getpixel((0, 0))
    assert corner_pixel == (254, 0, 0), "corner pixel: %r" % (corner_pixel, )


def test_nonuniformimage_setdata():
    ax = plt.gca()
    im = NonUniformImage(ax)
    x = np.arange(3, dtype=float)
    y = np.arange(4, dtype=float)
    z = np.arange(12, dtype=float).reshape((4, 3))
    im.set_data(x, y, z)
    x[0] = y[0] = z[0, 0] = 9.9
    assert im._A[0, 0] == im._Ax[0] == im._Ay[0] == 0, 'value changed'


def test_axesimage_setdata():
    ax = plt.gca()
    im = AxesImage(ax)
    z = np.arange(12, dtype=float).reshape((4, 3))
    im.set_data(z)
    z[0, 0] = 9.9
    assert im._A[0, 0] == 0, 'value changed'


def test_figureimage_setdata():
    fig = plt.gcf()
    im = FigureImage(fig)
    z = np.arange(12, dtype=float).reshape((4, 3))
    im.set_data(z)
    z[0, 0] = 9.9
    assert im._A[0, 0] == 0, 'value changed'


def test_pcolorimage_setdata():
    ax = plt.gca()
    im = PcolorImage(ax)
    x = np.arange(3, dtype=float)
    y = np.arange(4, dtype=float)
    z = np.arange(6, dtype=float).reshape((3, 2))
    im.set_data(x, y, z)
    x[0] = y[0] = z[0, 0] = 9.9
    assert im._A[0, 0] == im._Ax[0] == im._Ay[0] == 0, 'value changed'


def test_pcolorimage_extent():
    im = plt.hist2d([1, 2, 3], [3, 5, 6],
                    bins=[[0, 3, 7], [1, 2, 3]])[-1]
    assert im.get_extent() == (0, 7, 1, 3)


def test_minimized_rasterized():
    # This ensures that the rasterized content in the colorbars is
    # only as thick as the colorbar, and doesn't extend to other parts
    # of the image.  See #5814.  While the original bug exists only
    # in Postscript, the best way to detect it is to generate SVG
    # and then parse the output to make sure the two colorbar images
    # are the same size.
    from xml.etree import ElementTree

    np.random.seed(0)
    data = np.random.rand(10, 10)

    fig, ax = plt.subplots(1, 2)
    p1 = ax[0].pcolormesh(data)
    p2 = ax[1].pcolormesh(data)

    plt.colorbar(p1, ax=ax[0])
    plt.colorbar(p2, ax=ax[1])

    buff = io.BytesIO()
    plt.savefig(buff, format='svg')

    buff = io.BytesIO(buff.getvalue())
    tree = ElementTree.parse(buff)
    width = None
    for image in tree.iter('image'):
        if width is None:
            width = image['width']
        else:
            if image['width'] != width:
                assert False


@pytest.mark.network
def test_load_from_url():
    req = six.moves.urllib.request.urlopen(
        "http://matplotlib.org/_static/logo_sidebar_horiz.png")
    Z = plt.imread(req)


@image_comparison(baseline_images=['log_scale_image'],
                  remove_text=True)
# The recwarn fixture captures a warning in image_comparison.
def test_log_scale_image(recwarn):
    Z = np.zeros((10, 10))
    Z[::2] = 1

    fig = plt.figure()
    ax = fig.add_subplot(111)

    ax.imshow(Z, extent=[1, 100, 1, 100], cmap='viridis',
              vmax=1, vmin=-1)
    ax.set_yscale('log')


@image_comparison(baseline_images=['rotate_image'],
                  remove_text=True)
def test_rotate_image():
    delta = 0.25
    x = y = np.arange(-3.0, 3.0, delta)
    X, Y = np.meshgrid(x, y)
    Z1 = mlab.bivariate_normal(X, Y, 1.0, 1.0, 0.0, 0.0)
    Z2 = mlab.bivariate_normal(X, Y, 1.5, 0.5, 1, 1)
    Z = Z2 - Z1  # difference of Gaussians

    fig, ax1 = plt.subplots(1, 1)
    im1 = ax1.imshow(Z, interpolation='none', cmap='viridis',
                     origin='lower',
                     extent=[-2, 4, -3, 2], clip_on=True)

    trans_data2 = Affine2D().rotate_deg(30) + ax1.transData
    im1.set_transform(trans_data2)

    # display intended extent of the image
    x1, x2, y1, y2 = im1.get_extent()

    ax1.plot([x1, x2, x2, x1, x1], [y1, y1, y2, y2, y1], "r--", lw=3,
             transform=trans_data2)

    ax1.set_xlim(2, 5)
    ax1.set_ylim(0, 4)


def test_image_preserve_size():
    buff = io.BytesIO()

    im = np.zeros((481, 321))
    plt.imsave(buff, im)

    buff.seek(0)
    img = plt.imread(buff)

    assert img.shape[:2] == im.shape


def test_image_preserve_size2():
    n = 7
    data = np.identity(n, float)

    fig = plt.figure(figsize=(n, n), frameon=False)

    ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0])
    ax.set_axis_off()
    fig.add_axes(ax)
    ax.imshow(data, interpolation='nearest', origin='lower',aspect='auto')
    buff = io.BytesIO()
    fig.savefig(buff, dpi=1)

    buff.seek(0)
    img = plt.imread(buff)

    assert img.shape == (7, 7, 4)

    assert_array_equal(np.asarray(img[:, :, 0], bool),
                       np.identity(n, bool)[::-1])


@image_comparison(baseline_images=['mask_image_over_under'],
                  remove_text=True, extensions=['png'])
def test_mask_image_over_under():
    delta = 0.025
    x = y = np.arange(-3.0, 3.0, delta)
    X, Y = np.meshgrid(x, y)
    Z1 = mlab.bivariate_normal(X, Y, 1.0, 1.0, 0.0, 0.0)
    Z2 = mlab.bivariate_normal(X, Y, 1.5, 0.5, 1, 1)
    Z = 10*(Z2 - Z1)  # difference of Gaussians

    palette = copy(plt.cm.gray)
    palette.set_over('r', 1.0)
    palette.set_under('g', 1.0)
    palette.set_bad('b', 1.0)
    Zm = ma.masked_where(Z > 1.2, Z)
    fig, (ax1, ax2) = plt.subplots(1, 2)
    im = ax1.imshow(Zm, interpolation='bilinear',
                    cmap=palette,
                    norm=colors.Normalize(vmin=-1.0, vmax=1.0, clip=False),
                    origin='lower', extent=[-3, 3, -3, 3])
    ax1.set_title('Green=low, Red=high, Blue=bad')
    fig.colorbar(im, extend='both', orientation='horizontal',
                 ax=ax1, aspect=10)

    im = ax2.imshow(Zm, interpolation='nearest',
                    cmap=palette,
                    norm=colors.BoundaryNorm([-1, -0.5, -0.2, 0, 0.2, 0.5, 1],
                                             ncolors=256, clip=False),
                    origin='lower', extent=[-3, 3, -3, 3])
    ax2.set_title('With BoundaryNorm')
    fig.colorbar(im, extend='both', spacing='proportional',
                 orientation='horizontal', ax=ax2, aspect=10)


@image_comparison(baseline_images=['mask_image'],
                  remove_text=True)
def test_mask_image():
    # Test mask image two ways: Using nans and using a masked array.

    fig, (ax1, ax2) = plt.subplots(1, 2)

    A = np.ones((5, 5))
    A[1:2, 1:2] = np.nan

    ax1.imshow(A, interpolation='nearest')

    A = np.zeros((5, 5), dtype=bool)
    A[1:2, 1:2] = True
    A = np.ma.masked_array(np.ones((5, 5), dtype=np.uint16), A)

    ax2.imshow(A, interpolation='nearest')


@image_comparison(baseline_images=['imshow_endianess'],
                  remove_text=True, extensions=['png'])
def test_imshow_endianess():
    x = np.arange(10)
    X, Y = np.meshgrid(x, x)
    Z = ((X-5)**2 + (Y-5)**2)**0.5

    fig, (ax1, ax2) = plt.subplots(1, 2)

    kwargs = dict(origin="lower", interpolation='nearest',
                  cmap='viridis')

    ax1.imshow(Z.astype('<f8'), **kwargs)
    ax2.imshow(Z.astype('>f8'), **kwargs)


@image_comparison(baseline_images=['imshow_masked_interpolation'],
                  remove_text=True, style='mpl20')
def test_imshow_masked_interpolation():

    cm = copy(plt.get_cmap('viridis'))
    cm.set_over('r')
    cm.set_under('b')
    cm.set_bad('k')

    N = 20
    n = colors.Normalize(vmin=0, vmax=N*N-1)

    # data = np.random.random((N, N))*N*N
    data = np.arange(N*N, dtype='float').reshape(N, N)

    data[5, 5] = -1
    # This will cause crazy ringing for the higher-order
    # interpolations
    data[15, 5] = 1e5

    # data[3, 3] = np.nan

    data[15, 15] = np.inf

    mask = np.zeros_like(data).astype('bool')
    mask[5, 15] = True

    data = np.ma.masked_array(data, mask)

    fig, ax_grid = plt.subplots(3, 6)

    for interp, ax in zip(sorted(mimage._interpd_), ax_grid.ravel()):
        ax.set_title(interp)
        ax.imshow(data, norm=n, cmap=cm, interpolation=interp)
        ax.axis('off')


def test_imshow_no_warn_invalid():
    with warnings.catch_warnings(record=True) as warns:
        warnings.simplefilter("always")
        plt.imshow([[1, 2], [3, np.nan]])
    assert len(warns) == 0


@image_comparison(baseline_images=['imshow_flatfield'],
                  remove_text=True, style='mpl20',
                  extensions=['png'])
def test_imshow_flatfield():
    fig, ax = plt.subplots()
    im = ax.imshow(np.ones((5, 5)))
    im.set_clim(.5, 1.5)


@image_comparison(baseline_images=['imshow_bignumbers'],
                  remove_text=True, style='mpl20',
                  extensions=['png'])
def test_imshow_bignumbers():
    # putting a big number in an array of integers shouldn't
    # ruin the dynamic range of the resolved bits.
    fig, ax = plt.subplots()
    img = np.array([[1, 2, 1e12],[3, 1, 4]], dtype=np.uint64)
    pc = ax.imshow(img)
    pc.set_clim(0, 5)
    plt.show()


@pytest.mark.parametrize(
    "make_norm",
    [colors.Normalize,
     colors.LogNorm,
     lambda: colors.SymLogNorm(1),
     lambda: colors.PowerNorm(1)])
def test_empty_imshow(make_norm):
    fig, ax = plt.subplots()
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore", "Attempting to set identical left==right")
        im = ax.imshow([[]], norm=make_norm())
    im.set_extent([-5, 5, -5, 5])
    fig.canvas.draw()

    with pytest.raises(RuntimeError):
        im.make_image(fig._cachedRenderer)


def test_imshow_float128():
    fig, ax = plt.subplots()
    ax.imshow(np.zeros((3, 3), dtype=np.longdouble))


def test_imshow_bool():
    fig, ax = plt.subplots()
    ax.imshow(np.array([[True, False], [False, True]], dtype=bool))


def test_imshow_deprecated_interd_warn():
    im = plt.imshow([[1, 2], [3, np.nan]])
    for k in ('_interpd', '_interpdr', 'iterpnames'):
        with warnings.catch_warnings(record=True) as warns:
            getattr(im, k)
        assert len(warns) == 1


def test_full_invalid():
    x = np.ones((10, 10))
    x[:] = np.nan

    f, ax = plt.subplots()
    ax.imshow(x)

    f.canvas.draw()

social_vdgnn/lib/python3.6/site-packages/matplotlib/tests/test_image.py
eliatrevisan/IA-MPPI-LBM
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

try:
    # mock in python 3.3+
    from unittest import mock
except ImportError:
    import mock
import collections
import numpy as np
import pytest


from matplotlib.testing.decorators import image_comparison
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.transforms as mtransforms
import matplotlib.collections as mcollections
from matplotlib.legend_handler import HandlerTuple
import matplotlib.legend as mlegend
import inspect


# test that docstrigs are the same
def get_docstring_section(func, section):
    """ extract a section from the docstring of a function """
    ll = inspect.getdoc(func)
    lines = ll.splitlines()
    insec = False
    st = ''
    for ind in range(len(lines)):
        if lines[ind][:len(section)] == section and lines[ind+1][:3] == '---':
            insec = True
            ind = ind+1
        if insec:
            if len(lines[ind + 1]) > 3 and lines[ind + 1][0:3] == '---':
                insec = False
                break
            else:
                st += lines[ind] + '\n'
    return st


def test_legend_kwdocstrings():
    stleg = get_docstring_section(mpl.legend.Legend.__init__, 'Parameters')
    stax = get_docstring_section(mpl.axes.Axes.legend, 'Parameters')
    stfig = get_docstring_section(mpl.figure.Figure.legend, 'Parameters')
    assert stleg == stax
    assert stfig == stax
    assert stleg == stfig


def test_legend_ordereddict():
    # smoketest that ordereddict inputs work...

    X = np.random.randn(10)
    Y = np.random.randn(10)
    labels = ['a'] * 5 + ['b'] * 5
    colors = ['r'] * 5 + ['g'] * 5

    fig, ax = plt.subplots()
    for x, y, label, color in zip(X, Y, labels, colors):
        ax.scatter(x, y, label=label, c=color)

    handles, labels = ax.get_legend_handles_labels()
    legend = collections.OrderedDict(zip(labels, handles))
    ax.legend(legend.values(), legend.keys(), loc=6, bbox_to_anchor=(1, .5))


@image_comparison(baseline_images=['legend_auto1'], remove_text=True)
def test_legend_auto1():
    'Test automatic legend placement'
    fig = plt.figure()
    ax = fig.add_subplot(111)
    x = np.arange(100)
    ax.plot(x, 50 - x, 'o', label='y=1')
    ax.plot(x, x - 50, 'o', label='y=-1')
    ax.legend(loc=0)


@image_comparison(baseline_images=['legend_auto2'], remove_text=True)
def test_legend_auto2():
    'Test automatic legend placement'
    fig = plt.figure()
    ax = fig.add_subplot(111)
    x = np.arange(100)
    b1 = ax.bar(x, x, color='m')
    b2 = ax.bar(x, x[::-1], color='g')
    ax.legend([b1[0], b2[0]], ['up', 'down'], loc=0)


@image_comparison(baseline_images=['legend_auto3'])
def test_legend_auto3():
    'Test automatic legend placement'
    fig = plt.figure()
    ax = fig.add_subplot(111)
    x = [0.9, 0.1, 0.1, 0.9, 0.9, 0.5]
    y = [0.95, 0.95, 0.05, 0.05, 0.5, 0.5]
    ax.plot(x, y, 'o-', label='line')
    ax.set_xlim(0.0, 1.0)
    ax.set_ylim(0.0, 1.0)
    ax.legend(loc=0)


@image_comparison(baseline_images=['legend_various_labels'], remove_text=True)
def test_various_labels():
    # tests all sorts of label types
    fig = plt.figure()
    ax = fig.add_subplot(121)
    ax.plot(np.arange(4), 'o', label=1)
    ax.plot(np.linspace(4, 4.1), 'o', label='D\xe9velopp\xe9s')
    ax.plot(np.arange(4, 1, -1), 'o', label='__nolegend__')
    ax.legend(numpoints=1, loc=0)


@image_comparison(baseline_images=['legend_labels_first'], extensions=['png'],
                  remove_text=True)
def test_labels_first():
    # test labels to left of markers
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(np.arange(10), '-o', label=1)
    ax.plot(np.ones(10)*5, ':x', label="x")
    ax.plot(np.arange(20, 10, -1), 'd', label="diamond")
    ax.legend(loc=0, markerfirst=False)


@image_comparison(baseline_images=['legend_multiple_keys'], extensions=['png'],
                  remove_text=True)
def test_multiple_keys():
    # test legend entries with multiple keys
    fig = plt.figure()
    ax = fig.add_subplot(111)
    p1, = ax.plot([1, 2, 3], '-o')
    p2, = ax.plot([2, 3, 4], '-x')
    p3, = ax.plot([3, 4, 5], '-d')
    ax.legend([(p1, p2), (p2, p1), p3], ['two keys', 'pad=0', 'one key'],
              numpoints=1,
              handler_map={(p1, p2): HandlerTuple(ndivide=None),
                           (p2, p1): HandlerTuple(ndivide=None, pad=0)})


@image_comparison(baseline_images=['rgba_alpha'],
                  extensions=['png'], remove_text=True)
def test_alpha_rgba():
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(1, 1)
    ax.plot(range(10), lw=5)
    leg = plt.legend(['Longlabel that will go away'], loc=10)
    leg.legendPatch.set_facecolor([1, 0, 0, 0.5])


@image_comparison(baseline_images=['rcparam_alpha'],
                  extensions=['png'], remove_text=True)
def test_alpha_rcparam():
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(1, 1)
    ax.plot(range(10), lw=5)
    with mpl.rc_context(rc={'legend.framealpha': .75}):
        leg = plt.legend(['Longlabel that will go away'], loc=10)
        # this alpha is going to be over-ridden by the rcparam with
        # sets the alpha of the patch to be non-None which causes the alpha
        # value of the face color to be discarded.  This behavior may not be
        # ideal, but it is what it is and we should keep track of it changing
        leg.legendPatch.set_facecolor([1, 0, 0, 0.5])


@image_comparison(baseline_images=['fancy'], remove_text=True)
def test_fancy():
    # using subplot triggers some offsetbox functionality untested elsewhere
    plt.subplot(121)
    plt.scatter(np.arange(10), np.arange(10, 0, -1), label='XX\nXX')
    plt.plot([5] * 10, 'o--', label='XX')
    plt.errorbar(np.arange(10), np.arange(10), xerr=0.5,
                 yerr=0.5, label='XX')
    plt.legend(loc="center left", bbox_to_anchor=[1.0, 0.5],
               ncol=2, shadow=True, title="My legend", numpoints=1)


@image_comparison(baseline_images=['framealpha'], remove_text=True)
def test_framealpha():
    x = np.linspace(1, 100, 100)
    y = x
    plt.plot(x, y, label='mylabel', lw=10)
    plt.legend(framealpha=0.5)


@image_comparison(baseline_images=['scatter_rc3', 'scatter_rc1'],
                  remove_text=True)
def test_rc():
    # using subplot triggers some offsetbox functionality untested elsewhere
    plt.figure()
    ax = plt.subplot(121)
    ax.scatter(np.arange(10), np.arange(10, 0, -1), label='three')
    ax.legend(loc="center left", bbox_to_anchor=[1.0, 0.5],
              title="My legend")

    mpl.rcParams['legend.scatterpoints'] = 1
    plt.figure()
    ax = plt.subplot(121)
    ax.scatter(np.arange(10), np.arange(10, 0, -1), label='one')
    ax.legend(loc="center left", bbox_to_anchor=[1.0, 0.5],
              title="My legend")


@image_comparison(baseline_images=['legend_expand'], remove_text=True)
def test_legend_expand():
    'Test expand mode'
    legend_modes = [None, "expand"]
    fig, axes_list = plt.subplots(len(legend_modes), 1)
    x = np.arange(100)
    for ax, mode in zip(axes_list, legend_modes):
        ax.plot(x, 50 - x, 'o', label='y=1')
        l1 = ax.legend(loc=2, mode=mode)
        ax.add_artist(l1)
        ax.plot(x, x - 50, 'o', label='y=-1')
        l2 = ax.legend(loc=5, mode=mode)
        ax.add_artist(l2)
        ax.legend(loc=3, mode=mode, ncol=2)


@image_comparison(baseline_images=['hatching'], remove_text=True,
                  style='default')
def test_hatching():
    fig, ax = plt.subplots()

    # Patches
    patch = plt.Rectangle((0, 0), 0.3, 0.3, hatch='xx',
                          label='Patch\ndefault color\nfilled')
    ax.add_patch(patch)
    patch = plt.Rectangle((0.33, 0), 0.3, 0.3, hatch='||', edgecolor='C1',
                          label='Patch\nexplicit color\nfilled')
    ax.add_patch(patch)
    patch = plt.Rectangle((0, 0.4), 0.3, 0.3, hatch='xx', fill=False,
                          label='Patch\ndefault color\nunfilled')
    ax.add_patch(patch)
    patch = plt.Rectangle((0.33, 0.4), 0.3, 0.3, hatch='||', fill=False,
                          edgecolor='C1',
                          label='Patch\nexplicit color\nunfilled')
    ax.add_patch(patch)

    # Paths
    ax.fill_between([0, .15, .3], [.8, .8, .8], [.9, 1.0, .9],
                    hatch='+', label='Path\ndefault color')
    ax.fill_between([.33, .48, .63], [.8, .8, .8], [.9, 1.0, .9],
                    hatch='+', edgecolor='C2', label='Path\nexplicit color')

    ax.set_xlim(-0.01, 1.1)
    ax.set_ylim(-0.01, 1.1)
    ax.legend(handlelength=4, handleheight=4)


def test_legend_remove():
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    lines = ax.plot(range(10))
    leg = fig.legend(lines, "test")
    leg.remove()
    assert fig.legends == []
    leg = ax.legend("test")
    leg.remove()
    assert ax.get_legend() is None


class TestLegendFunction(object):
    # Tests the legend function on the Axes and pyplot.
    def test_legend_handle_label(self):
        lines = plt.plot(range(10))
        with mock.patch('matplotlib.legend.Legend') as Legend:
            plt.legend(lines, ['hello world'])
        Legend.assert_called_with(plt.gca(), lines, ['hello world'])

    def test_legend_no_args(self):
        lines = plt.plot(range(10), label='hello world')
        with mock.patch('matplotlib.legend.Legend') as Legend:
            plt.legend()
        Legend.assert_called_with(plt.gca(), lines, ['hello world'])

    def test_legend_label_args(self):
        lines = plt.plot(range(10), label='hello world')
        with mock.patch('matplotlib.legend.Legend') as Legend:
            plt.legend(['foobar'])
        Legend.assert_called_with(plt.gca(), lines, ['foobar'])

    def test_legend_three_args(self):
        lines = plt.plot(range(10), label='hello world')
        with mock.patch('matplotlib.legend.Legend') as Legend:
            plt.legend(lines, ['foobar'], loc='right')
        Legend.assert_called_with(plt.gca(), lines, ['foobar'], loc='right')

    def test_legend_handler_map(self):
        lines = plt.plot(range(10), label='hello world')
        with mock.patch('matplotlib.legend.'
                        '_get_legend_handles_labels') as handles_labels:
            handles_labels.return_value = lines, ['hello world']
            plt.legend(handler_map={'1': 2})
        handles_labels.assert_called_with([plt.gca()], {'1': 2})

    def test_kwargs(self):
        fig, ax = plt.subplots(1, 1)
        th = np.linspace(0, 2*np.pi, 1024)
        lns, = ax.plot(th, np.sin(th), label='sin', lw=5)
        lnc, = ax.plot(th, np.cos(th), label='cos', lw=5)
        with mock.patch('matplotlib.legend.Legend') as Legend:
            ax.legend(labels=('a', 'b'), handles=(lnc, lns))
        Legend.assert_called_with(ax, (lnc, lns), ('a', 'b'))

    def test_warn_args_kwargs(self):
        fig, ax = plt.subplots(1, 1)
        th = np.linspace(0, 2*np.pi, 1024)
        lns, = ax.plot(th, np.sin(th), label='sin', lw=5)
        lnc, = ax.plot(th, np.cos(th), label='cos', lw=5)
        with mock.patch('warnings.warn') as warn:
            ax.legend((lnc, lns), labels=('a', 'b'))

        warn.assert_called_with("You have mixed positional and keyword "
                                "arguments, some input may be "
                                "discarded.")

    def test_parasite(self):
        from mpl_toolkits.axes_grid1 import host_subplot

        host = host_subplot(111)
        par = host.twinx()

        p1, = host.plot([0, 1, 2], [0, 1, 2], label="Density")
        p2, = par.plot([0, 1, 2], [0, 3, 2], label="Temperature")

        with mock.patch('matplotlib.legend.Legend') as Legend:
            leg = plt.legend()
        Legend.assert_called_with(host, [p1, p2],
                ['Density', 'Temperature'])


class TestLegendFigureFunction(object):
    # Tests the legend function for figure
    def test_legend_handle_label(self):
        fig, ax = plt.subplots()
        lines = ax.plot(range(10))
        with mock.patch('matplotlib.legend.Legend') as Legend:
            fig.legend(lines, ['hello world'])
        Legend.assert_called_with(fig, lines, ['hello world'])

    def test_legend_no_args(self):
        fig, ax = plt.subplots()
        lines = ax.plot(range(10), label='hello world')
        with mock.patch('matplotlib.legend.Legend') as Legend:
            fig.legend()
        Legend.assert_called_with(fig, lines, ['hello world'])

    def test_legend_label_arg(self):
        fig, ax = plt.subplots()
        lines = ax.plot(range(10))
        with mock.patch('matplotlib.legend.Legend') as Legend:
            fig.legend(['foobar'])
        Legend.assert_called_with(fig, lines, ['foobar'])

    def test_legend_label_three_args(self):
        fig, ax = plt.subplots()
        lines = ax.plot(range(10))
        with mock.patch('matplotlib.legend.Legend') as Legend:
            fig.legend(lines, ['foobar'], 'right')
        Legend.assert_called_with(fig, lines, ['foobar'], 'right')

    def test_legend_label_three_args_pluskw(self):
        # test that third argument and loc=  called together give
        # Exception
        fig, ax = plt.subplots()
        lines = ax.plot(range(10))
        with pytest.raises(Exception):
            fig.legend(lines, ['foobar'], 'right', loc='left')

    def test_legend_kw_args(self):
        fig, axs = plt.subplots(1, 2)
        lines = axs[0].plot(range(10))
        lines2 = axs[1].plot(np.arange(10) * 2.)
        with mock.patch('matplotlib.legend.Legend') as Legend:
            fig.legend(loc='right', labels=('a', 'b'),
                    handles=(lines, lines2))
        Legend.assert_called_with(fig, (lines, lines2), ('a', 'b'),
                loc='right')

    def test_warn_args_kwargs(self):
        fig, axs = plt.subplots(1, 2)
        lines = axs[0].plot(range(10))
        lines2 = axs[1].plot(np.arange(10) * 2.)
        with mock.patch('warnings.warn') as warn:
            fig.legend((lines, lines2), labels=('a', 'b'))
        warn.assert_called_with("You have mixed positional and keyword "
                                "arguments, some input may be "
                                "discarded.")


@image_comparison(baseline_images=['legend_stackplot'], extensions=['png'])
def test_legend_stackplot():
    '''test legend for PolyCollection using stackplot'''
    # related to #1341, #1943, and PR #3303
    fig = plt.figure()
    ax = fig.add_subplot(111)
    x = np.linspace(0, 10, 10)
    y1 = 1.0 * x
    y2 = 2.0 * x + 1
    y3 = 3.0 * x + 2
    ax.stackplot(x, y1, y2, y3, labels=['y1', 'y2', 'y3'])
    ax.set_xlim((0, 10))
    ax.set_ylim((0, 70))
    ax.legend(loc=0)


def test_cross_figure_patch_legend():
    fig, ax = plt.subplots()
    fig2, ax2 = plt.subplots()

    brs = ax.bar(range(3), range(3))
    fig2.legend(brs, 'foo')


def test_nanscatter():
    fig, ax = plt.subplots()

    h = ax.scatter([np.nan], [np.nan], marker="o",
                   facecolor="r", edgecolor="r", s=3)

    ax.legend([h], ["scatter"])

    fig, ax = plt.subplots()
    for color in ['red', 'green', 'blue']:
        n = 750
        x, y = np.random.rand(2, n)
        scale = 200.0 * np.random.rand(n)
        ax.scatter(x, y, c=color, s=scale, label=color,
                   alpha=0.3, edgecolors='none')

    ax.legend()
    ax.grid(True)


def test_legend_repeatcheckok():
    fig, ax = plt.subplots()
    ax.scatter(0.0, 1.0, color='k', marker='o', label='test')
    ax.scatter(0.5, 0.0, color='r', marker='v', label='test')
    hl = ax.legend()
    hand, lab = mlegend._get_legend_handles_labels([ax])
    assert len(lab) == 2
    fig, ax = plt.subplots()
    ax.scatter(0.0, 1.0, color='k', marker='o', label='test')
    ax.scatter(0.5, 0.0, color='k', marker='v', label='test')
    hl = ax.legend()
    hand, lab = mlegend._get_legend_handles_labels([ax])
    assert len(lab) == 2


@image_comparison(baseline_images=['not_covering_scatter'], extensions=['png'])
def test_not_covering_scatter():
    colors = ['b', 'g', 'r']

    for n in range(3):
        plt.scatter([n], [n], color=colors[n])

    plt.legend(['foo', 'foo', 'foo'], loc='best')
    plt.gca().set_xlim(-0.5, 2.2)
    plt.gca().set_ylim(-0.5, 2.2)


@image_comparison(baseline_images=['not_covering_scatter_transform'],
                  extensions=['png'])
def test_not_covering_scatter_transform():
    # Offsets point to top left, the default auto position
    offset = mtransforms.Affine2D().translate(-20, 20)
    x = np.linspace(0, 30, 1000)
    plt.plot(x, x)

    plt.scatter([20], [10], transform=offset + plt.gca().transData)

    plt.legend(['foo', 'bar'], loc='best')


def test_linecollection_scaled_dashes():
    lines1 = [[(0, .5), (.5, 1)], [(.3, .6), (.2, .2)]]
    lines2 = [[[0.7, .2], [.8, .4]], [[.5, .7], [.6, .1]]]
    lines3 = [[[0.6, .2], [.8, .4]], [[.5, .7], [.1, .1]]]
    lc1 = mcollections.LineCollection(lines1, linestyles="--", lw=3)
    lc2 = mcollections.LineCollection(lines2, linestyles="-.")
    lc3 = mcollections.LineCollection(lines3, linestyles=":", lw=.5)

    fig, ax = plt.subplots()
    ax.add_collection(lc1)
    ax.add_collection(lc2)
    ax.add_collection(lc3)

    leg = ax.legend([lc1, lc2, lc3], ["line1", "line2", 'line 3'])
    h1, h2, h3 = leg.legendHandles

    for oh, lh in zip((lc1, lc2, lc3), (h1, h2, h3)):
        assert oh.get_linestyles()[0][1] == lh._dashSeq
        assert oh.get_linestyles()[0][0] == lh._dashOffset


def test_handler_numpoints():
    '''test legend handler with numponts less than or equal to 1'''
    # related to #6921 and PR #8478
    fig, ax = plt.subplots()
    ax.plot(range(5), label='test')
    ax.legend(numpoints=0.5)


def test_shadow_framealpha():
    # Test if framealpha is activated when shadow is True
    # and framealpha is not explicitly passed'''
    fig, ax = plt.subplots()
    ax.plot(range(100), label="test")
    leg = ax.legend(shadow=True, facecolor='w')
    assert leg.get_frame().get_alpha() == 1

social_vdgnn/lib/python3.6/site-packages/matplotlib/tests/test_legend.py
eliatrevisan/IA-MPPI-LBM
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six

import tempfile
import warnings

from numpy.testing import (assert_allclose, assert_almost_equal,
                           assert_array_equal, assert_array_almost_equal_nulp)
import numpy.ma.testutils as matest
import numpy as np
import datetime as datetime
import pytest

import matplotlib.mlab as mlab
import matplotlib.cbook as cbook


try:
    from mpl_toolkits.natgrid import _natgrid
    HAS_NATGRID = True
except ImportError:
    HAS_NATGRID = False


def test_colinear_pca():
    a = mlab.PCA._get_colinear()
    pca = mlab.PCA(a)

    assert_allclose(pca.fracs[2:], 0., atol=1e-8)
    assert_allclose(pca.Y[:, 2:], 0., atol=1e-8)


@pytest.mark.parametrize('input', [
    # test odd lengths
    [1, 2, 3],
    # test even lengths
    [1, 2, 3, 4],
    # derived from email sent by jason-sage to MPL-user on 20090914
    [1, 1, 2, 2, 1, 2, 4, 3, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9, 7, 6, 4, 5, 5],
],
ids=[
    'odd length',
    'even length',
    'custom data',
])
@pytest.mark.parametrize('percentile', [
    0,
    50,
    75,
    100,
    [0, 75, 100],
])
def test_prctile(input, percentile):
    assert_allclose(mlab.prctile(input, percentile),
                    np.percentile(input, percentile))


def test_norm():
    np.random.seed(0)
    N = 1000
    x = np.random.standard_normal(N)
    targ = np.linalg.norm(x)
    res = mlab._norm(x)
    assert_almost_equal(targ, res)


@pytest.mark.parametrize('xmin, xmax, N', [
    (.01, 1000., 6),
    (.03, 1313., 7),
    (.03, 1313., 0),
    (.03, 1313., 1),
], ids=[
    'tens',
    'primes',
    'none',
    'single',
])
def test_logspace(xmin, xmax, N):
    res = mlab.logspace(xmin, xmax, N)
    targ = np.logspace(np.log10(xmin), np.log10(xmax), N)
    assert_allclose(targ, res)
    assert res.size == N


class TestStride(object):
    def get_base(self, x):
        y = x
        while y.base is not None:
            y = y.base
        return y

    def calc_window_target(self, x, NFFT, noverlap=0, axis=0):
        '''This is an adaptation of the original window extraction
        algorithm.  This is here to test to make sure the new implementation
        has the same result'''
        step = NFFT - noverlap
        ind = np.arange(0, len(x) - NFFT + 1, step)
        n = len(ind)
        result = np.zeros((NFFT, n))

        # do the ffts of the slices
        for i in range(n):
            result[:, i] = x[ind[i]:ind[i]+NFFT]
        if axis == 1:
            result = result.T
        return result

    @pytest.mark.parametrize('shape', [(), (10, 1)], ids=['0D', '2D'])
    def test_stride_windows_invalid_input_shape(self, shape):
        x = np.arange(np.prod(shape)).reshape(shape)
        with pytest.raises(ValueError):
            mlab.stride_windows(x, 5)

    @pytest.mark.parametrize('n, noverlap',
                             [(0, None), (11, None), (2, 2), (2, 3)],
                             ids=['n less than 1', 'n greater than input',
                                  'noverlap greater than n',
                                  'noverlap equal to n'])
    def test_stride_windows_invalid_params(self, n, noverlap):
        x = np.arange(10)
        with pytest.raises(ValueError):
            mlab.stride_windows(x, n, noverlap)

    @pytest.mark.parametrize('shape', [(), (10, 1)], ids=['0D', '2D'])
    def test_stride_repeat_invalid_input_shape(self, shape):
        x = np.arange(np.prod(shape)).reshape(shape)
        with pytest.raises(ValueError):
            mlab.stride_repeat(x, 5)

    @pytest.mark.parametrize('axis', [-1, 2],
                             ids=['axis less than 0',
                                  'axis greater than input shape'])
    def test_stride_repeat_invalid_axis(self, axis):
        x = np.array(0)
        with pytest.raises(ValueError):
            mlab.stride_repeat(x, 5, axis=axis)

    def test_stride_repeat_n_lt_1_ValueError(self):
        x = np.arange(10)
        with pytest.raises(ValueError):
            mlab.stride_repeat(x, 0)

    @pytest.mark.parametrize('axis', [0, 1], ids=['axis0', 'axis1'])
    @pytest.mark.parametrize('n', [1, 5], ids=['n1', 'n5'])
    def test_stride_repeat(self, n, axis):
        x = np.arange(10)
        y = mlab.stride_repeat(x, n, axis=axis)

        expected_shape = [10, 10]
        expected_shape[axis] = n
        yr = np.repeat(np.expand_dims(x, axis), n, axis=axis)

        assert yr.shape == y.shape
        assert_array_equal(yr, y)
        assert tuple(expected_shape) == y.shape
        assert self.get_base(y) is x

    @pytest.mark.parametrize('axis', [0, 1], ids=['axis0', 'axis1'])
    @pytest.mark.parametrize('n, noverlap',
                             [(1, 0), (5, 0), (15, 2), (13, -3)],
                             ids=['n1-noverlap0', 'n5-noverlap0',
                                  'n15-noverlap2', 'n13-noverlapn3'])
    def test_stride_windows(self, n, noverlap, axis):
        x = np.arange(100)
        y = mlab.stride_windows(x, n, noverlap=noverlap, axis=axis)

        expected_shape = [0, 0]
        expected_shape[axis] = n
        expected_shape[1 - axis] = 100 // (n - noverlap)
        yt = self.calc_window_target(x, n, noverlap=noverlap, axis=axis)

        assert yt.shape == y.shape
        assert_array_equal(yt, y)
        assert tuple(expected_shape) == y.shape
        assert self.get_base(y) is x

    @pytest.mark.parametrize('axis', [0, 1], ids=['axis0', 'axis1'])
    def test_stride_windows_n32_noverlap0_unflatten(self, axis):
        n = 32
        x = np.arange(n)[np.newaxis]
        x1 = np.tile(x, (21, 1))
        x2 = x1.flatten()
        y = mlab.stride_windows(x2, n, axis=axis)

        if axis == 0:
            x1 = x1.T
        assert y.shape == x1.shape
        assert_array_equal(y, x1)

    def test_stride_ensure_integer_type(self):
        N = 100
        x = np.empty(N + 20, dtype='>f4')
        x.fill(np.NaN)
        y = x[10:-10]
        y.fill(0.3)
        # previous to #3845 lead to corrupt access
        y_strided = mlab.stride_windows(y, n=33, noverlap=0.6)
        assert_array_equal(y_strided, 0.3)
        # previous to #3845 lead to corrupt access
        y_strided = mlab.stride_windows(y, n=33.3, noverlap=0)
        assert_array_equal(y_strided, 0.3)
        # even previous to #3845 could not find any problematic
        # configuration however, let's be sure it's not accidentally
        # introduced
        y_strided = mlab.stride_repeat(y, n=33.815)
        assert_array_equal(y_strided, 0.3)


@pytest.fixture
def tempcsv():
    if six.PY2:
        fd = tempfile.TemporaryFile(suffix='csv', mode="wb+")
    else:
        fd = tempfile.TemporaryFile(suffix='csv', mode="w+", newline='')
    with fd:
        yield fd


def test_recarray_csv_roundtrip(tempcsv):
    expected = np.recarray((99,),
                           [(str('x'), float),
                            (str('y'), float),
                            (str('t'), float)])
    # initialising all values: uninitialised memory sometimes produces
    # floats that do not round-trip to string and back.
    expected['x'][:] = np.linspace(-1e9, -1, 99)
    expected['y'][:] = np.linspace(1, 1e9, 99)
    expected['t'][:] = np.linspace(0, 0.01, 99)

    mlab.rec2csv(expected, tempcsv)
    tempcsv.seek(0)
    actual = mlab.csv2rec(tempcsv)

    assert_allclose(expected['x'], actual['x'])
    assert_allclose(expected['y'], actual['y'])
    assert_allclose(expected['t'], actual['t'])


def test_rec2csv_bad_shape_ValueError(tempcsv):
    bad = np.recarray((99, 4), [(str('x'), float),
                                (str('y'), float)])

    # the bad recarray should trigger a ValueError for having ndim > 1.
    with pytest.raises(ValueError):
        mlab.rec2csv(bad, tempcsv)


def test_csv2rec_names_with_comments(tempcsv):
    tempcsv.write('# comment\n1,2,3\n4,5,6\n')
    tempcsv.seek(0)
    array = mlab.csv2rec(tempcsv, names='a,b,c')
    assert len(array) == 2
    assert len(array.dtype) == 3


@pytest.mark.parametrize('input, kwargs', [
    ('01/11/14\n'
     '03/05/76 12:00:01 AM\n'
     '07/09/83 5:17:34 PM\n'
     '06/20/2054 2:31:45 PM\n'
     '10/31/00 11:50:23 AM\n',
     {}),
    ('11/01/14\n'
     '05/03/76 12:00:01 AM\n'
     '09/07/83 5:17:34 PM\n'
     '20/06/2054 2:31:45 PM\n'
     '31/10/00 11:50:23 AM\n',
     {'dayfirst': True}),
    ('14/01/11\n'
     '76/03/05 12:00:01 AM\n'
     '83/07/09 5:17:34 PM\n'
     '2054/06/20 2:31:45 PM\n'
     '00/10/31 11:50:23 AM\n',
     {'yearfirst': True}),
], ids=['usdate', 'dayfirst', 'yearfirst'])
def test_csv2rec_dates(tempcsv, input, kwargs):
    tempcsv.write(input)
    expected = [datetime.datetime(2014, 1, 11, 0, 0),
                datetime.datetime(1976, 3, 5, 0, 0, 1),
                datetime.datetime(1983, 7, 9, 17, 17, 34),
                datetime.datetime(2054, 6, 20, 14, 31, 45),
                datetime.datetime(2000, 10, 31, 11, 50, 23)]
    tempcsv.seek(0)
    array = mlab.csv2rec(tempcsv, names='a', **kwargs)
    assert_array_equal(array['a'].tolist(), expected)


def test_rec2txt_basic():
    # str() calls around field names necessary b/c as of numpy 1.11
    # dtype doesn't like unicode names (caused by unicode_literals import)
    a = np.array([(1.0, 2, 'foo', 'bing'),
                  (2.0, 3, 'bar', 'blah')],
                 dtype=np.dtype([(str('x'), np.float32),
                                 (str('y'), np.int8),
                                 (str('s'), str, 3),
                                 (str('s2'), str, 4)]))
    truth = ('       x   y   s     s2\n'
             '   1.000   2   foo   bing   \n'
             '   2.000   3   bar   blah   ').splitlines()
    assert mlab.rec2txt(a).splitlines() == truth


class TestWindow(object):
    def setup(self):
        np.random.seed(0)
        n = 1000

        self.sig_rand = np.random.standard_normal(n) + 100.
        self.sig_ones = np.ones(n)

    def check_window_apply_repeat(self, x, window, NFFT, noverlap):
        '''This is an adaptation of the original window application
        algorithm.  This is here to test to make sure the new implementation
        has the same result'''
        step = NFFT - noverlap
        ind = np.arange(0, len(x) - NFFT + 1, step)
        n = len(ind)
        result = np.zeros((NFFT, n))

        if cbook.iterable(window):
            windowVals = window
        else:
            windowVals = window(np.ones((NFFT,), x.dtype))

        # do the ffts of the slices
        for i in range(n):
            result[:, i] = windowVals * x[ind[i]:ind[i]+NFFT]
        return result

    def test_window_none_rand(self):
        res = mlab.window_none(self.sig_ones)
        assert_array_equal(res, self.sig_ones)

    def test_window_none_ones(self):
        res = mlab.window_none(self.sig_rand)
        assert_array_equal(res, self.sig_rand)

    def test_window_hanning_rand(self):
        targ = np.hanning(len(self.sig_rand)) * self.sig_rand
        res = mlab.window_hanning(self.sig_rand)

        assert_allclose(targ, res, atol=1e-06)

    def test_window_hanning_ones(self):
        targ = np.hanning(len(self.sig_ones))
        res = mlab.window_hanning(self.sig_ones)

        assert_allclose(targ, res, atol=1e-06)

    def test_apply_window_1D_axis1_ValueError(self):
        x = self.sig_rand
        window = mlab.window_hanning
        with pytest.raises(ValueError):
            mlab.apply_window(x, window, axis=1, return_window=False)

    def test_apply_window_1D_els_wrongsize_ValueError(self):
        x = self.sig_rand
        window = mlab.window_hanning(np.ones(x.shape[0]-1))
        with pytest.raises(ValueError):
            mlab.apply_window(x, window)

    def test_apply_window_0D_ValueError(self):
        x = np.array(0)
        window = mlab.window_hanning
        with pytest.raises(ValueError):
            mlab.apply_window(x, window, axis=1, return_window=False)

    def test_apply_window_3D_ValueError(self):
        x = self.sig_rand[np.newaxis][np.newaxis]
        window = mlab.window_hanning
        with pytest.raises(ValueError):
            mlab.apply_window(x, window, axis=1, return_window=False)

    def test_apply_window_hanning_1D(self):
        x = self.sig_rand
        window = mlab.window_hanning
        window1 = mlab.window_hanning(np.ones(x.shape[0]))
        y, window2 = mlab.apply_window(x, window, return_window=True)
        yt = window(x)
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)
        assert_array_equal(window1, window2)

    def test_apply_window_hanning_1D_axis0(self):
        x = self.sig_rand
        window = mlab.window_hanning
        y = mlab.apply_window(x, window, axis=0, return_window=False)
        yt = window(x)
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)

    def test_apply_window_hanning_els_1D_axis0(self):
        x = self.sig_rand
        window = mlab.window_hanning(np.ones(x.shape[0]))
        window1 = mlab.window_hanning
        y = mlab.apply_window(x, window, axis=0, return_window=False)
        yt = window1(x)
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)

    def test_apply_window_hanning_2D_axis0(self):
        x = np.random.standard_normal([1000, 10]) + 100.
        window = mlab.window_hanning
        y = mlab.apply_window(x, window, axis=0, return_window=False)
        yt = np.zeros_like(x)
        for i in range(x.shape[1]):
            yt[:, i] = window(x[:, i])
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)

    def test_apply_window_hanning_els1_2D_axis0(self):
        x = np.random.standard_normal([1000, 10]) + 100.
        window = mlab.window_hanning(np.ones(x.shape[0]))
        window1 = mlab.window_hanning
        y = mlab.apply_window(x, window, axis=0, return_window=False)
        yt = np.zeros_like(x)
        for i in range(x.shape[1]):
            yt[:, i] = window1(x[:, i])
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)

    def test_apply_window_hanning_els2_2D_axis0(self):
        x = np.random.standard_normal([1000, 10]) + 100.
        window = mlab.window_hanning
        window1 = mlab.window_hanning(np.ones(x.shape[0]))
        y, window2 = mlab.apply_window(x, window, axis=0, return_window=True)
        yt = np.zeros_like(x)
        for i in range(x.shape[1]):
            yt[:, i] = window1*x[:, i]
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)
        assert_array_equal(window1, window2)

    def test_apply_window_hanning_els3_2D_axis0(self):
        x = np.random.standard_normal([1000, 10]) + 100.
        window = mlab.window_hanning
        window1 = mlab.window_hanning(np.ones(x.shape[0]))
        y, window2 = mlab.apply_window(x, window, axis=0, return_window=True)
        yt = mlab.apply_window(x, window1, axis=0, return_window=False)
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)
        assert_array_equal(window1, window2)

    def test_apply_window_hanning_2D_axis1(self):
        x = np.random.standard_normal([10, 1000]) + 100.
        window = mlab.window_hanning
        y = mlab.apply_window(x, window, axis=1, return_window=False)
        yt = np.zeros_like(x)
        for i in range(x.shape[0]):
            yt[i, :] = window(x[i, :])
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)

    def test_apply_window_hanning_2D__els1_axis1(self):
        x = np.random.standard_normal([10, 1000]) + 100.
        window = mlab.window_hanning(np.ones(x.shape[1]))
        window1 = mlab.window_hanning
        y = mlab.apply_window(x, window, axis=1, return_window=False)
        yt = np.zeros_like(x)
        for i in range(x.shape[0]):
            yt[i, :] = window1(x[i, :])
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)

    def test_apply_window_hanning_2D_els2_axis1(self):
        x = np.random.standard_normal([10, 1000]) + 100.
        window = mlab.window_hanning
        window1 = mlab.window_hanning(np.ones(x.shape[1]))
        y, window2 = mlab.apply_window(x, window, axis=1, return_window=True)
        yt = np.zeros_like(x)
        for i in range(x.shape[0]):
            yt[i, :] = window1 * x[i, :]
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)
        assert_array_equal(window1, window2)

    def test_apply_window_hanning_2D_els3_axis1(self):
        x = np.random.standard_normal([10, 1000]) + 100.
        window = mlab.window_hanning
        window1 = mlab.window_hanning(np.ones(x.shape[1]))
        y = mlab.apply_window(x, window, axis=1, return_window=False)
        yt = mlab.apply_window(x, window1, axis=1, return_window=False)
        assert yt.shape == y.shape
        assert x.shape == y.shape
        assert_allclose(yt, y, atol=1e-06)

    def test_apply_window_stride_windows_hanning_2D_n13_noverlapn3_axis0(self):
        x = self.sig_rand
        window = mlab.window_hanning
        yi = mlab.stride_windows(x, n=13, noverlap=2, axis=0)
        y = mlab.apply_window(yi, window, axis=0, return_window=False)
        yt = self.check_window_apply_repeat(x, window, 13, 2)
        assert yt.shape == y.shape
        assert x.shape != y.shape
        assert_allclose(yt, y, atol=1e-06)

    def test_apply_window_hanning_2D_stack_axis1(self):
        ydata = np.arange(32)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ycontrol1 = mlab.apply_window(ydata1, mlab.window_hanning)
        ycontrol2 = mlab.window_hanning(ydata2)
        ydata = np.vstack([ydata1, ydata2])
        ycontrol = np.vstack([ycontrol1, ycontrol2])
        ydata = np.tile(ydata, (20, 1))
        ycontrol = np.tile(ycontrol, (20, 1))
        result = mlab.apply_window(ydata, mlab.window_hanning, axis=1,
                                   return_window=False)
        assert_allclose(ycontrol, result, atol=1e-08)

    def test_apply_window_hanning_2D_stack_windows_axis1(self):
        ydata = np.arange(32)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ycontrol1 = mlab.apply_window(ydata1, mlab.window_hanning)
        ycontrol2 = mlab.window_hanning(ydata2)
        ydata = np.vstack([ydata1, ydata2])
        ycontrol = np.vstack([ycontrol1, ycontrol2])
        ydata = np.tile(ydata, (20, 1))
        ycontrol = np.tile(ycontrol, (20, 1))
        result = mlab.apply_window(ydata, mlab.window_hanning, axis=1,
                                   return_window=False)
        assert_allclose(ycontrol, result, atol=1e-08)

    def test_apply_window_hanning_2D_stack_windows_axis1_unflatten(self):
        n = 32
        ydata = np.arange(n)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ycontrol1 = mlab.apply_window(ydata1, mlab.window_hanning)
        ycontrol2 = mlab.window_hanning(ydata2)
        ydata = np.vstack([ydata1, ydata2])
        ycontrol = np.vstack([ycontrol1, ycontrol2])
        ydata = np.tile(ydata, (20, 1))
        ycontrol = np.tile(ycontrol, (20, 1))
        ydata = ydata.flatten()
        ydata1 = mlab.stride_windows(ydata, 32, noverlap=0, axis=0)
        result = mlab.apply_window(ydata1, mlab.window_hanning, axis=0,
                                   return_window=False)
        assert_allclose(ycontrol.T, result, atol=1e-08)


class TestDetrend(object):
    def setup(self):
        np.random.seed(0)
        n = 1000
        x = np.linspace(0., 100, n)

        self.sig_zeros = np.zeros(n)

        self.sig_off = self.sig_zeros + 100.
        self.sig_slope = np.linspace(-10., 90., n)

        self.sig_slope_mean = x - x.mean()

        sig_rand = np.random.standard_normal(n)
        sig_sin = np.sin(x*2*np.pi/(n/100))

        sig_rand -= sig_rand.mean()
        sig_sin -= sig_sin.mean()

        self.sig_base = sig_rand + sig_sin

        self.atol = 1e-08

    def test_detrend_none_0D_zeros(self):
        input = 0.
        targ = input
        res = mlab.detrend_none(input)
        assert input == targ

    def test_detrend_none_0D_zeros_axis1(self):
        input = 0.
        targ = input
        res = mlab.detrend_none(input, axis=1)
        assert input == targ

    def test_detrend_str_none_0D_zeros(self):
        input = 0.
        targ = input
        res = mlab.detrend(input, key='none')
        assert input == targ

    def test_detrend_detrend_none_0D_zeros(self):
        input = 0.
        targ = input
        res = mlab.detrend(input, key=mlab.detrend_none)
        assert input == targ

    def test_detrend_none_0D_off(self):
        input = 5.5
        targ = input
        res = mlab.detrend_none(input)
        assert input == targ

    def test_detrend_none_1D_off(self):
        input = self.sig_off
        targ = input
        res = mlab.detrend_none(input)
        assert_array_equal(res, targ)

    def test_detrend_none_1D_slope(self):
        input = self.sig_slope
        targ = input
        res = mlab.detrend_none(input)
        assert_array_equal(res, targ)

    def test_detrend_none_1D_base(self):
        input = self.sig_base
        targ = input
        res = mlab.detrend_none(input)
        assert_array_equal(res, targ)

    def test_detrend_none_1D_base_slope_off_list(self):
        input = self.sig_base + self.sig_slope + self.sig_off
        targ = input.tolist()
        res = mlab.detrend_none(input.tolist())
        assert res == targ

    def test_detrend_none_2D(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        input = np.vstack(arri)
        targ = input
        res = mlab.detrend_none(input)
        assert_array_equal(res, targ)

    def test_detrend_none_2D_T(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        input = np.vstack(arri)
        targ = input
        res = mlab.detrend_none(input.T)
        assert_array_equal(res.T, targ)

    def test_detrend_mean_0D_zeros(self):
        input = 0.
        targ = 0.
        res = mlab.detrend_mean(input)
        assert_almost_equal(res, targ)

    def test_detrend_str_mean_0D_zeros(self):
        input = 0.
        targ = 0.
        res = mlab.detrend(input, key='mean')
        assert_almost_equal(res, targ)

    def test_detrend_detrend_mean_0D_zeros(self):
        input = 0.
        targ = 0.
        res = mlab.detrend(input, key=mlab.detrend_mean)
        assert_almost_equal(res, targ)

    def test_detrend_mean_0D_off(self):
        input = 5.5
        targ = 0.
        res = mlab.detrend_mean(input)
        assert_almost_equal(res, targ)

    def test_detrend_str_mean_0D_off(self):
        input = 5.5
        targ = 0.
        res = mlab.detrend(input, key='mean')
        assert_almost_equal(res, targ)

    def test_detrend_detrend_mean_0D_off(self):
        input = 5.5
        targ = 0.
        res = mlab.detrend(input, key=mlab.detrend_mean)
        assert_almost_equal(res, targ)

    def test_detrend_mean_1D_zeros(self):
        input = self.sig_zeros
        targ = self.sig_zeros
        res = mlab.detrend_mean(input)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_mean_1D_base(self):
        input = self.sig_base
        targ = self.sig_base
        res = mlab.detrend_mean(input)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_mean_1D_base_off(self):
        input = self.sig_base + self.sig_off
        targ = self.sig_base
        res = mlab.detrend_mean(input)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_mean_1D_base_slope(self):
        input = self.sig_base + self.sig_slope
        targ = self.sig_base + self.sig_slope_mean
        res = mlab.detrend_mean(input)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_mean_1D_base_slope_off(self):
        input = self.sig_base + self.sig_slope + self.sig_off
        targ = self.sig_base + self.sig_slope_mean
        res = mlab.detrend_mean(input)
        assert_allclose(res, targ, atol=1e-08)

    def test_detrend_mean_1D_base_slope_off_axis0(self):
        input = self.sig_base + self.sig_slope + self.sig_off
        targ = self.sig_base + self.sig_slope_mean
        res = mlab.detrend_mean(input, axis=0)
        assert_allclose(res, targ, atol=1e-08)

    def test_detrend_mean_1D_base_slope_off_list(self):
        input = self.sig_base + self.sig_slope + self.sig_off
        targ = self.sig_base + self.sig_slope_mean
        res = mlab.detrend_mean(input.tolist())
        assert_allclose(res, targ, atol=1e-08)

    def test_detrend_mean_1D_base_slope_off_list_axis0(self):
        input = self.sig_base + self.sig_slope + self.sig_off
        targ = self.sig_base + self.sig_slope_mean
        res = mlab.detrend_mean(input.tolist(), axis=0)
        assert_allclose(res, targ, atol=1e-08)

    def test_demean_0D_off(self):
        input = 5.5
        targ = 0.
        res = mlab.demean(input, axis=None)
        assert_almost_equal(res, targ)

    def test_demean_1D_base_slope_off(self):
        input = self.sig_base + self.sig_slope + self.sig_off
        targ = self.sig_base + self.sig_slope_mean
        res = mlab.demean(input)
        assert_allclose(res, targ, atol=1e-08)

    def test_demean_1D_base_slope_off_axis0(self):
        input = self.sig_base + self.sig_slope + self.sig_off
        targ = self.sig_base + self.sig_slope_mean
        res = mlab.demean(input, axis=0)
        assert_allclose(res, targ, atol=1e-08)

    def test_demean_1D_base_slope_off_list(self):
        input = self.sig_base + self.sig_slope + self.sig_off
        targ = self.sig_base + self.sig_slope_mean
        res = mlab.demean(input.tolist())
        assert_allclose(res, targ, atol=1e-08)

    def test_detrend_mean_2D_default(self):
        arri = [self.sig_off,
                self.sig_base + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_base]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend_mean(input)
        assert_allclose(res, targ, atol=1e-08)

    def test_detrend_mean_2D_none(self):
        arri = [self.sig_off,
                self.sig_base + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_base]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend_mean(input, axis=None)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_detrend_mean_2D_none_T(self):
        arri = [self.sig_off,
                self.sig_base + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_base]
        input = np.vstack(arri).T
        targ = np.vstack(arrt)
        res = mlab.detrend_mean(input, axis=None)
        assert_allclose(res.T, targ,
                        atol=1e-08)

    def test_detrend_mean_2D_axis0(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri).T
        targ = np.vstack(arrt).T
        res = mlab.detrend_mean(input, axis=0)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_detrend_mean_2D_axis1(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend_mean(input, axis=1)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_detrend_mean_2D_axism1(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend_mean(input, axis=-1)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_detrend_2D_default(self):
        arri = [self.sig_off,
                self.sig_base + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_base]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend(input)
        assert_allclose(res, targ, atol=1e-08)

    def test_detrend_2D_none(self):
        arri = [self.sig_off,
                self.sig_base + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_base]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend(input, axis=None)
        assert_allclose(res, targ, atol=1e-08)

    def test_detrend_str_mean_2D_axis0(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri).T
        targ = np.vstack(arrt).T
        res = mlab.detrend(input, key='mean', axis=0)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_detrend_str_constant_2D_none_T(self):
        arri = [self.sig_off,
                self.sig_base + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_base]
        input = np.vstack(arri).T
        targ = np.vstack(arrt)
        res = mlab.detrend(input, key='constant', axis=None)
        assert_allclose(res.T, targ,
                        atol=1e-08)

    def test_detrend_str_default_2D_axis1(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend(input, key='default', axis=1)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_detrend_detrend_mean_2D_axis0(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri).T
        targ = np.vstack(arrt).T
        res = mlab.detrend(input, key=mlab.detrend_mean, axis=0)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_demean_2D_default(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri).T
        targ = np.vstack(arrt).T
        res = mlab.demean(input)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_demean_2D_none(self):
        arri = [self.sig_off,
                self.sig_base + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_base]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.demean(input, axis=None)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_demean_2D_axis0(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri).T
        targ = np.vstack(arrt).T
        res = mlab.demean(input, axis=0)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_demean_2D_axis1(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.demean(input, axis=1)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_demean_2D_axism1(self):
        arri = [self.sig_base,
                self.sig_base + self.sig_off,
                self.sig_base + self.sig_slope,
                self.sig_base + self.sig_off + self.sig_slope]
        arrt = [self.sig_base,
                self.sig_base,
                self.sig_base + self.sig_slope_mean,
                self.sig_base + self.sig_slope_mean]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.demean(input, axis=-1)
        assert_allclose(res, targ,
                        atol=1e-08)

    def test_detrend_bad_key_str_ValueError(self):
        input = self.sig_slope[np.newaxis]
        with pytest.raises(ValueError):
            mlab.detrend(input, key='spam')

    def test_detrend_bad_key_var_ValueError(self):
        input = self.sig_slope[np.newaxis]
        with pytest.raises(ValueError):
            mlab.detrend(input, key=5)

    def test_detrend_mean_0D_d0_ValueError(self):
        input = 5.5
        with pytest.raises(ValueError):
            mlab.detrend_mean(input, axis=0)

    def test_detrend_0D_d0_ValueError(self):
        input = 5.5
        with pytest.raises(ValueError):
            mlab.detrend(input, axis=0)

    def test_detrend_mean_1D_d1_ValueError(self):
        input = self.sig_slope
        with pytest.raises(ValueError):
            mlab.detrend_mean(input, axis=1)

    def test_detrend_1D_d1_ValueError(self):
        input = self.sig_slope
        with pytest.raises(ValueError):
            mlab.detrend(input, axis=1)

    def test_demean_1D_d1_ValueError(self):
        input = self.sig_slope
        with pytest.raises(ValueError):
            mlab.demean(input, axis=1)

    def test_detrend_mean_2D_d2_ValueError(self):
        input = self.sig_slope[np.newaxis]
        with pytest.raises(ValueError):
            mlab.detrend_mean(input, axis=2)

    def test_detrend_2D_d2_ValueError(self):
        input = self.sig_slope[np.newaxis]
        with pytest.raises(ValueError):
            mlab.detrend(input, axis=2)

    def test_demean_2D_d2_ValueError(self):
        input = self.sig_slope[np.newaxis]
        with pytest.raises(ValueError):
            mlab.demean(input, axis=2)

    def test_detrend_linear_0D_zeros(self):
        input = 0.
        targ = 0.
        res = mlab.detrend_linear(input)
        assert_almost_equal(res, targ)

    def test_detrend_linear_0D_off(self):
        input = 5.5
        targ = 0.
        res = mlab.detrend_linear(input)
        assert_almost_equal(res, targ)

    def test_detrend_str_linear_0D_off(self):
        input = 5.5
        targ = 0.
        res = mlab.detrend(input, key='linear')
        assert_almost_equal(res, targ)

    def test_detrend_detrend_linear_0D_off(self):
        input = 5.5
        targ = 0.
        res = mlab.detrend(input, key=mlab.detrend_linear)
        assert_almost_equal(res, targ)

    def test_detrend_linear_1d_off(self):
        input = self.sig_off
        targ = self.sig_zeros
        res = mlab.detrend_linear(input)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_linear_1d_slope(self):
        input = self.sig_slope
        targ = self.sig_zeros
        res = mlab.detrend_linear(input)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_linear_1d_slope_off(self):
        input = self.sig_slope + self.sig_off
        targ = self.sig_zeros
        res = mlab.detrend_linear(input)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_str_linear_1d_slope_off(self):
        input = self.sig_slope + self.sig_off
        targ = self.sig_zeros
        res = mlab.detrend(input, key='linear')
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_detrend_linear_1d_slope_off(self):
        input = self.sig_slope + self.sig_off
        targ = self.sig_zeros
        res = mlab.detrend(input, key=mlab.detrend_linear)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_linear_1d_slope_off_list(self):
        input = self.sig_slope + self.sig_off
        targ = self.sig_zeros
        res = mlab.detrend_linear(input.tolist())
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_linear_2D_ValueError(self):
        input = self.sig_slope[np.newaxis]
        with pytest.raises(ValueError):
            mlab.detrend_linear(input)

    def test_detrend_str_linear_2d_slope_off_axis0(self):
        arri = [self.sig_off,
                self.sig_slope,
                self.sig_slope + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_zeros,
                self.sig_zeros]
        input = np.vstack(arri).T
        targ = np.vstack(arrt).T
        res = mlab.detrend(input, key='linear', axis=0)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_detrend_linear_1d_slope_off_axis1(self):
        arri = [self.sig_off,
                self.sig_slope,
                self.sig_slope + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_zeros,
                self.sig_zeros]
        input = np.vstack(arri).T
        targ = np.vstack(arrt).T
        res = mlab.detrend(input, key=mlab.detrend_linear, axis=0)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_str_linear_2d_slope_off_axis0(self):
        arri = [self.sig_off,
                self.sig_slope,
                self.sig_slope + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_zeros,
                self.sig_zeros]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend(input, key='linear', axis=1)
        assert_allclose(res, targ, atol=self.atol)

    def test_detrend_detrend_linear_1d_slope_off_axis1(self):
        arri = [self.sig_off,
                self.sig_slope,
                self.sig_slope + self.sig_off]
        arrt = [self.sig_zeros,
                self.sig_zeros,
                self.sig_zeros]
        input = np.vstack(arri)
        targ = np.vstack(arrt)
        res = mlab.detrend(input, key=mlab.detrend_linear, axis=1)
        assert_allclose(res, targ, atol=self.atol)


@pytest.mark.parametrize('iscomplex', [False, True],
                         ids=['real', 'complex'], scope='class')
@pytest.mark.parametrize('sides', ['onesided', 'twosided', 'default'],
                         scope='class')
@pytest.mark.parametrize(
    'fstims,len_x,NFFT_density,nover_density,pad_to_density,pad_to_spectrum',
    [
        ([], None, -1, -1, -1, -1),
        ([4], None, -1, -1, -1, -1),
        ([4, 5, 10], None, -1, -1, -1, -1),
        ([], None, None, -1, -1, None),
        ([], None, -1, -1, None, None),
        ([], None, None, -1, None, None),
        ([], 1024, 512, -1, -1, 128),
        ([], 256, -1, -1, 33, 257),
        ([], 255, 33, -1, -1, None),
        ([], 256, 128, -1, 256, 256),
        ([], None, -1, 32, -1, -1),
   ],
    ids=[
        'nosig',
        'Fs4',
        'FsAll',
        'nosig_noNFFT',
        'nosig_nopad_to',
        'nosig_noNFFT_no_pad_to',
        'nosig_trim',
        'nosig_odd',
        'nosig_oddlen',
        'nosig_stretch',
        'nosig_overlap',
    ],
    scope='class')
class TestSpectral(object):
    @pytest.fixture(scope='class', autouse=True)
    def stim(self, request, fstims, iscomplex, sides, len_x, NFFT_density,
             nover_density, pad_to_density, pad_to_spectrum):
        Fs = 100.

        x = np.arange(0, 10, 1 / Fs)
        if len_x is not None:
            x = x[:len_x]

        # get the stimulus frequencies, defaulting to None
        fstims = [Fs / fstim for fstim in fstims]

        # get the constants, default to calculated values
        if NFFT_density is None:
            NFFT_density_real = 256
        elif NFFT_density < 0:
            NFFT_density_real = NFFT_density = 100
        else:
            NFFT_density_real = NFFT_density

        if nover_density is None:
            nover_density_real = 0
        elif nover_density < 0:
            nover_density_real = nover_density = NFFT_density_real // 2
        else:
            nover_density_real = nover_density

        if pad_to_density is None:
            pad_to_density_real = NFFT_density_real
        elif pad_to_density < 0:
            pad_to_density = int(2**np.ceil(np.log2(NFFT_density_real)))
            pad_to_density_real = pad_to_density
        else:
            pad_to_density_real = pad_to_density

        if pad_to_spectrum is None:
            pad_to_spectrum_real = len(x)
        elif pad_to_spectrum < 0:
            pad_to_spectrum_real = pad_to_spectrum = len(x)
        else:
            pad_to_spectrum_real = pad_to_spectrum

        if pad_to_spectrum is None:
            NFFT_spectrum_real = NFFT_spectrum = pad_to_spectrum_real
        else:
            NFFT_spectrum_real = NFFT_spectrum = len(x)
        nover_spectrum_real = nover_spectrum = 0

        NFFT_specgram = NFFT_density
        nover_specgram = nover_density
        pad_to_specgram = pad_to_density
        NFFT_specgram_real = NFFT_density_real
        nover_specgram_real = nover_density_real

        if sides == 'onesided' or (sides == 'default' and not iscomplex):
            # frequencies for specgram, psd, and csd
            # need to handle even and odd differently
            if pad_to_density_real % 2:
                freqs_density = np.linspace(0, Fs / 2,
                                            num=pad_to_density_real,
                                            endpoint=False)[::2]
            else:
                freqs_density = np.linspace(0, Fs / 2,
                                            num=pad_to_density_real // 2 + 1)

            # frequencies for complex, magnitude, angle, and phase spectrums
            # need to handle even and odd differently
            if pad_to_spectrum_real % 2:
                freqs_spectrum = np.linspace(0, Fs / 2,
                                             num=pad_to_spectrum_real,
                                             endpoint=False)[::2]
            else:
                freqs_spectrum = np.linspace(0, Fs / 2,
                                             num=pad_to_spectrum_real // 2 + 1)
        else:
            # frequencies for specgram, psd, and csd
            # need to handle even and odd differentl
            if pad_to_density_real % 2:
                freqs_density = np.linspace(-Fs / 2, Fs / 2,
                                            num=2 * pad_to_density_real,
                                            endpoint=False)[1::2]
            else:
                freqs_density = np.linspace(-Fs / 2, Fs / 2,
                                            num=pad_to_density_real,
                                            endpoint=False)

            # frequencies for complex, magnitude, angle, and phase spectrums
            # need to handle even and odd differently
            if pad_to_spectrum_real % 2:
                freqs_spectrum = np.linspace(-Fs / 2, Fs / 2,
                                             num=2 * pad_to_spectrum_real,
                                             endpoint=False)[1::2]
            else:
                freqs_spectrum = np.linspace(-Fs / 2, Fs / 2,
                                             num=pad_to_spectrum_real,
                                             endpoint=False)

        freqs_specgram = freqs_density
        # time points for specgram
        t_start = NFFT_specgram_real // 2
        t_stop = len(x) - NFFT_specgram_real // 2 + 1
        t_step = NFFT_specgram_real - nover_specgram_real
        t_specgram = x[t_start:t_stop:t_step]
        if NFFT_specgram_real % 2:
            t_specgram += 1 / Fs / 2
        if len(t_specgram) == 0:
            t_specgram = np.array([NFFT_specgram_real / (2 * Fs)])
        t_spectrum = np.array([NFFT_spectrum_real / (2 * Fs)])
        t_density = t_specgram

        y = np.zeros_like(x)
        for i, fstim in enumerate(fstims):
            y += np.sin(fstim * x * np.pi * 2) * 10**i

        if iscomplex:
            y = y.astype('complex')

        # Interestingly, the instance on which this fixture is called is not
        # the same as the one on which a test is run. So we need to modify the
        # class itself when using a class-scoped fixture.
        cls = request.cls

        cls.Fs = Fs
        cls.sides = sides
        cls.fstims = fstims

        cls.NFFT_density = NFFT_density
        cls.nover_density = nover_density
        cls.pad_to_density = pad_to_density

        cls.NFFT_spectrum = NFFT_spectrum
        cls.nover_spectrum = nover_spectrum
        cls.pad_to_spectrum = pad_to_spectrum

        cls.NFFT_specgram = NFFT_specgram
        cls.nover_specgram = nover_specgram
        cls.pad_to_specgram = pad_to_specgram

        cls.t_specgram = t_specgram
        cls.t_density = t_density
        cls.t_spectrum = t_spectrum
        cls.y = y

        cls.freqs_density = freqs_density
        cls.freqs_spectrum = freqs_spectrum
        cls.freqs_specgram = freqs_specgram

        cls.NFFT_density_real = NFFT_density_real

    def check_freqs(self, vals, targfreqs, resfreqs, fstims):
        assert resfreqs.argmin() == 0
        assert resfreqs.argmax() == len(resfreqs)-1
        assert_allclose(resfreqs, targfreqs, atol=1e-06)
        for fstim in fstims:
            i = np.abs(resfreqs - fstim).argmin()
            assert vals[i] > vals[i+2]
            assert vals[i] > vals[i-2]

    def check_maxfreq(self, spec, fsp, fstims):
        # skip the test if there are no frequencies
        if len(fstims) == 0:
            return

        # if twosided, do the test for each side
        if fsp.min() < 0:
            fspa = np.abs(fsp)
            zeroind = fspa.argmin()
            self.check_maxfreq(spec[:zeroind], fspa[:zeroind], fstims)
            self.check_maxfreq(spec[zeroind:], fspa[zeroind:], fstims)
            return

        fstimst = fstims[:]
        spect = spec.copy()

        # go through each peak and make sure it is correctly the maximum peak
        while fstimst:
            maxind = spect.argmax()
            maxfreq = fsp[maxind]
            assert_almost_equal(maxfreq, fstimst[-1])
            del fstimst[-1]
            spect[maxind-5:maxind+5] = 0

    def test_spectral_helper_raises_complex_same_data(self):
        # test that mode 'complex' cannot be used if x is not y
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, y=self.y+1, mode='complex')

    def test_spectral_helper_raises_magnitude_same_data(self):
        # test that mode 'magnitude' cannot be used if x is not y
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, y=self.y+1, mode='magnitude')

    def test_spectral_helper_raises_angle_same_data(self):
        # test that mode 'angle' cannot be used if x is not y
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, y=self.y+1, mode='angle')

    def test_spectral_helper_raises_phase_same_data(self):
        # test that mode 'phase' cannot be used if x is not y
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, y=self.y+1, mode='phase')

    def test_spectral_helper_raises_unknown_mode(self):
        # test that unknown value for mode cannot be used
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, mode='spam')

    def test_spectral_helper_raises_unknown_sides(self):
        # test that unknown value for sides cannot be used
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, y=self.y, sides='eggs')

    def test_spectral_helper_raises_noverlap_gt_NFFT(self):
        # test that noverlap cannot be larger than NFFT
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, y=self.y, NFFT=10, noverlap=20)

    def test_spectral_helper_raises_noverlap_eq_NFFT(self):
        # test that noverlap cannot be equal to NFFT
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, NFFT=10, noverlap=10)

    def test_spectral_helper_raises_winlen_ne_NFFT(self):
        # test that the window length cannot be different from NFFT
        with pytest.raises(ValueError):
            mlab._spectral_helper(x=self.y, y=self.y, NFFT=10,
                                  window=np.ones(9))

    def test_single_spectrum_helper_raises_mode_default(self):
        # test that mode 'default' cannot be used with _single_spectrum_helper
        with pytest.raises(ValueError):
            mlab._single_spectrum_helper(x=self.y, mode='default')

    def test_single_spectrum_helper_raises_mode_psd(self):
        # test that mode 'psd' cannot be used with _single_spectrum_helper
        with pytest.raises(ValueError):
            mlab._single_spectrum_helper(x=self.y, mode='psd')

    def test_spectral_helper_psd(self):
        freqs = self.freqs_density
        spec, fsp, t = mlab._spectral_helper(x=self.y, y=self.y,
                                             NFFT=self.NFFT_density,
                                             Fs=self.Fs,
                                             noverlap=self.nover_density,
                                             pad_to=self.pad_to_density,
                                             sides=self.sides,
                                             mode='psd')

        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_density, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]

    def test_spectral_helper_magnitude_specgram(self):
        freqs = self.freqs_specgram
        spec, fsp, t = mlab._spectral_helper(x=self.y, y=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='magnitude')

        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_specgram, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]

    def test_spectral_helper_magnitude_magnitude_spectrum(self):
        freqs = self.freqs_spectrum
        spec, fsp, t = mlab._spectral_helper(x=self.y, y=self.y,
                                             NFFT=self.NFFT_spectrum,
                                             Fs=self.Fs,
                                             noverlap=self.nover_spectrum,
                                             pad_to=self.pad_to_spectrum,
                                             sides=self.sides,
                                             mode='magnitude')

        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_spectrum, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == 1

    def test_csd(self):
        freqs = self.freqs_density
        spec, fsp = mlab.csd(x=self.y, y=self.y+1,
                             NFFT=self.NFFT_density,
                             Fs=self.Fs,
                             noverlap=self.nover_density,
                             pad_to=self.pad_to_density,
                             sides=self.sides)
        assert_allclose(fsp, freqs, atol=1e-06)
        assert spec.shape == freqs.shape

    def test_csd_padding(self):
        """Test zero padding of csd(). """
        if self.NFFT_density is None:  # for derived classes
            return
        sargs = dict(x=self.y, y=self.y+1, Fs=self.Fs, window=mlab.window_none,
                     sides=self.sides)

        spec0, _ = mlab.csd(NFFT=self.NFFT_density, **sargs)
        spec1, _ = mlab.csd(NFFT=self.NFFT_density*2, **sargs)
        assert_almost_equal(np.sum(np.conjugate(spec0)*spec0).real,
                            np.sum(np.conjugate(spec1/2)*spec1/2).real)

    def test_psd(self):
        freqs = self.freqs_density
        spec, fsp = mlab.psd(x=self.y,
                             NFFT=self.NFFT_density,
                             Fs=self.Fs,
                             noverlap=self.nover_density,
                             pad_to=self.pad_to_density,
                             sides=self.sides)
        assert spec.shape == freqs.shape
        self.check_freqs(spec, freqs, fsp, self.fstims)

    def test_psd_detrend_mean_func_offset(self):
        if self.NFFT_density is None:
            return
        freqs = self.freqs_density
        ydata = np.zeros(self.NFFT_density)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ydata = np.vstack([ydata1, ydata2])
        ydata = np.tile(ydata, (20, 1))
        ydatab = ydata.T.flatten()
        ydata = ydata.flatten()
        ycontrol = np.zeros_like(ydata)
        spec_g, fsp_g = mlab.psd(x=ydata,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend=mlab.detrend_mean)
        spec_b, fsp_b = mlab.psd(x=ydatab,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend=mlab.detrend_mean)
        spec_c, fsp_c = mlab.psd(x=ycontrol,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides)
        assert_array_equal(fsp_g, fsp_c)
        assert_array_equal(fsp_b, fsp_c)
        assert_allclose(spec_g, spec_c, atol=1e-08)
        # these should not be almost equal
        with pytest.raises(AssertionError):
            assert_allclose(spec_b, spec_c, atol=1e-08)

    def test_psd_detrend_mean_str_offset(self):
        if self.NFFT_density is None:
            return
        freqs = self.freqs_density
        ydata = np.zeros(self.NFFT_density)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ydata = np.vstack([ydata1, ydata2])
        ydata = np.tile(ydata, (20, 1))
        ydatab = ydata.T.flatten()
        ydata = ydata.flatten()
        ycontrol = np.zeros_like(ydata)
        spec_g, fsp_g = mlab.psd(x=ydata,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend='mean')
        spec_b, fsp_b = mlab.psd(x=ydatab,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend='mean')
        spec_c, fsp_c = mlab.psd(x=ycontrol,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides)
        assert_array_equal(fsp_g, fsp_c)
        assert_array_equal(fsp_b, fsp_c)
        assert_allclose(spec_g, spec_c, atol=1e-08)
        # these should not be almost equal
        with pytest.raises(AssertionError):
            assert_allclose(spec_b, spec_c, atol=1e-08)

    def test_psd_detrend_linear_func_trend(self):
        if self.NFFT_density is None:
            return
        freqs = self.freqs_density
        ydata = np.arange(self.NFFT_density)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ydata = np.vstack([ydata1, ydata2])
        ydata = np.tile(ydata, (20, 1))
        ydatab = ydata.T.flatten()
        ydata = ydata.flatten()
        ycontrol = np.zeros_like(ydata)
        spec_g, fsp_g = mlab.psd(x=ydata,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend=mlab.detrend_linear)
        spec_b, fsp_b = mlab.psd(x=ydatab,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend=mlab.detrend_linear)
        spec_c, fsp_c = mlab.psd(x=ycontrol,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides)
        assert_array_equal(fsp_g, fsp_c)
        assert_array_equal(fsp_b, fsp_c)
        assert_allclose(spec_g, spec_c, atol=1e-08)
        # these should not be almost equal
        with pytest.raises(AssertionError):
            assert_allclose(spec_b, spec_c, atol=1e-08)

    def test_psd_detrend_linear_str_trend(self):
        if self.NFFT_density is None:
            return
        freqs = self.freqs_density
        ydata = np.arange(self.NFFT_density)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ydata = np.vstack([ydata1, ydata2])
        ydata = np.tile(ydata, (20, 1))
        ydatab = ydata.T.flatten()
        ydata = ydata.flatten()
        ycontrol = np.zeros_like(ydata)
        spec_g, fsp_g = mlab.psd(x=ydata,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend='linear')
        spec_b, fsp_b = mlab.psd(x=ydatab,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend='linear')
        spec_c, fsp_c = mlab.psd(x=ycontrol,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides)
        assert_array_equal(fsp_g, fsp_c)
        assert_array_equal(fsp_b, fsp_c)
        assert_allclose(spec_g, spec_c, atol=1e-08)
        # these should not be almost equal
        with pytest.raises(AssertionError):
            assert_allclose(spec_b, spec_c, atol=1e-08)

    def test_psd_window_hanning(self):
        if self.NFFT_density is None:
            return
        freqs = self.freqs_density
        ydata = np.arange(self.NFFT_density)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ycontrol1, windowVals = mlab.apply_window(ydata1,
                                                  mlab.window_hanning,
                                                  return_window=True)
        ycontrol2 = mlab.window_hanning(ydata2)
        ydata = np.vstack([ydata1, ydata2])
        ycontrol = np.vstack([ycontrol1, ycontrol2])
        ydata = np.tile(ydata, (20, 1))
        ycontrol = np.tile(ycontrol, (20, 1))
        ydatab = ydata.T.flatten()
        ydataf = ydata.flatten()
        ycontrol = ycontrol.flatten()
        spec_g, fsp_g = mlab.psd(x=ydataf,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 window=mlab.window_hanning)
        spec_b, fsp_b = mlab.psd(x=ydatab,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 window=mlab.window_hanning)
        spec_c, fsp_c = mlab.psd(x=ycontrol,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 window=mlab.window_none)
        spec_c *= len(ycontrol1)/(np.abs(windowVals)**2).sum()
        assert_array_equal(fsp_g, fsp_c)
        assert_array_equal(fsp_b, fsp_c)
        assert_allclose(spec_g, spec_c, atol=1e-08)
        # these should not be almost equal
        with pytest.raises(AssertionError):
            assert_allclose(spec_b, spec_c, atol=1e-08)

    def test_psd_window_hanning_detrend_linear(self):
        if self.NFFT_density is None:
            return
        freqs = self.freqs_density
        ydata = np.arange(self.NFFT_density)
        ycontrol = np.zeros(self.NFFT_density)
        ydata1 = ydata+5
        ydata2 = ydata+3.3
        ycontrol1 = ycontrol
        ycontrol2 = ycontrol
        ycontrol1, windowVals = mlab.apply_window(ycontrol1,
                                                  mlab.window_hanning,
                                                  return_window=True)
        ycontrol2 = mlab.window_hanning(ycontrol2)
        ydata = np.vstack([ydata1, ydata2])
        ycontrol = np.vstack([ycontrol1, ycontrol2])
        ydata = np.tile(ydata, (20, 1))
        ycontrol = np.tile(ycontrol, (20, 1))
        ydatab = ydata.T.flatten()
        ydataf = ydata.flatten()
        ycontrol = ycontrol.flatten()
        spec_g, fsp_g = mlab.psd(x=ydataf,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend=mlab.detrend_linear,
                                 window=mlab.window_hanning)
        spec_b, fsp_b = mlab.psd(x=ydatab,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 detrend=mlab.detrend_linear,
                                 window=mlab.window_hanning)
        spec_c, fsp_c = mlab.psd(x=ycontrol,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=0,
                                 sides=self.sides,
                                 window=mlab.window_none)
        spec_c *= len(ycontrol1)/(np.abs(windowVals)**2).sum()
        assert_array_equal(fsp_g, fsp_c)
        assert_array_equal(fsp_b, fsp_c)
        assert_allclose(spec_g, spec_c, atol=1e-08)
        # these should not be almost equal
        with pytest.raises(AssertionError):
            assert_allclose(spec_b, spec_c, atol=1e-08)

    def test_psd_windowarray(self):
        freqs = self.freqs_density
        spec, fsp = mlab.psd(x=self.y,
                             NFFT=self.NFFT_density,
                             Fs=self.Fs,
                             noverlap=self.nover_density,
                             pad_to=self.pad_to_density,
                             sides=self.sides,
                             window=np.ones(self.NFFT_density_real))
        assert_allclose(fsp, freqs, atol=1e-06)
        assert spec.shape == freqs.shape

    def test_psd_windowarray_scale_by_freq(self):
        freqs = self.freqs_density
        win = mlab.window_hanning(np.ones(self.NFFT_density_real))

        spec, fsp = mlab.psd(x=self.y,
                             NFFT=self.NFFT_density,
                             Fs=self.Fs,
                             noverlap=self.nover_density,
                             pad_to=self.pad_to_density,
                             sides=self.sides,
                             window=mlab.window_hanning)
        spec_s, fsp_s = mlab.psd(x=self.y,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=self.nover_density,
                                 pad_to=self.pad_to_density,
                                 sides=self.sides,
                                 window=mlab.window_hanning,
                                 scale_by_freq=True)
        spec_n, fsp_n = mlab.psd(x=self.y,
                                 NFFT=self.NFFT_density,
                                 Fs=self.Fs,
                                 noverlap=self.nover_density,
                                 pad_to=self.pad_to_density,
                                 sides=self.sides,
                                 window=mlab.window_hanning,
                                 scale_by_freq=False)
        assert_array_equal(fsp, fsp_s)
        assert_array_equal(fsp, fsp_n)
        assert_array_equal(spec, spec_s)
        assert_allclose(spec_s*(win**2).sum(),
                        spec_n/self.Fs*win.sum()**2,
                        atol=1e-08)

    def test_complex_spectrum(self):
        freqs = self.freqs_spectrum
        spec, fsp = mlab.complex_spectrum(x=self.y,
                                          Fs=self.Fs,
                                          sides=self.sides,
                                          pad_to=self.pad_to_spectrum)
        assert_allclose(fsp, freqs, atol=1e-06)
        assert spec.shape == freqs.shape

    def test_magnitude_spectrum(self):
        freqs = self.freqs_spectrum
        spec, fsp = mlab.magnitude_spectrum(x=self.y,
                                            Fs=self.Fs,
                                            sides=self.sides,
                                            pad_to=self.pad_to_spectrum)
        assert spec.shape == freqs.shape
        self.check_maxfreq(spec, fsp, self.fstims)
        self.check_freqs(spec, freqs, fsp, self.fstims)

    def test_angle_spectrum(self):
        freqs = self.freqs_spectrum
        spec, fsp = mlab.angle_spectrum(x=self.y,
                                        Fs=self.Fs,
                                        sides=self.sides,
                                        pad_to=self.pad_to_spectrum)
        assert_allclose(fsp, freqs, atol=1e-06)
        assert spec.shape == freqs.shape

    def test_phase_spectrum(self):
        freqs = self.freqs_spectrum
        spec, fsp = mlab.phase_spectrum(x=self.y,
                                        Fs=self.Fs,
                                        sides=self.sides,
                                        pad_to=self.pad_to_spectrum)
        assert_allclose(fsp, freqs, atol=1e-06)
        assert spec.shape == freqs.shape

    def test_specgram_auto(self):
        freqs = self.freqs_specgram
        spec, fsp, t = mlab.specgram(x=self.y,
                                     NFFT=self.NFFT_specgram,
                                     Fs=self.Fs,
                                     noverlap=self.nover_specgram,
                                     pad_to=self.pad_to_specgram,
                                     sides=self.sides)
        specm = np.mean(spec, axis=1)

        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_specgram, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]

        # since we are using a single freq, all time slices
        # should be about the same
        if np.abs(spec.max()) != 0:
            assert_allclose(np.diff(spec, axis=1).max()/np.abs(spec.max()), 0,
                            atol=1e-02)
        self.check_freqs(specm, freqs, fsp, self.fstims)

    def test_specgram_default(self):
        freqs = self.freqs_specgram
        spec, fsp, t = mlab.specgram(x=self.y,
                                     NFFT=self.NFFT_specgram,
                                     Fs=self.Fs,
                                     noverlap=self.nover_specgram,
                                     pad_to=self.pad_to_specgram,
                                     sides=self.sides,
                                     mode='default')
        specm = np.mean(spec, axis=1)

        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_specgram, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]

        # since we are using a single freq, all time slices
        # should be about the same
        if np.abs(spec.max()) != 0:
            assert_allclose(np.diff(spec, axis=1).max()/np.abs(spec.max()), 0,
                            atol=1e-02)
        self.check_freqs(specm, freqs, fsp, self.fstims)

    def test_specgram_psd(self):
        freqs = self.freqs_specgram
        spec, fsp, t = mlab.specgram(x=self.y,
                                     NFFT=self.NFFT_specgram,
                                     Fs=self.Fs,
                                     noverlap=self.nover_specgram,
                                     pad_to=self.pad_to_specgram,
                                     sides=self.sides,
                                     mode='psd')
        specm = np.mean(spec, axis=1)

        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_specgram, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]
        # since we are using a single freq, all time slices
        # should be about the same
        if np.abs(spec.max()) != 0:
            assert_allclose(np.diff(spec, axis=1).max()/np.abs(spec.max()), 0,
                            atol=1e-02)
        self.check_freqs(specm, freqs, fsp, self.fstims)

    def test_specgram_complex(self):
        freqs = self.freqs_specgram
        spec, fsp, t = mlab.specgram(x=self.y,
                                     NFFT=self.NFFT_specgram,
                                     Fs=self.Fs,
                                     noverlap=self.nover_specgram,
                                     pad_to=self.pad_to_specgram,
                                     sides=self.sides,
                                     mode='complex')
        specm = np.mean(np.abs(spec), axis=1)
        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_specgram, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]

        self.check_freqs(specm, freqs, fsp, self.fstims)

    def test_specgram_magnitude(self):
        freqs = self.freqs_specgram
        spec, fsp, t = mlab.specgram(x=self.y,
                                     NFFT=self.NFFT_specgram,
                                     Fs=self.Fs,
                                     noverlap=self.nover_specgram,
                                     pad_to=self.pad_to_specgram,
                                     sides=self.sides,
                                     mode='magnitude')
        specm = np.mean(spec, axis=1)
        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_specgram, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]
        # since we are using a single freq, all time slices
        # should be about the same
        if np.abs(spec.max()) != 0:
            assert_allclose(np.diff(spec, axis=1).max()/np.abs(spec.max()), 0,
                            atol=1e-02)
        self.check_freqs(specm, freqs, fsp, self.fstims)

    def test_specgram_angle(self):
        freqs = self.freqs_specgram
        spec, fsp, t = mlab.specgram(x=self.y,
                                     NFFT=self.NFFT_specgram,
                                     Fs=self.Fs,
                                     noverlap=self.nover_specgram,
                                     pad_to=self.pad_to_specgram,
                                     sides=self.sides,
                                     mode='angle')
        specm = np.mean(spec, axis=1)
        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_specgram, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]

    def test_specgram_phase(self):
        freqs = self.freqs_specgram
        spec, fsp, t = mlab.specgram(x=self.y,
                                     NFFT=self.NFFT_specgram,
                                     Fs=self.Fs,
                                     noverlap=self.nover_specgram,
                                     pad_to=self.pad_to_specgram,
                                     sides=self.sides,
                                     mode='phase')
        specm = np.mean(spec, axis=1)

        assert_allclose(fsp, freqs, atol=1e-06)
        assert_allclose(t, self.t_specgram, atol=1e-06)

        assert spec.shape[0] == freqs.shape[0]
        assert spec.shape[1] == self.t_specgram.shape[0]

    def test_specgram_warn_only1seg(self):
        """Warning should be raised if len(x) <= NFFT. """
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always", category=UserWarning)
            mlab.specgram(x=self.y, NFFT=len(self.y), Fs=self.Fs)
        assert len(w) == 1
        assert issubclass(w[0].category, UserWarning)
        assert str(w[0].message).startswith("Only one segment is calculated")

    def test_psd_csd_equal(self):
        freqs = self.freqs_density
        Pxx, freqsxx = mlab.psd(x=self.y,
                                NFFT=self.NFFT_density,
                                Fs=self.Fs,
                                noverlap=self.nover_density,
                                pad_to=self.pad_to_density,
                                sides=self.sides)
        Pxy, freqsxy = mlab.csd(x=self.y, y=self.y,
                                NFFT=self.NFFT_density,
                                Fs=self.Fs,
                                noverlap=self.nover_density,
                                pad_to=self.pad_to_density,
                                sides=self.sides)
        assert_array_almost_equal_nulp(Pxx, Pxy)
        assert_array_equal(freqsxx, freqsxy)

    def test_specgram_auto_default_equal(self):
        '''test that mlab.specgram without mode and with mode 'default' and
        'psd' are all the same'''
        freqs = self.freqs_specgram
        speca, freqspeca, ta = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides)
        specb, freqspecb, tb = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='default')
        assert_array_equal(speca, specb)
        assert_array_equal(freqspeca, freqspecb)
        assert_array_equal(ta, tb)

    def test_specgram_auto_psd_equal(self):
        '''test that mlab.specgram without mode and with mode 'default' and
        'psd' are all the same'''
        freqs = self.freqs_specgram
        speca, freqspeca, ta = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides)
        specc, freqspecc, tc = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='psd')
        assert_array_equal(speca, specc)
        assert_array_equal(freqspeca, freqspecc)
        assert_array_equal(ta, tc)

    def test_specgram_complex_mag_equivalent(self):
        freqs = self.freqs_specgram
        specc, freqspecc, tc = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='complex')
        specm, freqspecm, tm = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='magnitude')

        assert_array_equal(freqspecc, freqspecm)
        assert_array_equal(tc, tm)
        assert_allclose(np.abs(specc), specm, atol=1e-06)

    def test_specgram_complex_angle_equivalent(self):
        freqs = self.freqs_specgram
        specc, freqspecc, tc = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='complex')
        speca, freqspeca, ta = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='angle')

        assert_array_equal(freqspecc, freqspeca)
        assert_array_equal(tc, ta)
        assert_allclose(np.angle(specc), speca, atol=1e-06)

    def test_specgram_complex_phase_equivalent(self):
        freqs = self.freqs_specgram
        specc, freqspecc, tc = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='complex')
        specp, freqspecp, tp = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='phase')

        assert_array_equal(freqspecc, freqspecp)
        assert_array_equal(tc, tp)
        assert_allclose(np.unwrap(np.angle(specc), axis=0), specp,
                        atol=1e-06)

    def test_specgram_angle_phase_equivalent(self):
        freqs = self.freqs_specgram
        speca, freqspeca, ta = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='angle')
        specp, freqspecp, tp = mlab.specgram(x=self.y,
                                             NFFT=self.NFFT_specgram,
                                             Fs=self.Fs,
                                             noverlap=self.nover_specgram,
                                             pad_to=self.pad_to_specgram,
                                             sides=self.sides,
                                             mode='phase')

        assert_array_equal(freqspeca, freqspecp)
        assert_array_equal(ta, tp)
        assert_allclose(np.unwrap(speca, axis=0), specp,
                        atol=1e-06)

    def test_psd_windowarray_equal(self):
        freqs = self.freqs_density
        win = mlab.window_hanning(np.ones(self.NFFT_density_real))
        speca, fspa = mlab.psd(x=self.y,
                               NFFT=self.NFFT_density,
                               Fs=self.Fs,
                               noverlap=self.nover_density,
                               pad_to=self.pad_to_density,
                               sides=self.sides,
                               window=win)
        specb, fspb = mlab.psd(x=self.y,
                               NFFT=self.NFFT_density,
                               Fs=self.Fs,
                               noverlap=self.nover_density,
                               pad_to=self.pad_to_density,
                               sides=self.sides)
        assert_array_equal(fspa, fspb)
        assert_allclose(speca, specb, atol=1e-08)


# extra test for cohere...
def test_cohere():
    N = 1024
    np.random.seed(19680801)
    x = np.random.randn(N)
    # phase offset
    y = np.roll(x, 20)
    # high-freq roll-off
    y = np.convolve(y, np.ones(20) / 20., mode='same')
    cohsq, f = mlab.cohere(x, y, NFFT=256, Fs=2, noverlap=128)
    assert_allclose(np.mean(cohsq), 0.837, atol=1.e-3)
    assert np.isreal(np.mean(cohsq))


def test_griddata_linear():
    # z is a linear function of x and y.
    def get_z(x, y):
        return 3.0*x - y

    # Passing 1D xi and yi arrays to griddata.
    x = np.asarray([0.0, 1.0, 0.0, 1.0, 0.5])
    y = np.asarray([0.0, 0.0, 1.0, 1.0, 0.5])
    z = get_z(x, y)
    xi = [0.2, 0.4, 0.6, 0.8]
    yi = [0.1, 0.3, 0.7, 0.9]
    zi = mlab.griddata(x, y, z, xi, yi, interp='linear')
    xi, yi = np.meshgrid(xi, yi)
    np.testing.assert_array_almost_equal(zi, get_z(xi, yi))

    # Passing 2D xi and yi arrays to griddata.
    zi = mlab.griddata(x, y, z, xi, yi, interp='linear')
    np.testing.assert_array_almost_equal(zi, get_z(xi, yi))

    # Masking z array.
    z_masked = np.ma.array(z, mask=[False, False, False, True, False])
    correct_zi_masked = np.ma.masked_where(xi + yi > 1.0, get_z(xi, yi))
    zi = mlab.griddata(x, y, z_masked, xi, yi, interp='linear')
    matest.assert_array_almost_equal(zi, correct_zi_masked)
    np.testing.assert_array_equal(np.ma.getmask(zi),
                                  np.ma.getmask(correct_zi_masked))


@pytest.mark.xfail(not HAS_NATGRID, reason='natgrid not installed')
def test_griddata_nn():
    # z is a linear function of x and y.
    def get_z(x, y):
        return 3.0*x - y

    # Passing 1D xi and yi arrays to griddata.
    x = np.asarray([0.0, 1.0, 0.0, 1.0, 0.5])
    y = np.asarray([0.0, 0.0, 1.0, 1.0, 0.5])
    z = get_z(x, y)
    xi = [0.2, 0.4, 0.6, 0.8]
    yi = [0.1, 0.3, 0.7, 0.9]
    correct_zi = [[0.49999252, 1.0999978, 1.7000030, 2.3000080],
                  [0.29999208, 0.8999978, 1.5000029, 2.1000059],
                  [-0.1000099, 0.4999943, 1.0999964, 1.6999979],
                  [-0.3000128, 0.2999894, 0.8999913, 1.4999933]]
    zi = mlab.griddata(x, y, z, xi, yi, interp='nn')
    np.testing.assert_array_almost_equal(zi, correct_zi, 5)

    # Decreasing xi or yi should raise ValueError.
    with pytest.raises(ValueError):
        mlab.griddata(x, y, z, xi[::-1], yi, interp='nn')
    with pytest.raises(ValueError):
        mlab.griddata(x, y, z, xi, yi[::-1], interp='nn')

    # Passing 2D xi and yi arrays to griddata.
    xi, yi = np.meshgrid(xi, yi)
    zi = mlab.griddata(x, y, z, xi, yi, interp='nn')
    np.testing.assert_array_almost_equal(zi, correct_zi, 5)

    # Masking z array.
    z_masked = np.ma.array(z, mask=[False, False, False, True, False])
    correct_zi_masked = np.ma.masked_where(xi + yi > 1.0, correct_zi)
    zi = mlab.griddata(x, y, z_masked, xi, yi, interp='nn')
    np.testing.assert_array_almost_equal(zi, correct_zi_masked, 5)
    np.testing.assert_array_equal(np.ma.getmask(zi),
                                  np.ma.getmask(correct_zi_masked))


#*****************************************************************
# These Tests where taken from SCIPY with some minor modifications
# this can be retreived from:
# https://github.com/scipy/scipy/blob/master/scipy/stats/tests/test_kdeoth.py
#*****************************************************************

class TestGaussianKDE(object):

    def test_kde_integer_input(self):
        """Regression test for #1181."""
        x1 = np.arange(5)
        kde = mlab.GaussianKDE(x1)
        y_expected = [0.13480721, 0.18222869, 0.19514935, 0.18222869,
                      0.13480721]
        np.testing.assert_array_almost_equal(kde(x1), y_expected, decimal=6)

    def test_gaussian_kde_covariance_caching(self):
        x1 = np.array([-7, -5, 1, 4, 5], dtype=float)
        xs = np.linspace(-10, 10, num=5)
        # These expected values are from scipy 0.10, before some changes to
        # gaussian_kde. They were not compared with any external reference.
        y_expected = [0.02463386, 0.04689208, 0.05395444, 0.05337754,
                      0.01664475]

        # set it to the default bandwidth.
        kde2 = mlab.GaussianKDE(x1, 'scott')
        y2 = kde2(xs)

        np.testing.assert_array_almost_equal(y_expected, y2, decimal=7)

    def test_kde_bandwidth_method(self):

        np.random.seed(8765678)
        n_basesample = 50
        xn = np.random.randn(n_basesample)

        # Default
        gkde = mlab.GaussianKDE(xn)
        # Supply a callable
        gkde2 = mlab.GaussianKDE(xn, 'scott')
        # Supply a scalar
        gkde3 = mlab.GaussianKDE(xn, bw_method=gkde.factor)

        xs = np.linspace(-7, 7, 51)
        kdepdf = gkde.evaluate(xs)
        kdepdf2 = gkde2.evaluate(xs)
        assert kdepdf.all() == kdepdf2.all()
        kdepdf3 = gkde3.evaluate(xs)
        assert kdepdf.all() == kdepdf3.all()


class TestGaussianKDECustom(object):
    def test_no_data(self):
        """Pass no data into the GaussianKDE class."""
        with pytest.raises(ValueError):
            mlab.GaussianKDE([])

    def test_single_dataset_element(self):
        """Pass a single dataset element into the GaussianKDE class."""
        with pytest.raises(ValueError):
            mlab.GaussianKDE([42])

    def test_silverman_multidim_dataset(self):
        """Use a multi-dimensional array as the dataset and test silverman's
        output"""
        x1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
        with pytest.raises(np.linalg.LinAlgError):
            mlab.GaussianKDE(x1, "silverman")

    def test_silverman_singledim_dataset(self):
        """Use a single dimension list as the dataset and test silverman's
        output."""
        x1 = np.array([-7, -5, 1, 4, 5])
        mygauss = mlab.GaussianKDE(x1, "silverman")
        y_expected = 0.76770389927475502
        assert_almost_equal(mygauss.covariance_factor(), y_expected, 7)

    def test_scott_multidim_dataset(self):
        """Use a multi-dimensional array as the dataset and test scott's output
        """
        x1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
        with pytest.raises(np.linalg.LinAlgError):
            mlab.GaussianKDE(x1, "scott")

    def test_scott_singledim_dataset(self):
        """Use a single-dimensional array as the dataset and test scott's
        output"""
        x1 = np.array([-7, -5, 1, 4, 5])
        mygauss = mlab.GaussianKDE(x1, "scott")
        y_expected = 0.72477966367769553
        assert_almost_equal(mygauss.covariance_factor(), y_expected, 7)

    def test_scalar_empty_dataset(self):
        """Use an empty array as the dataset and test the scalar's cov factor
        """
        with pytest.raises(ValueError):
            mlab.GaussianKDE([], bw_method=5)

    def test_scalar_covariance_dataset(self):
        """Use a dataset and test a scalar's cov factor
        """
        np.random.seed(8765678)
        n_basesample = 50
        multidim_data = [np.random.randn(n_basesample) for i in range(5)]

        kde = mlab.GaussianKDE(multidim_data, bw_method=0.5)
        assert kde.covariance_factor() == 0.5

    def test_callable_covariance_dataset(self):
        """Use a multi-dimensional array as the dataset and test the callable's
        cov factor"""
        np.random.seed(8765678)
        n_basesample = 50
        multidim_data = [np.random.randn(n_basesample) for i in range(5)]

        def callable_fun(x):
            return 0.55
        kde = mlab.GaussianKDE(multidim_data, bw_method=callable_fun)
        assert kde.covariance_factor() == 0.55

    def test_callable_singledim_dataset(self):
        """Use a single-dimensional array as the dataset and test the
        callable's cov factor"""
        np.random.seed(8765678)
        n_basesample = 50
        multidim_data = np.random.randn(n_basesample)

        kde = mlab.GaussianKDE(multidim_data, bw_method='silverman')
        y_expected = 0.48438841363348911
        assert_almost_equal(kde.covariance_factor(), y_expected, 7)

    def test_wrong_bw_method(self):
        """Test the error message that should be called when bw is invalid."""
        np.random.seed(8765678)
        n_basesample = 50
        data = np.random.randn(n_basesample)
        with pytest.raises(ValueError):
            mlab.GaussianKDE(data, bw_method="invalid")


class TestGaussianKDEEvaluate(object):

    def test_evaluate_diff_dim(self):
        """Test the evaluate method when the dim's of dataset and points are
        different dimensions"""
        x1 = np.arange(3, 10, 2)
        kde = mlab.GaussianKDE(x1)
        x2 = np.arange(3, 12, 2)
        y_expected = [
            0.08797252, 0.11774109, 0.11774109, 0.08797252, 0.0370153
        ]
        y = kde.evaluate(x2)
        np.testing.assert_array_almost_equal(y, y_expected, 7)

    def test_evaluate_inv_dim(self):
        """ Invert the dimensions. i.e., Give the dataset a dimension of
        1 [3,2,4], and the points will have a dimension of 3 [[3],[2],[4]].
        ValueError should be raised"""
        np.random.seed(8765678)
        n_basesample = 50
        multidim_data = np.random.randn(n_basesample)
        kde = mlab.GaussianKDE(multidim_data)
        x2 = [[1], [2], [3]]
        with pytest.raises(ValueError):
            kde.evaluate(x2)

    def test_evaluate_dim_and_num(self):
        """ Tests if evaluated against a one by one array"""
        x1 = np.arange(3, 10, 2)
        x2 = np.array([3])
        kde = mlab.GaussianKDE(x1)
        y_expected = [0.08797252]
        y = kde.evaluate(x2)
        np.testing.assert_array_almost_equal(y, y_expected, 7)

    def test_evaluate_point_dim_not_one(self):
        """Test"""
        x1 = np.arange(3, 10, 2)
        x2 = [np.arange(3, 10, 2), np.arange(3, 10, 2)]
        kde = mlab.GaussianKDE(x1)
        with pytest.raises(ValueError):
            kde.evaluate(x2)

    def test_evaluate_equal_dim_and_num_lt(self):
        """Test when line 3810 fails"""
        x1 = np.arange(3, 10, 2)
        x2 = np.arange(3, 8, 2)
        kde = mlab.GaussianKDE(x1)
        y_expected = [0.08797252, 0.11774109, 0.11774109]
        y = kde.evaluate(x2)
        np.testing.assert_array_almost_equal(y, y_expected, 7)


def test_contiguous_regions():
    a, b, c = 3, 4, 5
    # Starts and ends with True
    mask = [True]*a + [False]*b + [True]*c
    expected = [(0, a), (a+b, a+b+c)]
    assert mlab.contiguous_regions(mask) == expected
    d, e = 6, 7
    # Starts with True ends with False
    mask = mask + [False]*e
    assert mlab.contiguous_regions(mask) == expected
    # Starts with False ends with True
    mask = [False]*d + mask[:-e]
    expected = [(d, d+a), (d+a+b, d+a+b+c)]
    assert mlab.contiguous_regions(mask) == expected
    # Starts and ends with False
    mask = mask + [False]*e
    assert mlab.contiguous_regions(mask) == expected
    # No True in mask
    assert mlab.contiguous_regions([False]*5) == []
    # Empty mask
    assert mlab.contiguous_regions([]) == []


def test_psd_onesided_norm():
    u = np.array([0, 1, 2, 3, 1, 2, 1])
    dt = 1.0
    Su = np.abs(np.fft.fft(u) * dt)**2 / float(dt * u.size)
    P, f = mlab.psd(u, NFFT=u.size, Fs=1/dt, window=mlab.window_none,
                    detrend=mlab.detrend_none, noverlap=0, pad_to=None,
                    scale_by_freq=None,
                    sides='onesided')
    Su_1side = np.append([Su[0]], Su[1:4] + Su[4:][::-1])
    assert_allclose(P, Su_1side, atol=1e-06)


def test_psd_oversampling():
    """Test the case len(x) < NFFT for psd(). """
    u = np.array([0, 1, 2, 3, 1, 2, 1])
    dt = 1.0
    Su = np.abs(np.fft.fft(u) * dt)**2 / float(dt * u.size)
    P, f = mlab.psd(u, NFFT=u.size*2, Fs=1/dt, window=mlab.window_none,
                    detrend=mlab.detrend_none, noverlap=0, pad_to=None,
                    scale_by_freq=None,
                    sides='onesided')
    Su_1side = np.append([Su[0]], Su[1:4] + Su[4:][::-1])
    assert_almost_equal(np.sum(P), np.sum(Su_1side))  # same energy

social_vdgnn/lib/python3.6/site-packages/matplotlib/tests/test_mlab.py
eliatrevisan/IA-MPPI-LBM
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import sys

import numpy as np
from numpy.testing import assert_array_almost_equal
import matplotlib.pyplot as plt
from matplotlib.testing.decorators import image_comparison
import matplotlib.transforms as mtransforms


on_win = (sys.platform == 'win32')


def velocity_field():
    Y, X = np.mgrid[-3:3:100j, -3:3:100j]
    U = -1 - X**2 + Y
    V = 1 + X - Y**2
    return X, Y, U, V


def swirl_velocity_field():
    x = np.linspace(-3., 3., 100)
    y = np.linspace(-3., 3., 100)
    X, Y = np.meshgrid(x, y)
    a = 0.1
    U = np.cos(a) * (-Y) - np.sin(a) * X
    V = np.sin(a) * (-Y) + np.cos(a) * X
    return x, y, U, V


@image_comparison(baseline_images=['streamplot_startpoints'])
def test_startpoints():
    X, Y, U, V = velocity_field()
    start_x = np.linspace(X.min(), X.max(), 10)
    start_y = np.linspace(Y.min(), Y.max(), 10)
    start_points = list(zip(start_x, start_y))
    plt.streamplot(X, Y, U, V, start_points=start_points)
    plt.plot(start_x, start_y, 'ok')


@image_comparison(baseline_images=['streamplot_colormap'],
                  tol=.02)
def test_colormap():
    X, Y, U, V = velocity_field()
    plt.streamplot(X, Y, U, V, color=U, density=0.6, linewidth=2,
                   cmap=plt.cm.autumn)
    plt.colorbar()


@image_comparison(baseline_images=['streamplot_linewidth'])
def test_linewidth():
    X, Y, U, V = velocity_field()
    speed = np.sqrt(U*U + V*V)
    lw = 5*speed/speed.max()
    df = 25. / 30.   # Compatibility factor for old test image
    plt.streamplot(X, Y, U, V, density=[0.5 * df, 1. * df], color='k',
                   linewidth=lw)


@image_comparison(baseline_images=['streamplot_masks_and_nans'],
                  tol=0.04 if on_win else 0)
def test_masks_and_nans():
    X, Y, U, V = velocity_field()
    mask = np.zeros(U.shape, dtype=bool)
    mask[40:60, 40:60] = 1
    U[:20, :20] = np.nan
    U = np.ma.array(U, mask=mask)
    with np.errstate(invalid='ignore'):
        plt.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues)


@image_comparison(baseline_images=['streamplot_maxlength'],
                  extensions=['png'])
def test_maxlength():
    x, y, U, V = swirl_velocity_field()
    plt.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]],
                   linewidth=2, density=2)


@image_comparison(baseline_images=['streamplot_direction'],
                  extensions=['png'])
def test_direction():
    x, y, U, V = swirl_velocity_field()
    plt.streamplot(x, y, U, V, integration_direction='backward',
                   maxlength=1.5, start_points=[[1.5, 0.]],
                   linewidth=2, density=2)


def test_streamplot_limits():
    ax = plt.axes()
    x = np.linspace(-5, 10, 20)
    y = np.linspace(-2, 4, 10)
    y, x = np.meshgrid(y, x)
    trans = mtransforms.Affine2D().translate(25, 32) + ax.transData
    plt.barbs(x, y, np.sin(x), np.cos(y), transform=trans)
    # The calculated bounds are approximately the bounds of the original data,
    # this is because the entire path is taken into account when updating the
    # datalim.
    assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6),
                              decimal=1)

social_vdgnn/lib/python3.6/site-packages/matplotlib/tests/test_streamplot.py
eliatrevisan/IA-MPPI-LBM
"""
This module is to support *bbox_inches* option in savefig command.
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six

import warnings
from matplotlib.transforms import Bbox, TransformedBbox, Affine2D


def adjust_bbox(fig, bbox_inches, fixed_dpi=None):
    """
    Temporarily adjust the figure so that only the specified area
    (bbox_inches) is saved.

    It modifies fig.bbox, fig.bbox_inches,
    fig.transFigure._boxout, and fig.patch.  While the figure size
    changes, the scale of the original figure is conserved.  A
    function which restores the original values are returned.
    """

    origBbox = fig.bbox
    origBboxInches = fig.bbox_inches
    _boxout = fig.transFigure._boxout

    asp_list = []
    locator_list = []
    for ax in fig.axes:
        pos = ax.get_position(original=False).frozen()
        locator_list.append(ax.get_axes_locator())
        asp_list.append(ax.get_aspect())

        def _l(a, r, pos=pos):
            return pos
        ax.set_axes_locator(_l)
        ax.set_aspect("auto")

    def restore_bbox():

        for ax, asp, loc in zip(fig.axes, asp_list, locator_list):
            ax.set_aspect(asp)
            ax.set_axes_locator(loc)

        fig.bbox = origBbox
        fig.bbox_inches = origBboxInches
        fig.transFigure._boxout = _boxout
        fig.transFigure.invalidate()
        fig.patch.set_bounds(0, 0, 1, 1)

    if fixed_dpi is not None:
        tr = Affine2D().scale(fixed_dpi)
        dpi_scale = fixed_dpi / fig.dpi
    else:
        tr = Affine2D().scale(fig.dpi)
        dpi_scale = 1.

    _bbox = TransformedBbox(bbox_inches, tr)

    fig.bbox_inches = Bbox.from_bounds(0, 0,
                                       bbox_inches.width, bbox_inches.height)
    x0, y0 = _bbox.x0, _bbox.y0
    w1, h1 = fig.bbox.width * dpi_scale, fig.bbox.height * dpi_scale
    fig.transFigure._boxout = Bbox.from_bounds(-x0, -y0, w1, h1)
    fig.transFigure.invalidate()

    fig.bbox = TransformedBbox(fig.bbox_inches, tr)

    fig.patch.set_bounds(x0 / w1, y0 / h1,
                         fig.bbox.width / w1, fig.bbox.height / h1)

    return restore_bbox


def process_figure_for_rasterizing(fig, bbox_inches_restore, fixed_dpi=None):
    """
    This need to be called when figure dpi changes during the drawing
    (e.g., rasterizing). It recovers the bbox and re-adjust it with
    the new dpi.
    """

    bbox_inches, restore_bbox = bbox_inches_restore
    restore_bbox()
    r = adjust_bbox(fig, bbox_inches, fixed_dpi)

    return bbox_inches, r

social_vdgnn/lib/python3.6/site-packages/matplotlib/tight_bbox.py
eliatrevisan/IA-MPPI-LBM
"""
This module contains a class representing a Type 1 font.

This version reads pfa and pfb files and splits them for embedding in
pdf files. It also supports SlantFont and ExtendFont transformations,
similarly to pdfTeX and friends. There is no support yet for
subsetting.

Usage::

   >>> font = Type1Font(filename)
   >>> clear_part, encrypted_part, finale = font.parts
   >>> slanted_font = font.transform({'slant': 0.167})
   >>> extended_font = font.transform({'extend': 1.2})

Sources:

* Adobe Technical Note #5040, Supporting Downloadable PostScript
  Language Fonts.

* Adobe Type 1 Font Format, Adobe Systems Incorporated, third printing,
  v1.1, 1993. ISBN 0-201-57044-0.
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six
from six import unichr

import binascii
import io
import itertools
import numpy as np
import re
import struct
import sys

if six.PY3:
    def ord(x):
        return x


class Type1Font(object):
    """
    A class representing a Type-1 font, for use by backends.

    .. attribute:: parts

       A 3-tuple of the cleartext part, the encrypted part, and the
       finale of zeros.

    .. attribute:: prop

       A dictionary of font properties.
    """
    __slots__ = ('parts', 'prop')

    def __init__(self, input):
        """
        Initialize a Type-1 font. *input* can be either the file name of
        a pfb file or a 3-tuple of already-decoded Type-1 font parts.
        """
        if isinstance(input, tuple) and len(input) == 3:
            self.parts = input
        else:
            with open(input, 'rb') as file:
                data = self._read(file)
            self.parts = self._split(data)

        self._parse()

    def _read(self, file):
        """
        Read the font from a file, decoding into usable parts.
        """
        rawdata = file.read()
        if not rawdata.startswith(b'\x80'):
            return rawdata

        data = b''
        while len(rawdata) > 0:
            if not rawdata.startswith(b'\x80'):
                raise RuntimeError('Broken pfb file (expected byte 128, '
                                   'got %d)' % ord(rawdata[0]))
            type = ord(rawdata[1])
            if type in (1, 2):
                length, = struct.unpack(str('<i'), rawdata[2:6])
                segment = rawdata[6:6 + length]
                rawdata = rawdata[6 + length:]

            if type == 1:       # ASCII text: include verbatim
                data += segment
            elif type == 2:     # binary data: encode in hexadecimal
                data += binascii.hexlify(segment)
            elif type == 3:     # end of file
                break
            else:
                raise RuntimeError('Unknown segment type %d in pfb file' %
                                   type)

        return data

    def _split(self, data):
        """
        Split the Type 1 font into its three main parts.

        The three parts are: (1) the cleartext part, which ends in a
        eexec operator; (2) the encrypted part; (3) the fixed part,
        which contains 512 ASCII zeros possibly divided on various
        lines, a cleartomark operator, and possibly something else.
        """

        # Cleartext part: just find the eexec and skip whitespace
        idx = data.index(b'eexec')
        idx += len(b'eexec')
        while data[idx] in b' \t\r\n':
            idx += 1
        len1 = idx

        # Encrypted part: find the cleartomark operator and count
        # zeros backward
        idx = data.rindex(b'cleartomark') - 1
        zeros = 512
        while zeros and data[idx] in b'0' or data[idx] in b'\r\n':
            if data[idx] in b'0':
                zeros -= 1
            idx -= 1
        if zeros:
            raise RuntimeError('Insufficiently many zeros in Type 1 font')

        # Convert encrypted part to binary (if we read a pfb file, we
        # may end up converting binary to hexadecimal to binary again;
        # but if we read a pfa file, this part is already in hex, and
        # I am not quite sure if even the pfb format guarantees that
        # it will be in binary).
        binary = binascii.unhexlify(data[len1:idx+1])

        return data[:len1], binary, data[idx+1:]

    _whitespace_re = re.compile(br'[\0\t\r\014\n ]+')
    _token_re = re.compile(br'/{0,2}[^]\0\t\r\v\n ()<>{}/%[]+')
    _comment_re = re.compile(br'%[^\r\n\v]*')
    _instring_re = re.compile(br'[()\\]')

    # token types, compared via object identity (poor man's enum)
    _whitespace = object()
    _name = object()
    _string = object()
    _delimiter = object()
    _number = object()

    @classmethod
    def _tokens(cls, text):
        """
        A PostScript tokenizer. Yield (token, value) pairs such as
        (cls._whitespace, '   ') or (cls._name, '/Foobar').
        """
        pos = 0
        while pos < len(text):
            match = (cls._comment_re.match(text[pos:]) or
                     cls._whitespace_re.match(text[pos:]))
            if match:
                yield (cls._whitespace, match.group())
                pos += match.end()
            elif text[pos] == b'(':
                start = pos
                pos += 1
                depth = 1
                while depth:
                    match = cls._instring_re.search(text[pos:])
                    if match is None:
                        return
                    pos += match.end()
                    if match.group() == b'(':
                        depth += 1
                    elif match.group() == b')':
                        depth -= 1
                    else:  # a backslash - skip the next character
                        pos += 1
                yield (cls._string, text[start:pos])
            elif text[pos:pos + 2] in (b'<<', b'>>'):
                yield (cls._delimiter, text[pos:pos + 2])
                pos += 2
            elif text[pos] == b'<':
                start = pos
                pos += text[pos:].index(b'>')
                yield (cls._string, text[start:pos])
            else:
                match = cls._token_re.match(text[pos:])
                if match:
                    try:
                        float(match.group())
                        yield (cls._number, match.group())
                    except ValueError:
                        yield (cls._name, match.group())
                    pos += match.end()
                else:
                    yield (cls._delimiter, text[pos:pos + 1])
                    pos += 1

    def _parse(self):
        """
        Find the values of various font properties. This limited kind
        of parsing is described in Chapter 10 "Adobe Type Manager
        Compatibility" of the Type-1 spec.
        """
        # Start with reasonable defaults
        prop = {'weight': 'Regular', 'ItalicAngle': 0.0, 'isFixedPitch': False,
                'UnderlinePosition': -100, 'UnderlineThickness': 50}
        filtered = ((token, value)
                    for token, value in self._tokens(self.parts[0])
                    if token is not self._whitespace)
        # The spec calls this an ASCII format; in Python 2.x we could
        # just treat the strings and names as opaque bytes but let's
        # turn them into proper Unicode, and be lenient in case of high bytes.
        convert = lambda x: x.decode('ascii', 'replace')
        for token, value in filtered:
            if token is self._name and value.startswith(b'/'):
                key = convert(value[1:])
                token, value = next(filtered)
                if token is self._name:
                    if value in (b'true', b'false'):
                        value = value == b'true'
                    else:
                        value = convert(value.lstrip(b'/'))
                elif token is self._string:
                    value = convert(value.lstrip(b'(').rstrip(b')'))
                elif token is self._number:
                    if b'.' in value:
                        value = float(value)
                    else:
                        value = int(value)
                else:  # more complicated value such as an array
                    value = None
                if key != 'FontInfo' and value is not None:
                    prop[key] = value

        # Fill in the various *Name properties
        if 'FontName' not in prop:
            prop['FontName'] = (prop.get('FullName') or
                                prop.get('FamilyName') or
                                'Unknown')
        if 'FullName' not in prop:
            prop['FullName'] = prop['FontName']
        if 'FamilyName' not in prop:
            extras = r'(?i)([ -](regular|plain|italic|oblique|(semi)?bold|(ultra)?light|extra|condensed))+$'
            prop['FamilyName'] = re.sub(extras, '', prop['FullName'])

        self.prop = prop

    @classmethod
    def _transformer(cls, tokens, slant, extend):
        def fontname(name):
            result = name
            if slant:
                result += b'_Slant_' + str(int(1000 * slant)).encode('latin-1')
            if extend != 1.0:
                result += b'_Extend_' + str(int(1000 * extend)).encode('latin-1')
            return result

        def italicangle(angle):
            return str(float(angle) - np.arctan(slant) / np.pi * 180).encode('latin-1')

        def fontmatrix(array):
            array = array.lstrip(b'[').rstrip(b']').strip().split()
            array = [float(x) for x in array]
            oldmatrix = np.eye(3, 3)
            oldmatrix[0:3, 0] = array[::2]
            oldmatrix[0:3, 1] = array[1::2]
            modifier = np.array([[extend, 0, 0],
                                 [slant, 1, 0],
                                 [0, 0, 1]])
            newmatrix = np.dot(modifier, oldmatrix)
            array[::2] = newmatrix[0:3, 0]
            array[1::2] = newmatrix[0:3, 1]
            as_string = u'[' + u' '.join(str(x) for x in array) + u']'
            return as_string.encode('latin-1')

        def replace(fun):
            def replacer(tokens):
                token, value = next(tokens)      # name, e.g., /FontMatrix
                yield bytes(value)
                token, value = next(tokens)      # possible whitespace
                while token is cls._whitespace:
                    yield bytes(value)
                    token, value = next(tokens)
                if value != b'[':                # name/number/etc.
                    yield bytes(fun(value))
                else:                            # array, e.g., [1 2 3]
                    result = b''
                    while value != b']':
                        result += value
                        token, value = next(tokens)
                    result += value
                    yield fun(result)
            return replacer

        def suppress(tokens):
            for x in itertools.takewhile(lambda x: x[1] != b'def', tokens):
                pass
            yield b''

        table = {b'/FontName': replace(fontname),
                 b'/ItalicAngle': replace(italicangle),
                 b'/FontMatrix': replace(fontmatrix),
                 b'/UniqueID': suppress}

        for token, value in tokens:
            if token is cls._name and value in table:
                for value in table[value](itertools.chain([(token, value)],
                                                          tokens)):
                    yield value
            else:
                yield value

    def transform(self, effects):
        """
        Transform the font by slanting or extending. *effects* should
        be a dict where ``effects['slant']`` is the tangent of the
        angle that the font is to be slanted to the right (so negative
        values slant to the left) and ``effects['extend']`` is the
        multiplier by which the font is to be extended (so values less
        than 1.0 condense). Returns a new :class:`Type1Font` object.
        """
        with io.BytesIO() as buffer:
            tokenizer = self._tokens(self.parts[0])
            transformed =  self._transformer(tokenizer,
                                             slant=effects.get('slant', 0.0),
                                             extend=effects.get('extend', 1.0))
            list(map(buffer.write, transformed))
            return Type1Font((buffer.getvalue(), self.parts[1], self.parts[2]))

social_vdgnn/lib/python3.6/site-packages/matplotlib/type1font.py
eliatrevisan/IA-MPPI-LBM
import sys
import subprocess
import pkgutil
import types
import importlib
import warnings

import numpy as np
import numpy
import pytest

try:
    import ctypes
except ImportError:
    ctypes = None


def check_dir(module, module_name=None):
    """Returns a mapping of all objects with the wrong __module__ attribute."""
    if module_name is None:
        module_name = module.__name__
    results = {}
    for name in dir(module):
        item = getattr(module, name)
        if (hasattr(item, '__module__') and hasattr(item, '__name__')
                and item.__module__ != module_name):
            results[name] = item.__module__ + '.' + item.__name__
    return results


def test_numpy_namespace():
    # None of these objects are publicly documented to be part of the main
    # NumPy namespace (some are useful though, others need to be cleaned up)
    undocumented = {
        'Tester': 'numpy.testing._private.nosetester.NoseTester',
        '_add_newdoc_ufunc': 'numpy.core._multiarray_umath._add_newdoc_ufunc',
        'add_docstring': 'numpy.core._multiarray_umath.add_docstring',
        'add_newdoc': 'numpy.core.function_base.add_newdoc',
        'add_newdoc_ufunc': 'numpy.core._multiarray_umath._add_newdoc_ufunc',
        'byte_bounds': 'numpy.lib.utils.byte_bounds',
        'compare_chararrays': 'numpy.core._multiarray_umath.compare_chararrays',
        'deprecate': 'numpy.lib.utils.deprecate',
        'deprecate_with_doc': 'numpy.lib.utils.<lambda>',
        'disp': 'numpy.lib.function_base.disp',
        'fastCopyAndTranspose': 'numpy.core._multiarray_umath._fastCopyAndTranspose',
        'get_array_wrap': 'numpy.lib.shape_base.get_array_wrap',
        'get_include': 'numpy.lib.utils.get_include',
        'mafromtxt': 'numpy.lib.npyio.mafromtxt',
        'ndfromtxt': 'numpy.lib.npyio.ndfromtxt',
        'recfromcsv': 'numpy.lib.npyio.recfromcsv',
        'recfromtxt': 'numpy.lib.npyio.recfromtxt',
        'safe_eval': 'numpy.lib.utils.safe_eval',
        'set_string_function': 'numpy.core.arrayprint.set_string_function',
        'show_config': 'numpy.__config__.show',
        'who': 'numpy.lib.utils.who',
    }
    # These built-in types are re-exported by numpy.
    builtins = {
        'bool': 'builtins.bool',
        'complex': 'builtins.complex',
        'float': 'builtins.float',
        'int': 'builtins.int',
        'long': 'builtins.int',
        'object': 'builtins.object',
        'str': 'builtins.str',
        'unicode': 'builtins.str',
    }
    whitelist = dict(undocumented, **builtins)
    bad_results = check_dir(np)
    # pytest gives better error messages with the builtin assert than with
    # assert_equal
    assert bad_results == whitelist


@pytest.mark.parametrize('name', ['testing', 'Tester'])
def test_import_lazy_import(name):
    """Make sure we can actually use the modules we lazy load.

    While not exported as part of the public API, it was accessible.  With the
    use of __getattr__ and __dir__, this isn't always true It can happen that
    an infinite recursion may happen.

    This is the only way I found that would force the failure to appear on the
    badly implemented code.

    We also test for the presence of the lazily imported modules in dir

    """
    exe = (sys.executable, '-c', "import numpy; numpy." + name)
    result = subprocess.check_output(exe)
    assert not result

    # Make sure they are still in the __dir__
    assert name in dir(np)


def test_dir_testing():
    """Assert that output of dir has only one "testing/tester"
    attribute without duplicate"""
    assert len(dir(np)) == len(set(dir(np)))
    

def test_numpy_linalg():
    bad_results = check_dir(np.linalg)
    assert bad_results == {}


def test_numpy_fft():
    bad_results = check_dir(np.fft)
    assert bad_results == {}


@pytest.mark.skipif(ctypes is None,
                    reason="ctypes not available in this python")
def test_NPY_NO_EXPORT():
    cdll = ctypes.CDLL(np.core._multiarray_tests.__file__)
    # Make sure an arbitrary NPY_NO_EXPORT function is actually hidden
    f = getattr(cdll, 'test_not_exported', None)
    assert f is None, ("'test_not_exported' is mistakenly exported, "
                      "NPY_NO_EXPORT does not work")


# Historically NumPy has not used leading underscores for private submodules
# much.  This has resulted in lots of things that look like public modules
# (i.e. things that can be imported as `import numpy.somesubmodule.somefile`),
# but were never intended to be public.  The PUBLIC_MODULES list contains
# modules that are either public because they were meant to be, or because they
# contain public functions/objects that aren't present in any other namespace
# for whatever reason and therefore should be treated as public.
#
# The PRIVATE_BUT_PRESENT_MODULES list contains modules that look public (lack
# of underscores) but should not be used.  For many of those modules the
# current status is fine.  For others it may make sense to work on making them
# private, to clean up our public API and avoid confusion.
PUBLIC_MODULES = ['numpy.' + s for s in [
    "ctypeslib",
    "distutils",
    "distutils.cpuinfo",
    "distutils.exec_command",
    "distutils.misc_util",
    "distutils.log",
    "distutils.system_info",
    "doc",
    "doc.basics",
    "doc.broadcasting",
    "doc.byteswapping",
    "doc.constants",
    "doc.creation",
    "doc.dispatch",
    "doc.glossary",
    "doc.indexing",
    "doc.internals",
    "doc.misc",
    "doc.structured_arrays",
    "doc.subclassing",
    "doc.ufuncs",
    "dual",
    "f2py",
    "fft",
    "lib",
    "lib.format",  # was this meant to be public?
    "lib.mixins",
    "lib.recfunctions",
    "lib.scimath",
    "linalg",
    "ma",
    "ma.extras",
    "ma.mrecords",
    "matlib",
    "polynomial",
    "polynomial.chebyshev",
    "polynomial.hermite",
    "polynomial.hermite_e",
    "polynomial.laguerre",
    "polynomial.legendre",
    "polynomial.polynomial",
    "polynomial.polyutils",
    "random",
    "testing",
    "version",
]]


PUBLIC_ALIASED_MODULES = [
    "numpy.char",
    "numpy.emath",
    "numpy.rec",
]


PRIVATE_BUT_PRESENT_MODULES = ['numpy.' + s for s in [
    "compat",
    "compat.py3k",
    "conftest",
    "core",
    "core.arrayprint",
    "core.defchararray",
    "core.einsumfunc",
    "core.fromnumeric",
    "core.function_base",
    "core.getlimits",
    "core.machar",
    "core.memmap",
    "core.multiarray",
    "core.numeric",
    "core.numerictypes",
    "core.overrides",
    "core.records",
    "core.shape_base",
    "core.umath",
    "core.umath_tests",
    "distutils.ccompiler",
    "distutils.command",
    "distutils.command.autodist",
    "distutils.command.bdist_rpm",
    "distutils.command.build",
    "distutils.command.build_clib",
    "distutils.command.build_ext",
    "distutils.command.build_py",
    "distutils.command.build_scripts",
    "distutils.command.build_src",
    "distutils.command.config",
    "distutils.command.config_compiler",
    "distutils.command.develop",
    "distutils.command.egg_info",
    "distutils.command.install",
    "distutils.command.install_clib",
    "distutils.command.install_data",
    "distutils.command.install_headers",
    "distutils.command.sdist",
    "distutils.conv_template",
    "distutils.core",
    "distutils.extension",
    "distutils.fcompiler",
    "distutils.fcompiler.absoft",
    "distutils.fcompiler.compaq",
    "distutils.fcompiler.environment",
    "distutils.fcompiler.g95",
    "distutils.fcompiler.gnu",
    "distutils.fcompiler.hpux",
    "distutils.fcompiler.ibm",
    "distutils.fcompiler.intel",
    "distutils.fcompiler.lahey",
    "distutils.fcompiler.mips",
    "distutils.fcompiler.nag",
    "distutils.fcompiler.none",
    "distutils.fcompiler.pathf95",
    "distutils.fcompiler.pg",
    "distutils.fcompiler.nv",
    "distutils.fcompiler.sun",
    "distutils.fcompiler.vast",
    "distutils.from_template",
    "distutils.intelccompiler",
    "distutils.lib2def",
    "distutils.line_endings",
    "distutils.mingw32ccompiler",
    "distutils.msvccompiler",
    "distutils.npy_pkg_config",
    "distutils.numpy_distribution",
    "distutils.pathccompiler",
    "distutils.unixccompiler",
    "f2py.auxfuncs",
    "f2py.capi_maps",
    "f2py.cb_rules",
    "f2py.cfuncs",
    "f2py.common_rules",
    "f2py.crackfortran",
    "f2py.diagnose",
    "f2py.f2py2e",
    "f2py.f2py_testing",
    "f2py.f90mod_rules",
    "f2py.func2subr",
    "f2py.rules",
    "f2py.use_rules",
    "fft.helper",
    "lib.arraypad",
    "lib.arraysetops",
    "lib.arrayterator",
    "lib.financial",
    "lib.function_base",
    "lib.histograms",
    "lib.index_tricks",
    "lib.nanfunctions",
    "lib.npyio",
    "lib.polynomial",
    "lib.shape_base",
    "lib.stride_tricks",
    "lib.twodim_base",
    "lib.type_check",
    "lib.ufunclike",
    "lib.user_array",  # note: not in np.lib, but probably should just be deleted
    "lib.utils",
    "linalg.lapack_lite",
    "linalg.linalg",
    "ma.bench",
    "ma.core",
    "ma.testutils",
    "ma.timer_comparison",
    "matrixlib",
    "matrixlib.defmatrix",
    "random.mtrand",
    "random.bit_generator",
    "testing.print_coercion_tables",
    "testing.utils",
]]


def is_unexpected(name):
    """Check if this needs to be considered."""
    if '._' in name or '.tests' in name or '.setup' in name:
        return False

    if name in PUBLIC_MODULES:
        return False

    if name in PUBLIC_ALIASED_MODULES:
        return False

    if name in PRIVATE_BUT_PRESENT_MODULES:
        return False

    return True


# These are present in a directory with an __init__.py but cannot be imported
# code_generators/ isn't installed, but present for an inplace build
SKIP_LIST = [
    "numpy.core.code_generators",
    "numpy.core.code_generators.genapi",
    "numpy.core.code_generators.generate_umath",
    "numpy.core.code_generators.ufunc_docstrings",
    "numpy.core.code_generators.generate_numpy_api",
    "numpy.core.code_generators.generate_ufunc_api",
    "numpy.core.code_generators.numpy_api",
    "numpy.core.cversions",
    "numpy.core.generate_numpy_api",
    "numpy.distutils.msvc9compiler",
]


def test_all_modules_are_expected():
    """
    Test that we don't add anything that looks like a new public module by
    accident.  Check is based on filenames.
    """

    modnames = []
    for _, modname, ispkg in pkgutil.walk_packages(path=np.__path__,
                                                   prefix=np.__name__ + '.',
                                                   onerror=None):
        if is_unexpected(modname) and modname not in SKIP_LIST:
            # We have a name that is new.  If that's on purpose, add it to
            # PUBLIC_MODULES.  We don't expect to have to add anything to
            # PRIVATE_BUT_PRESENT_MODULES.  Use an underscore in the name!
            modnames.append(modname)

    if modnames:
        raise AssertionError("Found unexpected modules: {}".format(modnames))


# Stuff that clearly shouldn't be in the API and is detected by the next test
# below
SKIP_LIST_2 = [
    'numpy.math',
    'numpy.distutils.log.sys',
    'numpy.distutils.system_info.copy',
    'numpy.distutils.system_info.distutils',
    'numpy.distutils.system_info.log',
    'numpy.distutils.system_info.os',
    'numpy.distutils.system_info.platform',
    'numpy.distutils.system_info.re',
    'numpy.distutils.system_info.shutil',
    'numpy.distutils.system_info.subprocess',
    'numpy.distutils.system_info.sys',
    'numpy.distutils.system_info.tempfile',
    'numpy.distutils.system_info.textwrap',
    'numpy.distutils.system_info.warnings',
    'numpy.doc.constants.re',
    'numpy.doc.constants.textwrap',
    'numpy.lib.emath',
    'numpy.lib.math',
    'numpy.matlib.char',
    'numpy.matlib.rec',
    'numpy.matlib.emath',
    'numpy.matlib.math',
    'numpy.matlib.linalg',
    'numpy.matlib.fft',
    'numpy.matlib.random',
    'numpy.matlib.ctypeslib',
    'numpy.matlib.ma',
]


def test_all_modules_are_expected_2():
    """
    Method checking all objects. The pkgutil-based method in
    `test_all_modules_are_expected` does not catch imports into a namespace,
    only filenames.  So this test is more thorough, and checks this like:

        import .lib.scimath as emath

    To check if something in a module is (effectively) public, one can check if
    there's anything in that namespace that's a public function/object but is
    not exposed in a higher-level namespace.  For example for a `numpy.lib`
    submodule::

        mod = np.lib.mixins
        for obj in mod.__all__:
            if obj in np.__all__:
                continue
            elif obj in np.lib.__all__:
                continue

            else:
                print(obj)

    """

    def find_unexpected_members(mod_name):
        members = []
        module = importlib.import_module(mod_name)
        if hasattr(module, '__all__'):
            objnames = module.__all__
        else:
            objnames = dir(module)

        for objname in objnames:
            if not objname.startswith('_'):
                fullobjname = mod_name + '.' + objname
                if isinstance(getattr(module, objname), types.ModuleType):
                    if is_unexpected(fullobjname):
                        if fullobjname not in SKIP_LIST_2:
                            members.append(fullobjname)

        return members

    unexpected_members = find_unexpected_members("numpy")
    for modname in PUBLIC_MODULES:
        unexpected_members.extend(find_unexpected_members(modname))

    if unexpected_members:
        raise AssertionError("Found unexpected object(s) that look like "
                             "modules: {}".format(unexpected_members))


def test_api_importable():
    """
    Check that all submodules listed higher up in this file can be imported

    Note that if a PRIVATE_BUT_PRESENT_MODULES entry goes missing, it may
    simply need to be removed from the list (deprecation may or may not be
    needed - apply common sense).
    """
    def check_importable(module_name):
        try:
            importlib.import_module(module_name)
        except (ImportError, AttributeError):
            return False

        return True

    module_names = []
    for module_name in PUBLIC_MODULES:
        if not check_importable(module_name):
            module_names.append(module_name)

    if module_names:
        raise AssertionError("Modules in the public API that cannot be "
                             "imported: {}".format(module_names))

    for module_name in PUBLIC_ALIASED_MODULES:
        try:
            eval(module_name)
        except AttributeError:
            module_names.append(module_name)

    if module_names:
        raise AssertionError("Modules in the public API that were not "
                             "found: {}".format(module_names))

    with warnings.catch_warnings(record=True) as w:
        warnings.filterwarnings('always', category=DeprecationWarning)
        warnings.filterwarnings('always', category=ImportWarning)
        for module_name in PRIVATE_BUT_PRESENT_MODULES:
            if not check_importable(module_name):
                module_names.append(module_name)

    if module_names:
        raise AssertionError("Modules that are not really public but looked "
                             "public and can not be imported: "
                             "{}".format(module_names))

social_vdgnn/lib/python3.6/site-packages/numpy/tests/test_public_api.py
eliatrevisan/IA-MPPI-LBM
'''
datetime.tzinfo timezone definitions generated from the
Olson timezone database:

    ftp://elsie.nci.nih.gov/pub/tz*.tar.gz

See the datetime section of the Python Library Reference for information
on how to use these modules.
'''

import sys
import datetime
import os.path

from pytz.exceptions import AmbiguousTimeError
from pytz.exceptions import InvalidTimeError
from pytz.exceptions import NonExistentTimeError
from pytz.exceptions import UnknownTimeZoneError
from pytz.lazy import LazyDict, LazyList, LazySet  # noqa
from pytz.tzinfo import unpickler, BaseTzInfo
from pytz.tzfile import build_tzinfo


# The IANA (nee Olson) database is updated several times a year.
OLSON_VERSION = '2022a'
VERSION = '2022.1'  # pip compatible version number.
__version__ = VERSION

OLSEN_VERSION = OLSON_VERSION  # Old releases had this misspelling

__all__ = [
    'timezone', 'utc', 'country_timezones', 'country_names',
    'AmbiguousTimeError', 'InvalidTimeError',
    'NonExistentTimeError', 'UnknownTimeZoneError',
    'all_timezones', 'all_timezones_set',
    'common_timezones', 'common_timezones_set',
    'BaseTzInfo', 'FixedOffset',
]


if sys.version_info[0] > 2:  # Python 3.x

    # Python 3.x doesn't have unicode(), making writing code
    # for Python 2.3 and Python 3.x a pain.
    unicode = str

    def ascii(s):
        r"""
        >>> ascii('Hello')
        'Hello'
        >>> ascii('\N{TRADE MARK SIGN}') #doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
            ...
        UnicodeEncodeError: ...
        """
        if type(s) == bytes:
            s = s.decode('ASCII')
        else:
            s.encode('ASCII')  # Raise an exception if not ASCII
        return s  # But the string - not a byte string.

else:  # Python 2.x

    def ascii(s):
        r"""
        >>> ascii('Hello')
        'Hello'
        >>> ascii(u'Hello')
        'Hello'
        >>> ascii(u'\N{TRADE MARK SIGN}') #doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
            ...
        UnicodeEncodeError: ...
        """
        return s.encode('ASCII')


def open_resource(name):
    """Open a resource from the zoneinfo subdir for reading.

    Uses the pkg_resources module if available and no standard file
    found at the calculated location.

    It is possible to specify different location for zoneinfo
    subdir by using the PYTZ_TZDATADIR environment variable.
    """
    name_parts = name.lstrip('/').split('/')
    for part in name_parts:
        if part == os.path.pardir or os.path.sep in part:
            raise ValueError('Bad path segment: %r' % part)
    zoneinfo_dir = os.environ.get('PYTZ_TZDATADIR', None)
    if zoneinfo_dir is not None:
        filename = os.path.join(zoneinfo_dir, *name_parts)
    else:
        filename = os.path.join(os.path.dirname(__file__),
                                'zoneinfo', *name_parts)
        if not os.path.exists(filename):
            # http://bugs.launchpad.net/bugs/383171 - we avoid using this
            # unless absolutely necessary to help when a broken version of
            # pkg_resources is installed.
            try:
                from pkg_resources import resource_stream
            except ImportError:
                resource_stream = None

            if resource_stream is not None:
                return resource_stream(__name__, 'zoneinfo/' + name)
    return open(filename, 'rb')


def resource_exists(name):
    """Return true if the given resource exists"""
    try:
        if os.environ.get('PYTZ_SKIPEXISTSCHECK', ''):
            # In "standard" distributions, we can assume that
            # all the listed timezones are present. As an
            # import-speed optimization, you can set the
            # PYTZ_SKIPEXISTSCHECK flag to skip checking
            # for the presence of the resource file on disk.
            return True
        open_resource(name).close()
        return True
    except IOError:
        return False


_tzinfo_cache = {}


def timezone(zone):
    r''' Return a datetime.tzinfo implementation for the given timezone

    >>> from datetime import datetime, timedelta
    >>> utc = timezone('UTC')
    >>> eastern = timezone('US/Eastern')
    >>> eastern.zone
    'US/Eastern'
    >>> timezone(unicode('US/Eastern')) is eastern
    True
    >>> utc_dt = datetime(2002, 10, 27, 6, 0, 0, tzinfo=utc)
    >>> loc_dt = utc_dt.astimezone(eastern)
    >>> fmt = '%Y-%m-%d %H:%M:%S %Z (%z)'
    >>> loc_dt.strftime(fmt)
    '2002-10-27 01:00:00 EST (-0500)'
    >>> (loc_dt - timedelta(minutes=10)).strftime(fmt)
    '2002-10-27 00:50:00 EST (-0500)'
    >>> eastern.normalize(loc_dt - timedelta(minutes=10)).strftime(fmt)
    '2002-10-27 01:50:00 EDT (-0400)'
    >>> (loc_dt + timedelta(minutes=10)).strftime(fmt)
    '2002-10-27 01:10:00 EST (-0500)'

    Raises UnknownTimeZoneError if passed an unknown zone.

    >>> try:
    ...     timezone('Asia/Shangri-La')
    ... except UnknownTimeZoneError:
    ...     print('Unknown')
    Unknown

    >>> try:
    ...     timezone(unicode('\N{TRADE MARK SIGN}'))
    ... except UnknownTimeZoneError:
    ...     print('Unknown')
    Unknown

    '''
    if zone is None:
        raise UnknownTimeZoneError(None)

    if zone.upper() == 'UTC':
        return utc

    try:
        zone = ascii(zone)
    except UnicodeEncodeError:
        # All valid timezones are ASCII
        raise UnknownTimeZoneError(zone)

    zone = _case_insensitive_zone_lookup(_unmunge_zone(zone))
    if zone not in _tzinfo_cache:
        if zone in all_timezones_set:  # noqa
            fp = open_resource(zone)
            try:
                _tzinfo_cache[zone] = build_tzinfo(zone, fp)
            finally:
                fp.close()
        else:
            raise UnknownTimeZoneError(zone)

    return _tzinfo_cache[zone]


def _unmunge_zone(zone):
    """Undo the time zone name munging done by older versions of pytz."""
    return zone.replace('_plus_', '+').replace('_minus_', '-')


_all_timezones_lower_to_standard = None


def _case_insensitive_zone_lookup(zone):
    """case-insensitively matching timezone, else return zone unchanged"""
    global _all_timezones_lower_to_standard
    if _all_timezones_lower_to_standard is None:
        _all_timezones_lower_to_standard = dict((tz.lower(), tz) for tz in all_timezones)  # noqa
    return _all_timezones_lower_to_standard.get(zone.lower()) or zone  # noqa


ZERO = datetime.timedelta(0)
HOUR = datetime.timedelta(hours=1)


class UTC(BaseTzInfo):
    """UTC

    Optimized UTC implementation. It unpickles using the single module global
    instance defined beneath this class declaration.
    """
    zone = "UTC"

    _utcoffset = ZERO
    _dst = ZERO
    _tzname = zone

    def fromutc(self, dt):
        if dt.tzinfo is None:
            return self.localize(dt)
        return super(utc.__class__, self).fromutc(dt)

    def utcoffset(self, dt):
        return ZERO

    def tzname(self, dt):
        return "UTC"

    def dst(self, dt):
        return ZERO

    def __reduce__(self):
        return _UTC, ()

    def localize(self, dt, is_dst=False):
        '''Convert naive time to local time'''
        if dt.tzinfo is not None:
            raise ValueError('Not naive datetime (tzinfo is already set)')
        return dt.replace(tzinfo=self)

    def normalize(self, dt, is_dst=False):
        '''Correct the timezone information on the given datetime'''
        if dt.tzinfo is self:
            return dt
        if dt.tzinfo is None:
            raise ValueError('Naive time - no tzinfo set')
        return dt.astimezone(self)

    def __repr__(self):
        return "<UTC>"

    def __str__(self):
        return "UTC"


UTC = utc = UTC()  # UTC is a singleton


def _UTC():
    """Factory function for utc unpickling.

    Makes sure that unpickling a utc instance always returns the same
    module global.

    These examples belong in the UTC class above, but it is obscured; or in
    the README.rst, but we are not depending on Python 2.4 so integrating
    the README.rst examples with the unit tests is not trivial.

    >>> import datetime, pickle
    >>> dt = datetime.datetime(2005, 3, 1, 14, 13, 21, tzinfo=utc)
    >>> naive = dt.replace(tzinfo=None)
    >>> p = pickle.dumps(dt, 1)
    >>> naive_p = pickle.dumps(naive, 1)
    >>> len(p) - len(naive_p)
    17
    >>> new = pickle.loads(p)
    >>> new == dt
    True
    >>> new is dt
    False
    >>> new.tzinfo is dt.tzinfo
    True
    >>> utc is UTC is timezone('UTC')
    True
    >>> utc is timezone('GMT')
    False
    """
    return utc


_UTC.__safe_for_unpickling__ = True


def _p(*args):
    """Factory function for unpickling pytz tzinfo instances.

    Just a wrapper around tzinfo.unpickler to save a few bytes in each pickle
    by shortening the path.
    """
    return unpickler(*args)


_p.__safe_for_unpickling__ = True


class _CountryTimezoneDict(LazyDict):
    """Map ISO 3166 country code to a list of timezone names commonly used
    in that country.

    iso3166_code is the two letter code used to identify the country.

    >>> def print_list(list_of_strings):
    ...     'We use a helper so doctests work under Python 2.3 -> 3.x'
    ...     for s in list_of_strings:
    ...         print(s)

    >>> print_list(country_timezones['nz'])
    Pacific/Auckland
    Pacific/Chatham
    >>> print_list(country_timezones['ch'])
    Europe/Zurich
    >>> print_list(country_timezones['CH'])
    Europe/Zurich
    >>> print_list(country_timezones[unicode('ch')])
    Europe/Zurich
    >>> print_list(country_timezones['XXX'])
    Traceback (most recent call last):
    ...
    KeyError: 'XXX'

    Previously, this information was exposed as a function rather than a
    dictionary. This is still supported::

    >>> print_list(country_timezones('nz'))
    Pacific/Auckland
    Pacific/Chatham
    """
    def __call__(self, iso3166_code):
        """Backwards compatibility."""
        return self[iso3166_code]

    def _fill(self):
        data = {}
        zone_tab = open_resource('zone.tab')
        try:
            for line in zone_tab:
                line = line.decode('UTF-8')
                if line.startswith('#'):
                    continue
                code, coordinates, zone = line.split(None, 4)[:3]
                if zone not in all_timezones_set:  # noqa
                    continue
                try:
                    data[code].append(zone)
                except KeyError:
                    data[code] = [zone]
            self.data = data
        finally:
            zone_tab.close()


country_timezones = _CountryTimezoneDict()


class _CountryNameDict(LazyDict):
    '''Dictionary proving ISO3166 code -> English name.

    >>> print(country_names['au'])
    Australia
    '''
    def _fill(self):
        data = {}
        zone_tab = open_resource('iso3166.tab')
        try:
            for line in zone_tab.readlines():
                line = line.decode('UTF-8')
                if line.startswith('#'):
                    continue
                code, name = line.split(None, 1)
                data[code] = name.strip()
            self.data = data
        finally:
            zone_tab.close()


country_names = _CountryNameDict()


# Time-zone info based solely on fixed offsets

class _FixedOffset(datetime.tzinfo):

    zone = None  # to match the standard pytz API

    def __init__(self, minutes):
        if abs(minutes) >= 1440:
            raise ValueError("absolute offset is too large", minutes)
        self._minutes = minutes
        self._offset = datetime.timedelta(minutes=minutes)

    def utcoffset(self, dt):
        return self._offset

    def __reduce__(self):
        return FixedOffset, (self._minutes, )

    def dst(self, dt):
        return ZERO

    def tzname(self, dt):
        return None

    def __repr__(self):
        return 'pytz.FixedOffset(%d)' % self._minutes

    def localize(self, dt, is_dst=False):
        '''Convert naive time to local time'''
        if dt.tzinfo is not None:
            raise ValueError('Not naive datetime (tzinfo is already set)')
        return dt.replace(tzinfo=self)

    def normalize(self, dt, is_dst=False):
        '''Correct the timezone information on the given datetime'''
        if dt.tzinfo is self:
            return dt
        if dt.tzinfo is None:
            raise ValueError('Naive time - no tzinfo set')
        return dt.astimezone(self)


def FixedOffset(offset, _tzinfos={}):
    """return a fixed-offset timezone based off a number of minutes.

        >>> one = FixedOffset(-330)
        >>> one
        pytz.FixedOffset(-330)
        >>> str(one.utcoffset(datetime.datetime.now()))
        '-1 day, 18:30:00'
        >>> str(one.dst(datetime.datetime.now()))
        '0:00:00'

        >>> two = FixedOffset(1380)
        >>> two
        pytz.FixedOffset(1380)
        >>> str(two.utcoffset(datetime.datetime.now()))
        '23:00:00'
        >>> str(two.dst(datetime.datetime.now()))
        '0:00:00'

    The datetime.timedelta must be between the range of -1 and 1 day,
    non-inclusive.

        >>> FixedOffset(1440)
        Traceback (most recent call last):
        ...
        ValueError: ('absolute offset is too large', 1440)

        >>> FixedOffset(-1440)
        Traceback (most recent call last):
        ...
        ValueError: ('absolute offset is too large', -1440)

    An offset of 0 is special-cased to return UTC.

        >>> FixedOffset(0) is UTC
        True

    There should always be only one instance of a FixedOffset per timedelta.
    This should be true for multiple creation calls.

        >>> FixedOffset(-330) is one
        True
        >>> FixedOffset(1380) is two
        True

    It should also be true for pickling.

        >>> import pickle
        >>> pickle.loads(pickle.dumps(one)) is one
        True
        >>> pickle.loads(pickle.dumps(two)) is two
        True
    """
    if offset == 0:
        return UTC

    info = _tzinfos.get(offset)
    if info is None:
        # We haven't seen this one before. we need to save it.

        # Use setdefault to avoid a race condition and make sure we have
        # only one
        info = _tzinfos.setdefault(offset, _FixedOffset(offset))

    return info


FixedOffset.__safe_for_unpickling__ = True


def _test():
    import doctest
    sys.path.insert(0, os.pardir)
    import pytz
    return doctest.testmod(pytz)


if __name__ == '__main__':
    _test()
all_timezones = \
['Africa/Abidjan',
 'Africa/Accra',
 'Africa/Addis_Ababa',
 'Africa/Algiers',
 'Africa/Asmara',
 'Africa/Asmera',
 'Africa/Bamako',
 'Africa/Bangui',
 'Africa/Banjul',
 'Africa/Bissau',
 'Africa/Blantyre',
 'Africa/Brazzaville',
 'Africa/Bujumbura',
 'Africa/Cairo',
 'Africa/Casablanca',
 'Africa/Ceuta',
 'Africa/Conakry',
 'Africa/Dakar',
 'Africa/Dar_es_Salaam',
 'Africa/Djibouti',
 'Africa/Douala',
 'Africa/El_Aaiun',
 'Africa/Freetown',
 'Africa/Gaborone',
 'Africa/Harare',
 'Africa/Johannesburg',
 'Africa/Juba',
 'Africa/Kampala',
 'Africa/Khartoum',
 'Africa/Kigali',
 'Africa/Kinshasa',
 'Africa/Lagos',
 'Africa/Libreville',
 'Africa/Lome',
 'Africa/Luanda',
 'Africa/Lubumbashi',
 'Africa/Lusaka',
 'Africa/Malabo',
 'Africa/Maputo',
 'Africa/Maseru',
 'Africa/Mbabane',
 'Africa/Mogadishu',
 'Africa/Monrovia',
 'Africa/Nairobi',
 'Africa/Ndjamena',
 'Africa/Niamey',
 'Africa/Nouakchott',
 'Africa/Ouagadougou',
 'Africa/Porto-Novo',
 'Africa/Sao_Tome',
 'Africa/Timbuktu',
 'Africa/Tripoli',
 'Africa/Tunis',
 'Africa/Windhoek',
 'America/Adak',
 'America/Anchorage',
 'America/Anguilla',
 'America/Antigua',
 'America/Araguaina',
 'America/Argentina/Buenos_Aires',
 'America/Argentina/Catamarca',
 'America/Argentina/ComodRivadavia',
 'America/Argentina/Cordoba',
 'America/Argentina/Jujuy',
 'America/Argentina/La_Rioja',
 'America/Argentina/Mendoza',
 'America/Argentina/Rio_Gallegos',
 'America/Argentina/Salta',
 'America/Argentina/San_Juan',
 'America/Argentina/San_Luis',
 'America/Argentina/Tucuman',
 'America/Argentina/Ushuaia',
 'America/Aruba',
 'America/Asuncion',
 'America/Atikokan',
 'America/Atka',
 'America/Bahia',
 'America/Bahia_Banderas',
 'America/Barbados',
 'America/Belem',
 'America/Belize',
 'America/Blanc-Sablon',
 'America/Boa_Vista',
 'America/Bogota',
 'America/Boise',
 'America/Buenos_Aires',
 'America/Cambridge_Bay',
 'America/Campo_Grande',
 'America/Cancun',
 'America/Caracas',
 'America/Catamarca',
 'America/Cayenne',
 'America/Cayman',
 'America/Chicago',
 'America/Chihuahua',
 'America/Coral_Harbour',
 'America/Cordoba',
 'America/Costa_Rica',
 'America/Creston',
 'America/Cuiaba',
 'America/Curacao',
 'America/Danmarkshavn',
 'America/Dawson',
 'America/Dawson_Creek',
 'America/Denver',
 'America/Detroit',
 'America/Dominica',
 'America/Edmonton',
 'America/Eirunepe',
 'America/El_Salvador',
 'America/Ensenada',
 'America/Fort_Nelson',
 'America/Fort_Wayne',
 'America/Fortaleza',
 'America/Glace_Bay',
 'America/Godthab',
 'America/Goose_Bay',
 'America/Grand_Turk',
 'America/Grenada',
 'America/Guadeloupe',
 'America/Guatemala',
 'America/Guayaquil',
 'America/Guyana',
 'America/Halifax',
 'America/Havana',
 'America/Hermosillo',
 'America/Indiana/Indianapolis',
 'America/Indiana/Knox',
 'America/Indiana/Marengo',
 'America/Indiana/Petersburg',
 'America/Indiana/Tell_City',
 'America/Indiana/Vevay',
 'America/Indiana/Vincennes',
 'America/Indiana/Winamac',
 'America/Indianapolis',
 'America/Inuvik',
 'America/Iqaluit',
 'America/Jamaica',
 'America/Jujuy',
 'America/Juneau',
 'America/Kentucky/Louisville',
 'America/Kentucky/Monticello',
 'America/Knox_IN',
 'America/Kralendijk',
 'America/La_Paz',
 'America/Lima',
 'America/Los_Angeles',
 'America/Louisville',
 'America/Lower_Princes',
 'America/Maceio',
 'America/Managua',
 'America/Manaus',
 'America/Marigot',
 'America/Martinique',
 'America/Matamoros',
 'America/Mazatlan',
 'America/Mendoza',
 'America/Menominee',
 'America/Merida',
 'America/Metlakatla',
 'America/Mexico_City',
 'America/Miquelon',
 'America/Moncton',
 'America/Monterrey',
 'America/Montevideo',
 'America/Montreal',
 'America/Montserrat',
 'America/Nassau',
 'America/New_York',
 'America/Nipigon',
 'America/Nome',
 'America/Noronha',
 'America/North_Dakota/Beulah',
 'America/North_Dakota/Center',
 'America/North_Dakota/New_Salem',
 'America/Nuuk',
 'America/Ojinaga',
 'America/Panama',
 'America/Pangnirtung',
 'America/Paramaribo',
 'America/Phoenix',
 'America/Port-au-Prince',
 'America/Port_of_Spain',
 'America/Porto_Acre',
 'America/Porto_Velho',
 'America/Puerto_Rico',
 'America/Punta_Arenas',
 'America/Rainy_River',
 'America/Rankin_Inlet',
 'America/Recife',
 'America/Regina',
 'America/Resolute',
 'America/Rio_Branco',
 'America/Rosario',
 'America/Santa_Isabel',
 'America/Santarem',
 'America/Santiago',
 'America/Santo_Domingo',
 'America/Sao_Paulo',
 'America/Scoresbysund',
 'America/Shiprock',
 'America/Sitka',
 'America/St_Barthelemy',
 'America/St_Johns',
 'America/St_Kitts',
 'America/St_Lucia',
 'America/St_Thomas',
 'America/St_Vincent',
 'America/Swift_Current',
 'America/Tegucigalpa',
 'America/Thule',
 'America/Thunder_Bay',
 'America/Tijuana',
 'America/Toronto',
 'America/Tortola',
 'America/Vancouver',
 'America/Virgin',
 'America/Whitehorse',
 'America/Winnipeg',
 'America/Yakutat',
 'America/Yellowknife',
 'Antarctica/Casey',
 'Antarctica/Davis',
 'Antarctica/DumontDUrville',
 'Antarctica/Macquarie',
 'Antarctica/Mawson',
 'Antarctica/McMurdo',
 'Antarctica/Palmer',
 'Antarctica/Rothera',
 'Antarctica/South_Pole',
 'Antarctica/Syowa',
 'Antarctica/Troll',
 'Antarctica/Vostok',
 'Arctic/Longyearbyen',
 'Asia/Aden',
 'Asia/Almaty',
 'Asia/Amman',
 'Asia/Anadyr',
 'Asia/Aqtau',
 'Asia/Aqtobe',
 'Asia/Ashgabat',
 'Asia/Ashkhabad',
 'Asia/Atyrau',
 'Asia/Baghdad',
 'Asia/Bahrain',
 'Asia/Baku',
 'Asia/Bangkok',
 'Asia/Barnaul',
 'Asia/Beirut',
 'Asia/Bishkek',
 'Asia/Brunei',
 'Asia/Calcutta',
 'Asia/Chita',
 'Asia/Choibalsan',
 'Asia/Chongqing',
 'Asia/Chungking',
 'Asia/Colombo',
 'Asia/Dacca',
 'Asia/Damascus',
 'Asia/Dhaka',
 'Asia/Dili',
 'Asia/Dubai',
 'Asia/Dushanbe',
 'Asia/Famagusta',
 'Asia/Gaza',
 'Asia/Harbin',
 'Asia/Hebron',
 'Asia/Ho_Chi_Minh',
 'Asia/Hong_Kong',
 'Asia/Hovd',
 'Asia/Irkutsk',
 'Asia/Istanbul',
 'Asia/Jakarta',
 'Asia/Jayapura',
 'Asia/Jerusalem',
 'Asia/Kabul',
 'Asia/Kamchatka',
 'Asia/Karachi',
 'Asia/Kashgar',
 'Asia/Kathmandu',
 'Asia/Katmandu',
 'Asia/Khandyga',
 'Asia/Kolkata',
 'Asia/Krasnoyarsk',
 'Asia/Kuala_Lumpur',
 'Asia/Kuching',
 'Asia/Kuwait',
 'Asia/Macao',
 'Asia/Macau',
 'Asia/Magadan',
 'Asia/Makassar',
 'Asia/Manila',
 'Asia/Muscat',
 'Asia/Nicosia',
 'Asia/Novokuznetsk',
 'Asia/Novosibirsk',
 'Asia/Omsk',
 'Asia/Oral',
 'Asia/Phnom_Penh',
 'Asia/Pontianak',
 'Asia/Pyongyang',
 'Asia/Qatar',
 'Asia/Qostanay',
 'Asia/Qyzylorda',
 'Asia/Rangoon',
 'Asia/Riyadh',
 'Asia/Saigon',
 'Asia/Sakhalin',
 'Asia/Samarkand',
 'Asia/Seoul',
 'Asia/Shanghai',
 'Asia/Singapore',
 'Asia/Srednekolymsk',
 'Asia/Taipei',
 'Asia/Tashkent',
 'Asia/Tbilisi',
 'Asia/Tehran',
 'Asia/Tel_Aviv',
 'Asia/Thimbu',
 'Asia/Thimphu',
 'Asia/Tokyo',
 'Asia/Tomsk',
 'Asia/Ujung_Pandang',
 'Asia/Ulaanbaatar',
 'Asia/Ulan_Bator',
 'Asia/Urumqi',
 'Asia/Ust-Nera',
 'Asia/Vientiane',
 'Asia/Vladivostok',
 'Asia/Yakutsk',
 'Asia/Yangon',
 'Asia/Yekaterinburg',
 'Asia/Yerevan',
 'Atlantic/Azores',
 'Atlantic/Bermuda',
 'Atlantic/Canary',
 'Atlantic/Cape_Verde',
 'Atlantic/Faeroe',
 'Atlantic/Faroe',
 'Atlantic/Jan_Mayen',
 'Atlantic/Madeira',
 'Atlantic/Reykjavik',
 'Atlantic/South_Georgia',
 'Atlantic/St_Helena',
 'Atlantic/Stanley',
 'Australia/ACT',
 'Australia/Adelaide',
 'Australia/Brisbane',
 'Australia/Broken_Hill',
 'Australia/Canberra',
 'Australia/Currie',
 'Australia/Darwin',
 'Australia/Eucla',
 'Australia/Hobart',
 'Australia/LHI',
 'Australia/Lindeman',
 'Australia/Lord_Howe',
 'Australia/Melbourne',
 'Australia/NSW',
 'Australia/North',
 'Australia/Perth',
 'Australia/Queensland',
 'Australia/South',
 'Australia/Sydney',
 'Australia/Tasmania',
 'Australia/Victoria',
 'Australia/West',
 'Australia/Yancowinna',
 'Brazil/Acre',
 'Brazil/DeNoronha',
 'Brazil/East',
 'Brazil/West',
 'CET',
 'CST6CDT',
 'Canada/Atlantic',
 'Canada/Central',
 'Canada/Eastern',
 'Canada/Mountain',
 'Canada/Newfoundland',
 'Canada/Pacific',
 'Canada/Saskatchewan',
 'Canada/Yukon',
 'Chile/Continental',
 'Chile/EasterIsland',
 'Cuba',
 'EET',
 'EST',
 'EST5EDT',
 'Egypt',
 'Eire',
 'Etc/GMT',
 'Etc/GMT+0',
 'Etc/GMT+1',
 'Etc/GMT+10',
 'Etc/GMT+11',
 'Etc/GMT+12',
 'Etc/GMT+2',
 'Etc/GMT+3',
 'Etc/GMT+4',
 'Etc/GMT+5',
 'Etc/GMT+6',
 'Etc/GMT+7',
 'Etc/GMT+8',
 'Etc/GMT+9',
 'Etc/GMT-0',
 'Etc/GMT-1',
 'Etc/GMT-10',
 'Etc/GMT-11',
 'Etc/GMT-12',
 'Etc/GMT-13',
 'Etc/GMT-14',
 'Etc/GMT-2',
 'Etc/GMT-3',
 'Etc/GMT-4',
 'Etc/GMT-5',
 'Etc/GMT-6',
 'Etc/GMT-7',
 'Etc/GMT-8',
 'Etc/GMT-9',
 'Etc/GMT0',
 'Etc/Greenwich',
 'Etc/UCT',
 'Etc/UTC',
 'Etc/Universal',
 'Etc/Zulu',
 'Europe/Amsterdam',
 'Europe/Andorra',
 'Europe/Astrakhan',
 'Europe/Athens',
 'Europe/Belfast',
 'Europe/Belgrade',
 'Europe/Berlin',
 'Europe/Bratislava',
 'Europe/Brussels',
 'Europe/Bucharest',
 'Europe/Budapest',
 'Europe/Busingen',
 'Europe/Chisinau',
 'Europe/Copenhagen',
 'Europe/Dublin',
 'Europe/Gibraltar',
 'Europe/Guernsey',
 'Europe/Helsinki',
 'Europe/Isle_of_Man',
 'Europe/Istanbul',
 'Europe/Jersey',
 'Europe/Kaliningrad',
 'Europe/Kiev',
 'Europe/Kirov',
 'Europe/Lisbon',
 'Europe/Ljubljana',
 'Europe/London',
 'Europe/Luxembourg',
 'Europe/Madrid',
 'Europe/Malta',
 'Europe/Mariehamn',
 'Europe/Minsk',
 'Europe/Monaco',
 'Europe/Moscow',
 'Europe/Nicosia',
 'Europe/Oslo',
 'Europe/Paris',
 'Europe/Podgorica',
 'Europe/Prague',
 'Europe/Riga',
 'Europe/Rome',
 'Europe/Samara',
 'Europe/San_Marino',
 'Europe/Sarajevo',
 'Europe/Saratov',
 'Europe/Simferopol',
 'Europe/Skopje',
 'Europe/Sofia',
 'Europe/Stockholm',
 'Europe/Tallinn',
 'Europe/Tirane',
 'Europe/Tiraspol',
 'Europe/Ulyanovsk',
 'Europe/Uzhgorod',
 'Europe/Vaduz',
 'Europe/Vatican',
 'Europe/Vienna',
 'Europe/Vilnius',
 'Europe/Volgograd',
 'Europe/Warsaw',
 'Europe/Zagreb',
 'Europe/Zaporozhye',
 'Europe/Zurich',
 'GB',
 'GB-Eire',
 'GMT',
 'GMT+0',
 'GMT-0',
 'GMT0',
 'Greenwich',
 'HST',
 'Hongkong',
 'Iceland',
 'Indian/Antananarivo',
 'Indian/Chagos',
 'Indian/Christmas',
 'Indian/Cocos',
 'Indian/Comoro',
 'Indian/Kerguelen',
 'Indian/Mahe',
 'Indian/Maldives',
 'Indian/Mauritius',
 'Indian/Mayotte',
 'Indian/Reunion',
 'Iran',
 'Israel',
 'Jamaica',
 'Japan',
 'Kwajalein',
 'Libya',
 'MET',
 'MST',
 'MST7MDT',
 'Mexico/BajaNorte',
 'Mexico/BajaSur',
 'Mexico/General',
 'NZ',
 'NZ-CHAT',
 'Navajo',
 'PRC',
 'PST8PDT',
 'Pacific/Apia',
 'Pacific/Auckland',
 'Pacific/Bougainville',
 'Pacific/Chatham',
 'Pacific/Chuuk',
 'Pacific/Easter',
 'Pacific/Efate',
 'Pacific/Enderbury',
 'Pacific/Fakaofo',
 'Pacific/Fiji',
 'Pacific/Funafuti',
 'Pacific/Galapagos',
 'Pacific/Gambier',
 'Pacific/Guadalcanal',
 'Pacific/Guam',
 'Pacific/Honolulu',
 'Pacific/Johnston',
 'Pacific/Kanton',
 'Pacific/Kiritimati',
 'Pacific/Kosrae',
 'Pacific/Kwajalein',
 'Pacific/Majuro',
 'Pacific/Marquesas',
 'Pacific/Midway',
 'Pacific/Nauru',
 'Pacific/Niue',
 'Pacific/Norfolk',
 'Pacific/Noumea',
 'Pacific/Pago_Pago',
 'Pacific/Palau',
 'Pacific/Pitcairn',
 'Pacific/Pohnpei',
 'Pacific/Ponape',
 'Pacific/Port_Moresby',
 'Pacific/Rarotonga',
 'Pacific/Saipan',
 'Pacific/Samoa',
 'Pacific/Tahiti',
 'Pacific/Tarawa',
 'Pacific/Tongatapu',
 'Pacific/Truk',
 'Pacific/Wake',
 'Pacific/Wallis',
 'Pacific/Yap',
 'Poland',
 'Portugal',
 'ROC',
 'ROK',
 'Singapore',
 'Turkey',
 'UCT',
 'US/Alaska',
 'US/Aleutian',
 'US/Arizona',
 'US/Central',
 'US/East-Indiana',
 'US/Eastern',
 'US/Hawaii',
 'US/Indiana-Starke',
 'US/Michigan',
 'US/Mountain',
 'US/Pacific',
 'US/Samoa',
 'UTC',
 'Universal',
 'W-SU',
 'WET',
 'Zulu']
all_timezones = LazyList(
        tz for tz in all_timezones if resource_exists(tz))
        
all_timezones_set = LazySet(all_timezones)
common_timezones = \
['Africa/Abidjan',
 'Africa/Accra',
 'Africa/Addis_Ababa',
 'Africa/Algiers',
 'Africa/Asmara',
 'Africa/Bamako',
 'Africa/Bangui',
 'Africa/Banjul',
 'Africa/Bissau',
 'Africa/Blantyre',
 'Africa/Brazzaville',
 'Africa/Bujumbura',
 'Africa/Cairo',
 'Africa/Casablanca',
 'Africa/Ceuta',
 'Africa/Conakry',
 'Africa/Dakar',
 'Africa/Dar_es_Salaam',
 'Africa/Djibouti',
 'Africa/Douala',
 'Africa/El_Aaiun',
 'Africa/Freetown',
 'Africa/Gaborone',
 'Africa/Harare',
 'Africa/Johannesburg',
 'Africa/Juba',
 'Africa/Kampala',
 'Africa/Khartoum',
 'Africa/Kigali',
 'Africa/Kinshasa',
 'Africa/Lagos',
 'Africa/Libreville',
 'Africa/Lome',
 'Africa/Luanda',
 'Africa/Lubumbashi',
 'Africa/Lusaka',
 'Africa/Malabo',
 'Africa/Maputo',
 'Africa/Maseru',
 'Africa/Mbabane',
 'Africa/Mogadishu',
 'Africa/Monrovia',
 'Africa/Nairobi',
 'Africa/Ndjamena',
 'Africa/Niamey',
 'Africa/Nouakchott',
 'Africa/Ouagadougou',
 'Africa/Porto-Novo',
 'Africa/Sao_Tome',
 'Africa/Tripoli',
 'Africa/Tunis',
 'Africa/Windhoek',
 'America/Adak',
 'America/Anchorage',
 'America/Anguilla',
 'America/Antigua',
 'America/Araguaina',
 'America/Argentina/Buenos_Aires',
 'America/Argentina/Catamarca',
 'America/Argentina/Cordoba',
 'America/Argentina/Jujuy',
 'America/Argentina/La_Rioja',
 'America/Argentina/Mendoza',
 'America/Argentina/Rio_Gallegos',
 'America/Argentina/Salta',
 'America/Argentina/San_Juan',
 'America/Argentina/San_Luis',
 'America/Argentina/Tucuman',
 'America/Argentina/Ushuaia',
 'America/Aruba',
 'America/Asuncion',
 'America/Atikokan',
 'America/Bahia',
 'America/Bahia_Banderas',
 'America/Barbados',
 'America/Belem',
 'America/Belize',
 'America/Blanc-Sablon',
 'America/Boa_Vista',
 'America/Bogota',
 'America/Boise',
 'America/Cambridge_Bay',
 'America/Campo_Grande',
 'America/Cancun',
 'America/Caracas',
 'America/Cayenne',
 'America/Cayman',
 'America/Chicago',
 'America/Chihuahua',
 'America/Costa_Rica',
 'America/Creston',
 'America/Cuiaba',
 'America/Curacao',
 'America/Danmarkshavn',
 'America/Dawson',
 'America/Dawson_Creek',
 'America/Denver',
 'America/Detroit',
 'America/Dominica',
 'America/Edmonton',
 'America/Eirunepe',
 'America/El_Salvador',
 'America/Fort_Nelson',
 'America/Fortaleza',
 'America/Glace_Bay',
 'America/Goose_Bay',
 'America/Grand_Turk',
 'America/Grenada',
 'America/Guadeloupe',
 'America/Guatemala',
 'America/Guayaquil',
 'America/Guyana',
 'America/Halifax',
 'America/Havana',
 'America/Hermosillo',
 'America/Indiana/Indianapolis',
 'America/Indiana/Knox',
 'America/Indiana/Marengo',
 'America/Indiana/Petersburg',
 'America/Indiana/Tell_City',
 'America/Indiana/Vevay',
 'America/Indiana/Vincennes',
 'America/Indiana/Winamac',
 'America/Inuvik',
 'America/Iqaluit',
 'America/Jamaica',
 'America/Juneau',
 'America/Kentucky/Louisville',
 'America/Kentucky/Monticello',
 'America/Kralendijk',
 'America/La_Paz',
 'America/Lima',
 'America/Los_Angeles',
 'America/Lower_Princes',
 'America/Maceio',
 'America/Managua',
 'America/Manaus',
 'America/Marigot',
 'America/Martinique',
 'America/Matamoros',
 'America/Mazatlan',
 'America/Menominee',
 'America/Merida',
 'America/Metlakatla',
 'America/Mexico_City',
 'America/Miquelon',
 'America/Moncton',
 'America/Monterrey',
 'America/Montevideo',
 'America/Montserrat',
 'America/Nassau',
 'America/New_York',
 'America/Nipigon',
 'America/Nome',
 'America/Noronha',
 'America/North_Dakota/Beulah',
 'America/North_Dakota/Center',
 'America/North_Dakota/New_Salem',
 'America/Nuuk',
 'America/Ojinaga',
 'America/Panama',
 'America/Pangnirtung',
 'America/Paramaribo',
 'America/Phoenix',
 'America/Port-au-Prince',
 'America/Port_of_Spain',
 'America/Porto_Velho',
 'America/Puerto_Rico',
 'America/Punta_Arenas',
 'America/Rainy_River',
 'America/Rankin_Inlet',
 'America/Recife',
 'America/Regina',
 'America/Resolute',
 'America/Rio_Branco',
 'America/Santarem',
 'America/Santiago',
 'America/Santo_Domingo',
 'America/Sao_Paulo',
 'America/Scoresbysund',
 'America/Sitka',
 'America/St_Barthelemy',
 'America/St_Johns',
 'America/St_Kitts',
 'America/St_Lucia',
 'America/St_Thomas',
 'America/St_Vincent',
 'America/Swift_Current',
 'America/Tegucigalpa',
 'America/Thule',
 'America/Thunder_Bay',
 'America/Tijuana',
 'America/Toronto',
 'America/Tortola',
 'America/Vancouver',
 'America/Whitehorse',
 'America/Winnipeg',
 'America/Yakutat',
 'America/Yellowknife',
 'Antarctica/Casey',
 'Antarctica/Davis',
 'Antarctica/DumontDUrville',
 'Antarctica/Macquarie',
 'Antarctica/Mawson',
 'Antarctica/McMurdo',
 'Antarctica/Palmer',
 'Antarctica/Rothera',
 'Antarctica/Syowa',
 'Antarctica/Troll',
 'Antarctica/Vostok',
 'Arctic/Longyearbyen',
 'Asia/Aden',
 'Asia/Almaty',
 'Asia/Amman',
 'Asia/Anadyr',
 'Asia/Aqtau',
 'Asia/Aqtobe',
 'Asia/Ashgabat',
 'Asia/Atyrau',
 'Asia/Baghdad',
 'Asia/Bahrain',
 'Asia/Baku',
 'Asia/Bangkok',
 'Asia/Barnaul',
 'Asia/Beirut',
 'Asia/Bishkek',
 'Asia/Brunei',
 'Asia/Chita',
 'Asia/Choibalsan',
 'Asia/Colombo',
 'Asia/Damascus',
 'Asia/Dhaka',
 'Asia/Dili',
 'Asia/Dubai',
 'Asia/Dushanbe',
 'Asia/Famagusta',
 'Asia/Gaza',
 'Asia/Hebron',
 'Asia/Ho_Chi_Minh',
 'Asia/Hong_Kong',
 'Asia/Hovd',
 'Asia/Irkutsk',
 'Asia/Jakarta',
 'Asia/Jayapura',
 'Asia/Jerusalem',
 'Asia/Kabul',
 'Asia/Kamchatka',
 'Asia/Karachi',
 'Asia/Kathmandu',
 'Asia/Khandyga',
 'Asia/Kolkata',
 'Asia/Krasnoyarsk',
 'Asia/Kuala_Lumpur',
 'Asia/Kuching',
 'Asia/Kuwait',
 'Asia/Macau',
 'Asia/Magadan',
 'Asia/Makassar',
 'Asia/Manila',
 'Asia/Muscat',
 'Asia/Nicosia',
 'Asia/Novokuznetsk',
 'Asia/Novosibirsk',
 'Asia/Omsk',
 'Asia/Oral',
 'Asia/Phnom_Penh',
 'Asia/Pontianak',
 'Asia/Pyongyang',
 'Asia/Qatar',
 'Asia/Qostanay',
 'Asia/Qyzylorda',
 'Asia/Riyadh',
 'Asia/Sakhalin',
 'Asia/Samarkand',
 'Asia/Seoul',
 'Asia/Shanghai',
 'Asia/Singapore',
 'Asia/Srednekolymsk',
 'Asia/Taipei',
 'Asia/Tashkent',
 'Asia/Tbilisi',
 'Asia/Tehran',
 'Asia/Thimphu',
 'Asia/Tokyo',
 'Asia/Tomsk',
 'Asia/Ulaanbaatar',
 'Asia/Urumqi',
 'Asia/Ust-Nera',
 'Asia/Vientiane',
 'Asia/Vladivostok',
 'Asia/Yakutsk',
 'Asia/Yangon',
 'Asia/Yekaterinburg',
 'Asia/Yerevan',
 'Atlantic/Azores',
 'Atlantic/Bermuda',
 'Atlantic/Canary',
 'Atlantic/Cape_Verde',
 'Atlantic/Faroe',
 'Atlantic/Madeira',
 'Atlantic/Reykjavik',
 'Atlantic/South_Georgia',
 'Atlantic/St_Helena',
 'Atlantic/Stanley',
 'Australia/Adelaide',
 'Australia/Brisbane',
 'Australia/Broken_Hill',
 'Australia/Darwin',
 'Australia/Eucla',
 'Australia/Hobart',
 'Australia/Lindeman',
 'Australia/Lord_Howe',
 'Australia/Melbourne',
 'Australia/Perth',
 'Australia/Sydney',
 'Canada/Atlantic',
 'Canada/Central',
 'Canada/Eastern',
 'Canada/Mountain',
 'Canada/Newfoundland',
 'Canada/Pacific',
 'Europe/Amsterdam',
 'Europe/Andorra',
 'Europe/Astrakhan',
 'Europe/Athens',
 'Europe/Belgrade',
 'Europe/Berlin',
 'Europe/Bratislava',
 'Europe/Brussels',
 'Europe/Bucharest',
 'Europe/Budapest',
 'Europe/Busingen',
 'Europe/Chisinau',
 'Europe/Copenhagen',
 'Europe/Dublin',
 'Europe/Gibraltar',
 'Europe/Guernsey',
 'Europe/Helsinki',
 'Europe/Isle_of_Man',
 'Europe/Istanbul',
 'Europe/Jersey',
 'Europe/Kaliningrad',
 'Europe/Kiev',
 'Europe/Kirov',
 'Europe/Lisbon',
 'Europe/Ljubljana',
 'Europe/London',
 'Europe/Luxembourg',
 'Europe/Madrid',
 'Europe/Malta',
 'Europe/Mariehamn',
 'Europe/Minsk',
 'Europe/Monaco',
 'Europe/Moscow',
 'Europe/Oslo',
 'Europe/Paris',
 'Europe/Podgorica',
 'Europe/Prague',
 'Europe/Riga',
 'Europe/Rome',
 'Europe/Samara',
 'Europe/San_Marino',
 'Europe/Sarajevo',
 'Europe/Saratov',
 'Europe/Simferopol',
 'Europe/Skopje',
 'Europe/Sofia',
 'Europe/Stockholm',
 'Europe/Tallinn',
 'Europe/Tirane',
 'Europe/Ulyanovsk',
 'Europe/Uzhgorod',
 'Europe/Vaduz',
 'Europe/Vatican',
 'Europe/Vienna',
 'Europe/Vilnius',
 'Europe/Volgograd',
 'Europe/Warsaw',
 'Europe/Zagreb',
 'Europe/Zaporozhye',
 'Europe/Zurich',
 'GMT',
 'Indian/Antananarivo',
 'Indian/Chagos',
 'Indian/Christmas',
 'Indian/Cocos',
 'Indian/Comoro',
 'Indian/Kerguelen',
 'Indian/Mahe',
 'Indian/Maldives',
 'Indian/Mauritius',
 'Indian/Mayotte',
 'Indian/Reunion',
 'Pacific/Apia',
 'Pacific/Auckland',
 'Pacific/Bougainville',
 'Pacific/Chatham',
 'Pacific/Chuuk',
 'Pacific/Easter',
 'Pacific/Efate',
 'Pacific/Fakaofo',
 'Pacific/Fiji',
 'Pacific/Funafuti',
 'Pacific/Galapagos',
 'Pacific/Gambier',
 'Pacific/Guadalcanal',
 'Pacific/Guam',
 'Pacific/Honolulu',
 'Pacific/Kanton',
 'Pacific/Kiritimati',
 'Pacific/Kosrae',
 'Pacific/Kwajalein',
 'Pacific/Majuro',
 'Pacific/Marquesas',
 'Pacific/Midway',
 'Pacific/Nauru',
 'Pacific/Niue',
 'Pacific/Norfolk',
 'Pacific/Noumea',
 'Pacific/Pago_Pago',
 'Pacific/Palau',
 'Pacific/Pitcairn',
 'Pacific/Pohnpei',
 'Pacific/Port_Moresby',
 'Pacific/Rarotonga',
 'Pacific/Saipan',
 'Pacific/Tahiti',
 'Pacific/Tarawa',
 'Pacific/Tongatapu',
 'Pacific/Wake',
 'Pacific/Wallis',
 'US/Alaska',
 'US/Arizona',
 'US/Central',
 'US/Eastern',
 'US/Hawaii',
 'US/Mountain',
 'US/Pacific',
 'UTC']
common_timezones = LazyList(
            tz for tz in common_timezones if tz in all_timezones)
        
common_timezones_set = LazySet(common_timezones)

social_vdgnn/lib/python3.6/site-packages/pytz/__init__.py
eliatrevisan/IA-MPPI-LBM
import numpy
from pytest import raises as assert_raises

import scipy.ndimage as sndi


def test_binary_erosion_noninteger_iterations():
    # regression test for gh-9905, gh-9909: ValueError for
    # non integer iterations
    data = numpy.ones([1])
    assert_raises(TypeError, sndi.binary_erosion, data, iterations=0.5)
    assert_raises(TypeError, sndi.binary_erosion, data, iterations=1.5)


def test_binary_dilation_noninteger_iterations():
    # regression test for gh-9905, gh-9909: ValueError for
    # non integer iterations
    data = numpy.ones([1])
    assert_raises(TypeError, sndi.binary_dilation, data, iterations=0.5)
    assert_raises(TypeError, sndi.binary_dilation, data, iterations=1.5)


def test_binary_opening_noninteger_iterations():
    # regression test for gh-9905, gh-9909: ValueError for
    # non integer iterations
    data = numpy.ones([1])
    assert_raises(TypeError, sndi.binary_opening, data, iterations=0.5)
    assert_raises(TypeError, sndi.binary_opening, data, iterations=1.5)


def test_binary_closing_noninteger_iterations():
    # regression test for gh-9905, gh-9909: ValueError for
    # non integer iterations
    data = numpy.ones([1])
    assert_raises(TypeError, sndi.binary_closing, data, iterations=0.5)
    assert_raises(TypeError, sndi.binary_closing, data, iterations=1.5)


def test_binary_closing_noninteger_brute_force_passes_when_true():
    # regression test for gh-9905, gh-9909: ValueError for
    # non integer iterations
    data = numpy.ones([1])

    assert sndi.binary_erosion(
        data, iterations=2, brute_force=1.5
    ) == sndi.binary_erosion(data, iterations=2, brute_force=bool(1.5))
    assert sndi.binary_erosion(
        data, iterations=2, brute_force=0.0
    ) == sndi.binary_erosion(data, iterations=2, brute_force=bool(0.0))

social_vdgnn/lib/python3.6/site-packages/scipy/ndimage/tests/test_morphology.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
from numpy.testing import assert_equal, assert_almost_equal
import scipy.special as sc


def test_ndtr():
    assert_equal(sc.ndtr(0), 0.5)
    assert_almost_equal(sc.ndtr(1), 0.84134474606)


class TestNdtri:

    def test_zero(self):
        assert sc.ndtri(0.5) == 0.0

    def test_asymptotes(self):
        assert_equal(sc.ndtri([0.0, 1.0]), [-np.inf, np.inf])

    def test_outside_of_domain(self):
        assert all(np.isnan(sc.ndtri([-1.5, 1.5])))

social_vdgnn/lib/python3.6/site-packages/scipy/special/tests/test_ndtr.py
eliatrevisan/IA-MPPI-LBM
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Wrapper around plugin to conditionally enable it."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorboard.plugins import base_plugin


class BeholderPluginLoader(base_plugin.TBLoader):
  """BeholderPlugin factory.

  This class checks for `tensorflow` install and dependency.
  """

  def load(self, context):
    """Returns the plugin, if possible.

    Args:
      context: The TBContext flags.

    Returns:
      A BeholderPlugin instance or None if it couldn't be loaded.
    """
    try:
      # pylint: disable=g-import-not-at-top,unused-import
      import tensorflow
    except ImportError:
      return
    # pylint: disable=g-import-not-at-top
    from tensorboard.plugins.beholder.beholder_plugin import BeholderPlugin
    return BeholderPlugin(context)

social_vdgnn/lib/python3.6/site-packages/tensorboard/plugins/beholder/beholder_plugin_loader.py
eliatrevisan/IA-MPPI-LBM
import abc
import collections
import collections.abc
import operator
import sys
import types as _types
import typing

# After PEP 560, internal typing API was substantially reworked.
# This is especially important for Protocol class which uses internal APIs
# quite extensively.
PEP_560 = sys.version_info[:3] >= (3, 7, 0)

if PEP_560:
    GenericMeta = type
else:
    # 3.6
    from typing import GenericMeta, _type_vars  # noqa


# Please keep __all__ alphabetized within each category.
__all__ = [
    # Super-special typing primitives.
    'ClassVar',
    'Concatenate',
    'Final',
    'LiteralString',
    'ParamSpec',
    'Self',
    'Type',
    'TypeVarTuple',
    'Unpack',

    # ABCs (from collections.abc).
    'Awaitable',
    'AsyncIterator',
    'AsyncIterable',
    'Coroutine',
    'AsyncGenerator',
    'AsyncContextManager',
    'ChainMap',

    # Concrete collection types.
    'ContextManager',
    'Counter',
    'Deque',
    'DefaultDict',
    'OrderedDict',
    'TypedDict',

    # Structural checks, a.k.a. protocols.
    'SupportsIndex',

    # One-off things.
    'Annotated',
    'assert_never',
    'dataclass_transform',
    'final',
    'IntVar',
    'is_typeddict',
    'Literal',
    'NewType',
    'overload',
    'Protocol',
    'reveal_type',
    'runtime',
    'runtime_checkable',
    'Text',
    'TypeAlias',
    'TypeGuard',
    'TYPE_CHECKING',
    'Never',
    'NoReturn',
    'Required',
    'NotRequired',
]

if PEP_560:
    __all__.extend(["get_args", "get_origin", "get_type_hints"])

# The functions below are modified copies of typing internal helpers.
# They are needed by _ProtocolMeta and they provide support for PEP 646.


def _no_slots_copy(dct):
    dict_copy = dict(dct)
    if '__slots__' in dict_copy:
        for slot in dict_copy['__slots__']:
            dict_copy.pop(slot, None)
    return dict_copy


_marker = object()


def _check_generic(cls, parameters, elen=_marker):
    """Check correct count for parameters of a generic cls (internal helper).
    This gives a nice error message in case of count mismatch.
    """
    if not elen:
        raise TypeError(f"{cls} is not a generic class")
    if elen is _marker:
        if not hasattr(cls, "__parameters__") or not cls.__parameters__:
            raise TypeError(f"{cls} is not a generic class")
        elen = len(cls.__parameters__)
    alen = len(parameters)
    if alen != elen:
        if hasattr(cls, "__parameters__"):
            parameters = [p for p in cls.__parameters__ if not _is_unpack(p)]
            num_tv_tuples = sum(isinstance(p, TypeVarTuple) for p in parameters)
            if (num_tv_tuples > 0) and (alen >= elen - num_tv_tuples):
                return
        raise TypeError(f"Too {'many' if alen > elen else 'few'} parameters for {cls};"
                        f" actual {alen}, expected {elen}")


if sys.version_info >= (3, 10):
    def _should_collect_from_parameters(t):
        return isinstance(
            t, (typing._GenericAlias, _types.GenericAlias, _types.UnionType)
        )
elif sys.version_info >= (3, 9):
    def _should_collect_from_parameters(t):
        return isinstance(t, (typing._GenericAlias, _types.GenericAlias))
else:
    def _should_collect_from_parameters(t):
        return isinstance(t, typing._GenericAlias) and not t._special


def _collect_type_vars(types, typevar_types=None):
    """Collect all type variable contained in types in order of
    first appearance (lexicographic order). For example::

        _collect_type_vars((T, List[S, T])) == (T, S)
    """
    if typevar_types is None:
        typevar_types = typing.TypeVar
    tvars = []
    for t in types:
        if (
            isinstance(t, typevar_types) and
            t not in tvars and
            not _is_unpack(t)
        ):
            tvars.append(t)
        if _should_collect_from_parameters(t):
            tvars.extend([t for t in t.__parameters__ if t not in tvars])
    return tuple(tvars)


# 3.6.2+
if hasattr(typing, 'NoReturn'):
    NoReturn = typing.NoReturn
# 3.6.0-3.6.1
else:
    class _NoReturn(typing._FinalTypingBase, _root=True):
        """Special type indicating functions that never return.
        Example::

          from typing import NoReturn

          def stop() -> NoReturn:
              raise Exception('no way')

        This type is invalid in other positions, e.g., ``List[NoReturn]``
        will fail in static type checkers.
        """
        __slots__ = ()

        def __instancecheck__(self, obj):
            raise TypeError("NoReturn cannot be used with isinstance().")

        def __subclasscheck__(self, cls):
            raise TypeError("NoReturn cannot be used with issubclass().")

    NoReturn = _NoReturn(_root=True)

# Some unconstrained type variables.  These are used by the container types.
# (These are not for export.)
T = typing.TypeVar('T')  # Any type.
KT = typing.TypeVar('KT')  # Key type.
VT = typing.TypeVar('VT')  # Value type.
T_co = typing.TypeVar('T_co', covariant=True)  # Any type covariant containers.
T_contra = typing.TypeVar('T_contra', contravariant=True)  # Ditto contravariant.

ClassVar = typing.ClassVar

# On older versions of typing there is an internal class named "Final".
# 3.8+
if hasattr(typing, 'Final') and sys.version_info[:2] >= (3, 7):
    Final = typing.Final
# 3.7
elif sys.version_info[:2] >= (3, 7):
    class _FinalForm(typing._SpecialForm, _root=True):

        def __repr__(self):
            return 'typing_extensions.' + self._name

        def __getitem__(self, parameters):
            item = typing._type_check(parameters,
                                      f'{self._name} accepts only single type')
            return typing._GenericAlias(self, (item,))

    Final = _FinalForm('Final',
                       doc="""A special typing construct to indicate that a name
                       cannot be re-assigned or overridden in a subclass.
                       For example:

                           MAX_SIZE: Final = 9000
                           MAX_SIZE += 1  # Error reported by type checker

                           class Connection:
                               TIMEOUT: Final[int] = 10
                           class FastConnector(Connection):
                               TIMEOUT = 1  # Error reported by type checker

                       There is no runtime checking of these properties.""")
# 3.6
else:
    class _Final(typing._FinalTypingBase, _root=True):
        """A special typing construct to indicate that a name
        cannot be re-assigned or overridden in a subclass.
        For example:

            MAX_SIZE: Final = 9000
            MAX_SIZE += 1  # Error reported by type checker

            class Connection:
                TIMEOUT: Final[int] = 10
            class FastConnector(Connection):
                TIMEOUT = 1  # Error reported by type checker

        There is no runtime checking of these properties.
        """

        __slots__ = ('__type__',)

        def __init__(self, tp=None, **kwds):
            self.__type__ = tp

        def __getitem__(self, item):
            cls = type(self)
            if self.__type__ is None:
                return cls(typing._type_check(item,
                           f'{cls.__name__[1:]} accepts only single type.'),
                           _root=True)
            raise TypeError(f'{cls.__name__[1:]} cannot be further subscripted')

        def _eval_type(self, globalns, localns):
            new_tp = typing._eval_type(self.__type__, globalns, localns)
            if new_tp == self.__type__:
                return self
            return type(self)(new_tp, _root=True)

        def __repr__(self):
            r = super().__repr__()
            if self.__type__ is not None:
                r += f'[{typing._type_repr(self.__type__)}]'
            return r

        def __hash__(self):
            return hash((type(self).__name__, self.__type__))

        def __eq__(self, other):
            if not isinstance(other, _Final):
                return NotImplemented
            if self.__type__ is not None:
                return self.__type__ == other.__type__
            return self is other

    Final = _Final(_root=True)


if sys.version_info >= (3, 11):
    final = typing.final
else:
    # @final exists in 3.8+, but we backport it for all versions
    # before 3.11 to keep support for the __final__ attribute.
    # See https://bugs.python.org/issue46342
    def final(f):
        """This decorator can be used to indicate to type checkers that
        the decorated method cannot be overridden, and decorated class
        cannot be subclassed. For example:

            class Base:
                @final
                def done(self) -> None:
                    ...
            class Sub(Base):
                def done(self) -> None:  # Error reported by type checker
                    ...
            @final
            class Leaf:
                ...
            class Other(Leaf):  # Error reported by type checker
                ...

        There is no runtime checking of these properties. The decorator
        sets the ``__final__`` attribute to ``True`` on the decorated object
        to allow runtime introspection.
        """
        try:
            f.__final__ = True
        except (AttributeError, TypeError):
            # Skip the attribute silently if it is not writable.
            # AttributeError happens if the object has __slots__ or a
            # read-only property, TypeError if it's a builtin class.
            pass
        return f


def IntVar(name):
    return typing.TypeVar(name)


# 3.8+:
if hasattr(typing, 'Literal'):
    Literal = typing.Literal
# 3.7:
elif sys.version_info[:2] >= (3, 7):
    class _LiteralForm(typing._SpecialForm, _root=True):

        def __repr__(self):
            return 'typing_extensions.' + self._name

        def __getitem__(self, parameters):
            return typing._GenericAlias(self, parameters)

    Literal = _LiteralForm('Literal',
                           doc="""A type that can be used to indicate to type checkers
                           that the corresponding value has a value literally equivalent
                           to the provided parameter. For example:

                               var: Literal[4] = 4

                           The type checker understands that 'var' is literally equal to
                           the value 4 and no other value.

                           Literal[...] cannot be subclassed. There is no runtime
                           checking verifying that the parameter is actually a value
                           instead of a type.""")
# 3.6:
else:
    class _Literal(typing._FinalTypingBase, _root=True):
        """A type that can be used to indicate to type checkers that the
        corresponding value has a value literally equivalent to the
        provided parameter. For example:

            var: Literal[4] = 4

        The type checker understands that 'var' is literally equal to the
        value 4 and no other value.

        Literal[...] cannot be subclassed. There is no runtime checking
        verifying that the parameter is actually a value instead of a type.
        """

        __slots__ = ('__values__',)

        def __init__(self, values=None, **kwds):
            self.__values__ = values

        def __getitem__(self, values):
            cls = type(self)
            if self.__values__ is None:
                if not isinstance(values, tuple):
                    values = (values,)
                return cls(values, _root=True)
            raise TypeError(f'{cls.__name__[1:]} cannot be further subscripted')

        def _eval_type(self, globalns, localns):
            return self

        def __repr__(self):
            r = super().__repr__()
            if self.__values__ is not None:
                r += f'[{", ".join(map(typing._type_repr, self.__values__))}]'
            return r

        def __hash__(self):
            return hash((type(self).__name__, self.__values__))

        def __eq__(self, other):
            if not isinstance(other, _Literal):
                return NotImplemented
            if self.__values__ is not None:
                return self.__values__ == other.__values__
            return self is other

    Literal = _Literal(_root=True)


_overload_dummy = typing._overload_dummy  # noqa
overload = typing.overload


# This is not a real generic class.  Don't use outside annotations.
Type = typing.Type

# Various ABCs mimicking those in collections.abc.
# A few are simply re-exported for completeness.


class _ExtensionsGenericMeta(GenericMeta):
    def __subclasscheck__(self, subclass):
        """This mimics a more modern GenericMeta.__subclasscheck__() logic
        (that does not have problems with recursion) to work around interactions
        between collections, typing, and typing_extensions on older
        versions of Python, see https://github.com/python/typing/issues/501.
        """
        if self.__origin__ is not None:
            if sys._getframe(1).f_globals['__name__'] not in ['abc', 'functools']:
                raise TypeError("Parameterized generics cannot be used with class "
                                "or instance checks")
            return False
        if not self.__extra__:
            return super().__subclasscheck__(subclass)
        res = self.__extra__.__subclasshook__(subclass)
        if res is not NotImplemented:
            return res
        if self.__extra__ in subclass.__mro__:
            return True
        for scls in self.__extra__.__subclasses__():
            if isinstance(scls, GenericMeta):
                continue
            if issubclass(subclass, scls):
                return True
        return False


Awaitable = typing.Awaitable
Coroutine = typing.Coroutine
AsyncIterable = typing.AsyncIterable
AsyncIterator = typing.AsyncIterator

# 3.6.1+
if hasattr(typing, 'Deque'):
    Deque = typing.Deque
# 3.6.0
else:
    class Deque(collections.deque, typing.MutableSequence[T],
                metaclass=_ExtensionsGenericMeta,
                extra=collections.deque):
        __slots__ = ()

        def __new__(cls, *args, **kwds):
            if cls._gorg is Deque:
                return collections.deque(*args, **kwds)
            return typing._generic_new(collections.deque, cls, *args, **kwds)

ContextManager = typing.ContextManager
# 3.6.2+
if hasattr(typing, 'AsyncContextManager'):
    AsyncContextManager = typing.AsyncContextManager
# 3.6.0-3.6.1
else:
    from _collections_abc import _check_methods as _check_methods_in_mro  # noqa

    class AsyncContextManager(typing.Generic[T_co]):
        __slots__ = ()

        async def __aenter__(self):
            return self

        @abc.abstractmethod
        async def __aexit__(self, exc_type, exc_value, traceback):
            return None

        @classmethod
        def __subclasshook__(cls, C):
            if cls is AsyncContextManager:
                return _check_methods_in_mro(C, "__aenter__", "__aexit__")
            return NotImplemented

DefaultDict = typing.DefaultDict

# 3.7.2+
if hasattr(typing, 'OrderedDict'):
    OrderedDict = typing.OrderedDict
# 3.7.0-3.7.2
elif (3, 7, 0) <= sys.version_info[:3] < (3, 7, 2):
    OrderedDict = typing._alias(collections.OrderedDict, (KT, VT))
# 3.6
else:
    class OrderedDict(collections.OrderedDict, typing.MutableMapping[KT, VT],
                      metaclass=_ExtensionsGenericMeta,
                      extra=collections.OrderedDict):

        __slots__ = ()

        def __new__(cls, *args, **kwds):
            if cls._gorg is OrderedDict:
                return collections.OrderedDict(*args, **kwds)
            return typing._generic_new(collections.OrderedDict, cls, *args, **kwds)

# 3.6.2+
if hasattr(typing, 'Counter'):
    Counter = typing.Counter
# 3.6.0-3.6.1
else:
    class Counter(collections.Counter,
                  typing.Dict[T, int],
                  metaclass=_ExtensionsGenericMeta, extra=collections.Counter):

        __slots__ = ()

        def __new__(cls, *args, **kwds):
            if cls._gorg is Counter:
                return collections.Counter(*args, **kwds)
            return typing._generic_new(collections.Counter, cls, *args, **kwds)

# 3.6.1+
if hasattr(typing, 'ChainMap'):
    ChainMap = typing.ChainMap
elif hasattr(collections, 'ChainMap'):
    class ChainMap(collections.ChainMap, typing.MutableMapping[KT, VT],
                   metaclass=_ExtensionsGenericMeta,
                   extra=collections.ChainMap):

        __slots__ = ()

        def __new__(cls, *args, **kwds):
            if cls._gorg is ChainMap:
                return collections.ChainMap(*args, **kwds)
            return typing._generic_new(collections.ChainMap, cls, *args, **kwds)

# 3.6.1+
if hasattr(typing, 'AsyncGenerator'):
    AsyncGenerator = typing.AsyncGenerator
# 3.6.0
else:
    class AsyncGenerator(AsyncIterator[T_co], typing.Generic[T_co, T_contra],
                         metaclass=_ExtensionsGenericMeta,
                         extra=collections.abc.AsyncGenerator):
        __slots__ = ()

NewType = typing.NewType
Text = typing.Text
TYPE_CHECKING = typing.TYPE_CHECKING


def _gorg(cls):
    """This function exists for compatibility with old typing versions."""
    assert isinstance(cls, GenericMeta)
    if hasattr(cls, '_gorg'):
        return cls._gorg
    while cls.__origin__ is not None:
        cls = cls.__origin__
    return cls


_PROTO_WHITELIST = ['Callable', 'Awaitable',
                    'Iterable', 'Iterator', 'AsyncIterable', 'AsyncIterator',
                    'Hashable', 'Sized', 'Container', 'Collection', 'Reversible',
                    'ContextManager', 'AsyncContextManager']


def _get_protocol_attrs(cls):
    attrs = set()
    for base in cls.__mro__[:-1]:  # without object
        if base.__name__ in ('Protocol', 'Generic'):
            continue
        annotations = getattr(base, '__annotations__', {})
        for attr in list(base.__dict__.keys()) + list(annotations.keys()):
            if (not attr.startswith('_abc_') and attr not in (
                    '__abstractmethods__', '__annotations__', '__weakref__',
                    '_is_protocol', '_is_runtime_protocol', '__dict__',
                    '__args__', '__slots__',
                    '__next_in_mro__', '__parameters__', '__origin__',
                    '__orig_bases__', '__extra__', '__tree_hash__',
                    '__doc__', '__subclasshook__', '__init__', '__new__',
                    '__module__', '_MutableMapping__marker', '_gorg')):
                attrs.add(attr)
    return attrs


def _is_callable_members_only(cls):
    return all(callable(getattr(cls, attr, None)) for attr in _get_protocol_attrs(cls))


# 3.8+
if hasattr(typing, 'Protocol'):
    Protocol = typing.Protocol
# 3.7
elif PEP_560:

    def _no_init(self, *args, **kwargs):
        if type(self)._is_protocol:
            raise TypeError('Protocols cannot be instantiated')

    class _ProtocolMeta(abc.ABCMeta):
        # This metaclass is a bit unfortunate and exists only because of the lack
        # of __instancehook__.
        def __instancecheck__(cls, instance):
            # We need this method for situations where attributes are
            # assigned in __init__.
            if ((not getattr(cls, '_is_protocol', False) or
                 _is_callable_members_only(cls)) and
                    issubclass(instance.__class__, cls)):
                return True
            if cls._is_protocol:
                if all(hasattr(instance, attr) and
                       (not callable(getattr(cls, attr, None)) or
                        getattr(instance, attr) is not None)
                       for attr in _get_protocol_attrs(cls)):
                    return True
            return super().__instancecheck__(instance)

    class Protocol(metaclass=_ProtocolMeta):
        # There is quite a lot of overlapping code with typing.Generic.
        # Unfortunately it is hard to avoid this while these live in two different
        # modules. The duplicated code will be removed when Protocol is moved to typing.
        """Base class for protocol classes. Protocol classes are defined as::

            class Proto(Protocol):
                def meth(self) -> int:
                    ...

        Such classes are primarily used with static type checkers that recognize
        structural subtyping (static duck-typing), for example::

            class C:
                def meth(self) -> int:
                    return 0

            def func(x: Proto) -> int:
                return x.meth()

            func(C())  # Passes static type check

        See PEP 544 for details. Protocol classes decorated with
        @typing_extensions.runtime act as simple-minded runtime protocol that checks
        only the presence of given attributes, ignoring their type signatures.

        Protocol classes can be generic, they are defined as::

            class GenProto(Protocol[T]):
                def meth(self) -> T:
                    ...
        """
        __slots__ = ()
        _is_protocol = True

        def __new__(cls, *args, **kwds):
            if cls is Protocol:
                raise TypeError("Type Protocol cannot be instantiated; "
                                "it can only be used as a base class")
            return super().__new__(cls)

        @typing._tp_cache
        def __class_getitem__(cls, params):
            if not isinstance(params, tuple):
                params = (params,)
            if not params and cls is not typing.Tuple:
                raise TypeError(
                    f"Parameter list to {cls.__qualname__}[...] cannot be empty")
            msg = "Parameters to generic types must be types."
            params = tuple(typing._type_check(p, msg) for p in params)  # noqa
            if cls is Protocol:
                # Generic can only be subscripted with unique type variables.
                if not all(isinstance(p, typing.TypeVar) for p in params):
                    i = 0
                    while isinstance(params[i], typing.TypeVar):
                        i += 1
                    raise TypeError(
                        "Parameters to Protocol[...] must all be type variables."
                        f" Parameter {i + 1} is {params[i]}")
                if len(set(params)) != len(params):
                    raise TypeError(
                        "Parameters to Protocol[...] must all be unique")
            else:
                # Subscripting a regular Generic subclass.
                _check_generic(cls, params, len(cls.__parameters__))
            return typing._GenericAlias(cls, params)

        def __init_subclass__(cls, *args, **kwargs):
            tvars = []
            if '__orig_bases__' in cls.__dict__:
                error = typing.Generic in cls.__orig_bases__
            else:
                error = typing.Generic in cls.__bases__
            if error:
                raise TypeError("Cannot inherit from plain Generic")
            if '__orig_bases__' in cls.__dict__:
                tvars = typing._collect_type_vars(cls.__orig_bases__)
                # Look for Generic[T1, ..., Tn] or Protocol[T1, ..., Tn].
                # If found, tvars must be a subset of it.
                # If not found, tvars is it.
                # Also check for and reject plain Generic,
                # and reject multiple Generic[...] and/or Protocol[...].
                gvars = None
                for base in cls.__orig_bases__:
                    if (isinstance(base, typing._GenericAlias) and
                            base.__origin__ in (typing.Generic, Protocol)):
                        # for error messages
                        the_base = base.__origin__.__name__
                        if gvars is not None:
                            raise TypeError(
                                "Cannot inherit from Generic[...]"
                                " and/or Protocol[...] multiple types.")
                        gvars = base.__parameters__
                if gvars is None:
                    gvars = tvars
                else:
                    tvarset = set(tvars)
                    gvarset = set(gvars)
                    if not tvarset <= gvarset:
                        s_vars = ', '.join(str(t) for t in tvars if t not in gvarset)
                        s_args = ', '.join(str(g) for g in gvars)
                        raise TypeError(f"Some type variables ({s_vars}) are"
                                        f" not listed in {the_base}[{s_args}]")
                    tvars = gvars
            cls.__parameters__ = tuple(tvars)

            # Determine if this is a protocol or a concrete subclass.
            if not cls.__dict__.get('_is_protocol', None):
                cls._is_protocol = any(b is Protocol for b in cls.__bases__)

            # Set (or override) the protocol subclass hook.
            def _proto_hook(other):
                if not cls.__dict__.get('_is_protocol', None):
                    return NotImplemented
                if not getattr(cls, '_is_runtime_protocol', False):
                    if sys._getframe(2).f_globals['__name__'] in ['abc', 'functools']:
                        return NotImplemented
                    raise TypeError("Instance and class checks can only be used with"
                                    " @runtime protocols")
                if not _is_callable_members_only(cls):
                    if sys._getframe(2).f_globals['__name__'] in ['abc', 'functools']:
                        return NotImplemented
                    raise TypeError("Protocols with non-method members"
                                    " don't support issubclass()")
                if not isinstance(other, type):
                    # Same error as for issubclass(1, int)
                    raise TypeError('issubclass() arg 1 must be a class')
                for attr in _get_protocol_attrs(cls):
                    for base in other.__mro__:
                        if attr in base.__dict__:
                            if base.__dict__[attr] is None:
                                return NotImplemented
                            break
                        annotations = getattr(base, '__annotations__', {})
                        if (isinstance(annotations, typing.Mapping) and
                                attr in annotations and
                                isinstance(other, _ProtocolMeta) and
                                other._is_protocol):
                            break
                    else:
                        return NotImplemented
                return True
            if '__subclasshook__' not in cls.__dict__:
                cls.__subclasshook__ = _proto_hook

            # We have nothing more to do for non-protocols.
            if not cls._is_protocol:
                return

            # Check consistency of bases.
            for base in cls.__bases__:
                if not (base in (object, typing.Generic) or
                        base.__module__ == 'collections.abc' and
                        base.__name__ in _PROTO_WHITELIST or
                        isinstance(base, _ProtocolMeta) and base._is_protocol):
                    raise TypeError('Protocols can only inherit from other'
                                    f' protocols, got {repr(base)}')
            cls.__init__ = _no_init
# 3.6
else:
    from typing import _next_in_mro, _type_check  # noqa

    def _no_init(self, *args, **kwargs):
        if type(self)._is_protocol:
            raise TypeError('Protocols cannot be instantiated')

    class _ProtocolMeta(GenericMeta):
        """Internal metaclass for Protocol.

        This exists so Protocol classes can be generic without deriving
        from Generic.
        """
        def __new__(cls, name, bases, namespace,
                    tvars=None, args=None, origin=None, extra=None, orig_bases=None):
            # This is just a version copied from GenericMeta.__new__ that
            # includes "Protocol" special treatment. (Comments removed for brevity.)
            assert extra is None  # Protocols should not have extra
            if tvars is not None:
                assert origin is not None
                assert all(isinstance(t, typing.TypeVar) for t in tvars), tvars
            else:
                tvars = _type_vars(bases)
                gvars = None
                for base in bases:
                    if base is typing.Generic:
                        raise TypeError("Cannot inherit from plain Generic")
                    if (isinstance(base, GenericMeta) and
                            base.__origin__ in (typing.Generic, Protocol)):
                        if gvars is not None:
                            raise TypeError(
                                "Cannot inherit from Generic[...] or"
                                " Protocol[...] multiple times.")
                        gvars = base.__parameters__
                if gvars is None:
                    gvars = tvars
                else:
                    tvarset = set(tvars)
                    gvarset = set(gvars)
                    if not tvarset <= gvarset:
                        s_vars = ", ".join(str(t) for t in tvars if t not in gvarset)
                        s_args = ", ".join(str(g) for g in gvars)
                        cls_name = "Generic" if any(b.__origin__ is typing.Generic
                                                    for b in bases) else "Protocol"
                        raise TypeError(f"Some type variables ({s_vars}) are"
                                        f" not listed in {cls_name}[{s_args}]")
                    tvars = gvars

            initial_bases = bases
            if (extra is not None and type(extra) is abc.ABCMeta and
                    extra not in bases):
                bases = (extra,) + bases
            bases = tuple(_gorg(b) if isinstance(b, GenericMeta) else b
                          for b in bases)
            if any(isinstance(b, GenericMeta) and b is not typing.Generic for b in bases):
                bases = tuple(b for b in bases if b is not typing.Generic)
            namespace.update({'__origin__': origin, '__extra__': extra})
            self = super(GenericMeta, cls).__new__(cls, name, bases, namespace,
                                                   _root=True)
            super(GenericMeta, self).__setattr__('_gorg',
                                                 self if not origin else
                                                 _gorg(origin))
            self.__parameters__ = tvars
            self.__args__ = tuple(... if a is typing._TypingEllipsis else
                                  () if a is typing._TypingEmpty else
                                  a for a in args) if args else None
            self.__next_in_mro__ = _next_in_mro(self)
            if orig_bases is None:
                self.__orig_bases__ = initial_bases
            elif origin is not None:
                self._abc_registry = origin._abc_registry
                self._abc_cache = origin._abc_cache
            if hasattr(self, '_subs_tree'):
                self.__tree_hash__ = (hash(self._subs_tree()) if origin else
                                      super(GenericMeta, self).__hash__())
            return self

        def __init__(cls, *args, **kwargs):
            super().__init__(*args, **kwargs)
            if not cls.__dict__.get('_is_protocol', None):
                cls._is_protocol = any(b is Protocol or
                                       isinstance(b, _ProtocolMeta) and
                                       b.__origin__ is Protocol
                                       for b in cls.__bases__)
            if cls._is_protocol:
                for base in cls.__mro__[1:]:
                    if not (base in (object, typing.Generic) or
                            base.__module__ == 'collections.abc' and
                            base.__name__ in _PROTO_WHITELIST or
                            isinstance(base, typing.TypingMeta) and base._is_protocol or
                            isinstance(base, GenericMeta) and
                            base.__origin__ is typing.Generic):
                        raise TypeError(f'Protocols can only inherit from other'
                                        f' protocols, got {repr(base)}')

                cls.__init__ = _no_init

            def _proto_hook(other):
                if not cls.__dict__.get('_is_protocol', None):
                    return NotImplemented
                if not isinstance(other, type):
                    # Same error as for issubclass(1, int)
                    raise TypeError('issubclass() arg 1 must be a class')
                for attr in _get_protocol_attrs(cls):
                    for base in other.__mro__:
                        if attr in base.__dict__:
                            if base.__dict__[attr] is None:
                                return NotImplemented
                            break
                        annotations = getattr(base, '__annotations__', {})
                        if (isinstance(annotations, typing.Mapping) and
                                attr in annotations and
                                isinstance(other, _ProtocolMeta) and
                                other._is_protocol):
                            break
                    else:
                        return NotImplemented
                return True
            if '__subclasshook__' not in cls.__dict__:
                cls.__subclasshook__ = _proto_hook

        def __instancecheck__(self, instance):
            # We need this method for situations where attributes are
            # assigned in __init__.
            if ((not getattr(self, '_is_protocol', False) or
                    _is_callable_members_only(self)) and
                    issubclass(instance.__class__, self)):
                return True
            if self._is_protocol:
                if all(hasattr(instance, attr) and
                        (not callable(getattr(self, attr, None)) or
                         getattr(instance, attr) is not None)
                        for attr in _get_protocol_attrs(self)):
                    return True
            return super(GenericMeta, self).__instancecheck__(instance)

        def __subclasscheck__(self, cls):
            if self.__origin__ is not None:
                if sys._getframe(1).f_globals['__name__'] not in ['abc', 'functools']:
                    raise TypeError("Parameterized generics cannot be used with class "
                                    "or instance checks")
                return False
            if (self.__dict__.get('_is_protocol', None) and
                    not self.__dict__.get('_is_runtime_protocol', None)):
                if sys._getframe(1).f_globals['__name__'] in ['abc',
                                                              'functools',
                                                              'typing']:
                    return False
                raise TypeError("Instance and class checks can only be used with"
                                " @runtime protocols")
            if (self.__dict__.get('_is_runtime_protocol', None) and
                    not _is_callable_members_only(self)):
                if sys._getframe(1).f_globals['__name__'] in ['abc',
                                                              'functools',
                                                              'typing']:
                    return super(GenericMeta, self).__subclasscheck__(cls)
                raise TypeError("Protocols with non-method members"
                                " don't support issubclass()")
            return super(GenericMeta, self).__subclasscheck__(cls)

        @typing._tp_cache
        def __getitem__(self, params):
            # We also need to copy this from GenericMeta.__getitem__ to get
            # special treatment of "Protocol". (Comments removed for brevity.)
            if not isinstance(params, tuple):
                params = (params,)
            if not params and _gorg(self) is not typing.Tuple:
                raise TypeError(
                    f"Parameter list to {self.__qualname__}[...] cannot be empty")
            msg = "Parameters to generic types must be types."
            params = tuple(_type_check(p, msg) for p in params)
            if self in (typing.Generic, Protocol):
                if not all(isinstance(p, typing.TypeVar) for p in params):
                    raise TypeError(
                        f"Parameters to {repr(self)}[...] must all be type variables")
                if len(set(params)) != len(params):
                    raise TypeError(
                        f"Parameters to {repr(self)}[...] must all be unique")
                tvars = params
                args = params
            elif self in (typing.Tuple, typing.Callable):
                tvars = _type_vars(params)
                args = params
            elif self.__origin__ in (typing.Generic, Protocol):
                raise TypeError(f"Cannot subscript already-subscripted {repr(self)}")
            else:
                _check_generic(self, params, len(self.__parameters__))
                tvars = _type_vars(params)
                args = params

            prepend = (self,) if self.__origin__ is None else ()
            return self.__class__(self.__name__,
                                  prepend + self.__bases__,
                                  _no_slots_copy(self.__dict__),
                                  tvars=tvars,
                                  args=args,
                                  origin=self,
                                  extra=self.__extra__,
                                  orig_bases=self.__orig_bases__)

    class Protocol(metaclass=_ProtocolMeta):
        """Base class for protocol classes. Protocol classes are defined as::

          class Proto(Protocol):
              def meth(self) -> int:
                  ...

        Such classes are primarily used with static type checkers that recognize
        structural subtyping (static duck-typing), for example::

          class C:
              def meth(self) -> int:
                  return 0

          def func(x: Proto) -> int:
              return x.meth()

          func(C())  # Passes static type check

        See PEP 544 for details. Protocol classes decorated with
        @typing_extensions.runtime act as simple-minded runtime protocol that checks
        only the presence of given attributes, ignoring their type signatures.

        Protocol classes can be generic, they are defined as::

          class GenProto(Protocol[T]):
              def meth(self) -> T:
                  ...
        """
        __slots__ = ()
        _is_protocol = True

        def __new__(cls, *args, **kwds):
            if _gorg(cls) is Protocol:
                raise TypeError("Type Protocol cannot be instantiated; "
                                "it can be used only as a base class")
            return typing._generic_new(cls.__next_in_mro__, cls, *args, **kwds)


# 3.8+
if hasattr(typing, 'runtime_checkable'):
    runtime_checkable = typing.runtime_checkable
# 3.6-3.7
else:
    def runtime_checkable(cls):
        """Mark a protocol class as a runtime protocol, so that it
        can be used with isinstance() and issubclass(). Raise TypeError
        if applied to a non-protocol class.

        This allows a simple-minded structural check very similar to the
        one-offs in collections.abc such as Hashable.
        """
        if not isinstance(cls, _ProtocolMeta) or not cls._is_protocol:
            raise TypeError('@runtime_checkable can be only applied to protocol classes,'
                            f' got {cls!r}')
        cls._is_runtime_protocol = True
        return cls


# Exists for backwards compatibility.
runtime = runtime_checkable


# 3.8+
if hasattr(typing, 'SupportsIndex'):
    SupportsIndex = typing.SupportsIndex
# 3.6-3.7
else:
    @runtime_checkable
    class SupportsIndex(Protocol):
        __slots__ = ()

        @abc.abstractmethod
        def __index__(self) -> int:
            pass


if hasattr(typing, "Required"):
    # The standard library TypedDict in Python 3.8 does not store runtime information
    # about which (if any) keys are optional.  See https://bugs.python.org/issue38834
    # The standard library TypedDict in Python 3.9.0/1 does not honour the "total"
    # keyword with old-style TypedDict().  See https://bugs.python.org/issue42059
    # The standard library TypedDict below Python 3.11 does not store runtime
    # information about optional and required keys when using Required or NotRequired.
    TypedDict = typing.TypedDict
    _TypedDictMeta = typing._TypedDictMeta
    is_typeddict = typing.is_typeddict
else:
    def _check_fails(cls, other):
        try:
            if sys._getframe(1).f_globals['__name__'] not in ['abc',
                                                              'functools',
                                                              'typing']:
                # Typed dicts are only for static structural subtyping.
                raise TypeError('TypedDict does not support instance and class checks')
        except (AttributeError, ValueError):
            pass
        return False

    def _dict_new(*args, **kwargs):
        if not args:
            raise TypeError('TypedDict.__new__(): not enough arguments')
        _, args = args[0], args[1:]  # allow the "cls" keyword be passed
        return dict(*args, **kwargs)

    _dict_new.__text_signature__ = '($cls, _typename, _fields=None, /, **kwargs)'

    def _typeddict_new(*args, total=True, **kwargs):
        if not args:
            raise TypeError('TypedDict.__new__(): not enough arguments')
        _, args = args[0], args[1:]  # allow the "cls" keyword be passed
        if args:
            typename, args = args[0], args[1:]  # allow the "_typename" keyword be passed
        elif '_typename' in kwargs:
            typename = kwargs.pop('_typename')
            import warnings
            warnings.warn("Passing '_typename' as keyword argument is deprecated",
                          DeprecationWarning, stacklevel=2)
        else:
            raise TypeError("TypedDict.__new__() missing 1 required positional "
                            "argument: '_typename'")
        if args:
            try:
                fields, = args  # allow the "_fields" keyword be passed
            except ValueError:
                raise TypeError('TypedDict.__new__() takes from 2 to 3 '
                                f'positional arguments but {len(args) + 2} '
                                'were given')
        elif '_fields' in kwargs and len(kwargs) == 1:
            fields = kwargs.pop('_fields')
            import warnings
            warnings.warn("Passing '_fields' as keyword argument is deprecated",
                          DeprecationWarning, stacklevel=2)
        else:
            fields = None

        if fields is None:
            fields = kwargs
        elif kwargs:
            raise TypeError("TypedDict takes either a dict or keyword arguments,"
                            " but not both")

        ns = {'__annotations__': dict(fields)}
        try:
            # Setting correct module is necessary to make typed dict classes pickleable.
            ns['__module__'] = sys._getframe(1).f_globals.get('__name__', '__main__')
        except (AttributeError, ValueError):
            pass

        return _TypedDictMeta(typename, (), ns, total=total)

    _typeddict_new.__text_signature__ = ('($cls, _typename, _fields=None,'
                                         ' /, *, total=True, **kwargs)')

    class _TypedDictMeta(type):
        def __init__(cls, name, bases, ns, total=True):
            super().__init__(name, bases, ns)

        def __new__(cls, name, bases, ns, total=True):
            # Create new typed dict class object.
            # This method is called directly when TypedDict is subclassed,
            # or via _typeddict_new when TypedDict is instantiated. This way
            # TypedDict supports all three syntaxes described in its docstring.
            # Subclasses and instances of TypedDict return actual dictionaries
            # via _dict_new.
            ns['__new__'] = _typeddict_new if name == 'TypedDict' else _dict_new
            tp_dict = super().__new__(cls, name, (dict,), ns)

            annotations = {}
            own_annotations = ns.get('__annotations__', {})
            msg = "TypedDict('Name', {f0: t0, f1: t1, ...}); each t must be a type"
            own_annotations = {
                n: typing._type_check(tp, msg) for n, tp in own_annotations.items()
            }
            required_keys = set()
            optional_keys = set()

            for base in bases:
                annotations.update(base.__dict__.get('__annotations__', {}))
                required_keys.update(base.__dict__.get('__required_keys__', ()))
                optional_keys.update(base.__dict__.get('__optional_keys__', ()))

            annotations.update(own_annotations)
            if PEP_560:
                for annotation_key, annotation_type in own_annotations.items():
                    annotation_origin = get_origin(annotation_type)
                    if annotation_origin is Annotated:
                        annotation_args = get_args(annotation_type)
                        if annotation_args:
                            annotation_type = annotation_args[0]
                            annotation_origin = get_origin(annotation_type)

                    if annotation_origin is Required:
                        required_keys.add(annotation_key)
                    elif annotation_origin is NotRequired:
                        optional_keys.add(annotation_key)
                    elif total:
                        required_keys.add(annotation_key)
                    else:
                        optional_keys.add(annotation_key)
            else:
                own_annotation_keys = set(own_annotations.keys())
                if total:
                    required_keys.update(own_annotation_keys)
                else:
                    optional_keys.update(own_annotation_keys)

            tp_dict.__annotations__ = annotations
            tp_dict.__required_keys__ = frozenset(required_keys)
            tp_dict.__optional_keys__ = frozenset(optional_keys)
            if not hasattr(tp_dict, '__total__'):
                tp_dict.__total__ = total
            return tp_dict

        __instancecheck__ = __subclasscheck__ = _check_fails

    TypedDict = _TypedDictMeta('TypedDict', (dict,), {})
    TypedDict.__module__ = __name__
    TypedDict.__doc__ = \
        """A simple typed name space. At runtime it is equivalent to a plain dict.

        TypedDict creates a dictionary type that expects all of its
        instances to have a certain set of keys, with each key
        associated with a value of a consistent type. This expectation
        is not checked at runtime but is only enforced by type checkers.
        Usage::

            class Point2D(TypedDict):
                x: int
                y: int
                label: str

            a: Point2D = {'x': 1, 'y': 2, 'label': 'good'}  # OK
            b: Point2D = {'z': 3, 'label': 'bad'}           # Fails type check

            assert Point2D(x=1, y=2, label='first') == dict(x=1, y=2, label='first')

        The type info can be accessed via the Point2D.__annotations__ dict, and
        the Point2D.__required_keys__ and Point2D.__optional_keys__ frozensets.
        TypedDict supports two additional equivalent forms::

            Point2D = TypedDict('Point2D', x=int, y=int, label=str)
            Point2D = TypedDict('Point2D', {'x': int, 'y': int, 'label': str})

        The class syntax is only supported in Python 3.6+, while two other
        syntax forms work for Python 2.7 and 3.2+
        """

    if hasattr(typing, "_TypedDictMeta"):
        _TYPEDDICT_TYPES = (typing._TypedDictMeta, _TypedDictMeta)
    else:
        _TYPEDDICT_TYPES = (_TypedDictMeta,)

    def is_typeddict(tp):
        """Check if an annotation is a TypedDict class

        For example::
            class Film(TypedDict):
                title: str
                year: int

            is_typeddict(Film)  # => True
            is_typeddict(Union[list, str])  # => False
        """
        return isinstance(tp, tuple(_TYPEDDICT_TYPES))

if hasattr(typing, "Required"):
    get_type_hints = typing.get_type_hints
elif PEP_560:
    import functools
    import types

    # replaces _strip_annotations()
    def _strip_extras(t):
        """Strips Annotated, Required and NotRequired from a given type."""
        if isinstance(t, _AnnotatedAlias):
            return _strip_extras(t.__origin__)
        if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired):
            return _strip_extras(t.__args__[0])
        if isinstance(t, typing._GenericAlias):
            stripped_args = tuple(_strip_extras(a) for a in t.__args__)
            if stripped_args == t.__args__:
                return t
            return t.copy_with(stripped_args)
        if hasattr(types, "GenericAlias") and isinstance(t, types.GenericAlias):
            stripped_args = tuple(_strip_extras(a) for a in t.__args__)
            if stripped_args == t.__args__:
                return t
            return types.GenericAlias(t.__origin__, stripped_args)
        if hasattr(types, "UnionType") and isinstance(t, types.UnionType):
            stripped_args = tuple(_strip_extras(a) for a in t.__args__)
            if stripped_args == t.__args__:
                return t
            return functools.reduce(operator.or_, stripped_args)

        return t

    def get_type_hints(obj, globalns=None, localns=None, include_extras=False):
        """Return type hints for an object.

        This is often the same as obj.__annotations__, but it handles
        forward references encoded as string literals, adds Optional[t] if a
        default value equal to None is set and recursively replaces all
        'Annotated[T, ...]', 'Required[T]' or 'NotRequired[T]' with 'T'
        (unless 'include_extras=True').

        The argument may be a module, class, method, or function. The annotations
        are returned as a dictionary. For classes, annotations include also
        inherited members.

        TypeError is raised if the argument is not of a type that can contain
        annotations, and an empty dictionary is returned if no annotations are
        present.

        BEWARE -- the behavior of globalns and localns is counterintuitive
        (unless you are familiar with how eval() and exec() work).  The
        search order is locals first, then globals.

        - If no dict arguments are passed, an attempt is made to use the
          globals from obj (or the respective module's globals for classes),
          and these are also used as the locals.  If the object does not appear
          to have globals, an empty dictionary is used.

        - If one dict argument is passed, it is used for both globals and
          locals.

        - If two dict arguments are passed, they specify globals and
          locals, respectively.
        """
        if hasattr(typing, "Annotated"):
            hint = typing.get_type_hints(
                obj, globalns=globalns, localns=localns, include_extras=True
            )
        else:
            hint = typing.get_type_hints(obj, globalns=globalns, localns=localns)
        if include_extras:
            return hint
        return {k: _strip_extras(t) for k, t in hint.items()}


# Python 3.9+ has PEP 593 (Annotated)
if hasattr(typing, 'Annotated'):
    Annotated = typing.Annotated
    # Not exported and not a public API, but needed for get_origin() and get_args()
    # to work.
    _AnnotatedAlias = typing._AnnotatedAlias
# 3.7-3.8
elif PEP_560:
    class _AnnotatedAlias(typing._GenericAlias, _root=True):
        """Runtime representation of an annotated type.

        At its core 'Annotated[t, dec1, dec2, ...]' is an alias for the type 't'
        with extra annotations. The alias behaves like a normal typing alias,
        instantiating is the same as instantiating the underlying type, binding
        it to types is also the same.
        """
        def __init__(self, origin, metadata):
            if isinstance(origin, _AnnotatedAlias):
                metadata = origin.__metadata__ + metadata
                origin = origin.__origin__
            super().__init__(origin, origin)
            self.__metadata__ = metadata

        def copy_with(self, params):
            assert len(params) == 1
            new_type = params[0]
            return _AnnotatedAlias(new_type, self.__metadata__)

        def __repr__(self):
            return (f"typing_extensions.Annotated[{typing._type_repr(self.__origin__)}, "
                    f"{', '.join(repr(a) for a in self.__metadata__)}]")

        def __reduce__(self):
            return operator.getitem, (
                Annotated, (self.__origin__,) + self.__metadata__
            )

        def __eq__(self, other):
            if not isinstance(other, _AnnotatedAlias):
                return NotImplemented
            if self.__origin__ != other.__origin__:
                return False
            return self.__metadata__ == other.__metadata__

        def __hash__(self):
            return hash((self.__origin__, self.__metadata__))

    class Annotated:
        """Add context specific metadata to a type.

        Example: Annotated[int, runtime_check.Unsigned] indicates to the
        hypothetical runtime_check module that this type is an unsigned int.
        Every other consumer of this type can ignore this metadata and treat
        this type as int.

        The first argument to Annotated must be a valid type (and will be in
        the __origin__ field), the remaining arguments are kept as a tuple in
        the __extra__ field.

        Details:

        - It's an error to call `Annotated` with less than two arguments.
        - Nested Annotated are flattened::

            Annotated[Annotated[T, Ann1, Ann2], Ann3] == Annotated[T, Ann1, Ann2, Ann3]

        - Instantiating an annotated type is equivalent to instantiating the
        underlying type::

            Annotated[C, Ann1](5) == C(5)

        - Annotated can be used as a generic type alias::

            Optimized = Annotated[T, runtime.Optimize()]
            Optimized[int] == Annotated[int, runtime.Optimize()]

            OptimizedList = Annotated[List[T], runtime.Optimize()]
            OptimizedList[int] == Annotated[List[int], runtime.Optimize()]
        """

        __slots__ = ()

        def __new__(cls, *args, **kwargs):
            raise TypeError("Type Annotated cannot be instantiated.")

        @typing._tp_cache
        def __class_getitem__(cls, params):
            if not isinstance(params, tuple) or len(params) < 2:
                raise TypeError("Annotated[...] should be used "
                                "with at least two arguments (a type and an "
                                "annotation).")
            allowed_special_forms = (ClassVar, Final)
            if get_origin(params[0]) in allowed_special_forms:
                origin = params[0]
            else:
                msg = "Annotated[t, ...]: t must be a type."
                origin = typing._type_check(params[0], msg)
            metadata = tuple(params[1:])
            return _AnnotatedAlias(origin, metadata)

        def __init_subclass__(cls, *args, **kwargs):
            raise TypeError(
                f"Cannot subclass {cls.__module__}.Annotated"
            )
# 3.6
else:

    def _is_dunder(name):
        """Returns True if name is a __dunder_variable_name__."""
        return len(name) > 4 and name.startswith('__') and name.endswith('__')

    # Prior to Python 3.7 types did not have `copy_with`. A lot of the equality
    # checks, argument expansion etc. are done on the _subs_tre. As a result we
    # can't provide a get_type_hints function that strips out annotations.

    class AnnotatedMeta(typing.GenericMeta):
        """Metaclass for Annotated"""

        def __new__(cls, name, bases, namespace, **kwargs):
            if any(b is not object for b in bases):
                raise TypeError("Cannot subclass " + str(Annotated))
            return super().__new__(cls, name, bases, namespace, **kwargs)

        @property
        def __metadata__(self):
            return self._subs_tree()[2]

        def _tree_repr(self, tree):
            cls, origin, metadata = tree
            if not isinstance(origin, tuple):
                tp_repr = typing._type_repr(origin)
            else:
                tp_repr = origin[0]._tree_repr(origin)
            metadata_reprs = ", ".join(repr(arg) for arg in metadata)
            return f'{cls}[{tp_repr}, {metadata_reprs}]'

        def _subs_tree(self, tvars=None, args=None):  # noqa
            if self is Annotated:
                return Annotated
            res = super()._subs_tree(tvars=tvars, args=args)
            # Flatten nested Annotated
            if isinstance(res[1], tuple) and res[1][0] is Annotated:
                sub_tp = res[1][1]
                sub_annot = res[1][2]
                return (Annotated, sub_tp, sub_annot + res[2])
            return res

        def _get_cons(self):
            """Return the class used to create instance of this type."""
            if self.__origin__ is None:
                raise TypeError("Cannot get the underlying type of a "
                                "non-specialized Annotated type.")
            tree = self._subs_tree()
            while isinstance(tree, tuple) and tree[0] is Annotated:
                tree = tree[1]
            if isinstance(tree, tuple):
                return tree[0]
            else:
                return tree

        @typing._tp_cache
        def __getitem__(self, params):
            if not isinstance(params, tuple):
                params = (params,)
            if self.__origin__ is not None:  # specializing an instantiated type
                return super().__getitem__(params)
            elif not isinstance(params, tuple) or len(params) < 2:
                raise TypeError("Annotated[...] should be instantiated "
                                "with at least two arguments (a type and an "
                                "annotation).")
            else:
                if (
                    isinstance(params[0], typing._TypingBase) and
                    type(params[0]).__name__ == "_ClassVar"
                ):
                    tp = params[0]
                else:
                    msg = "Annotated[t, ...]: t must be a type."
                    tp = typing._type_check(params[0], msg)
                metadata = tuple(params[1:])
            return self.__class__(
                self.__name__,
                self.__bases__,
                _no_slots_copy(self.__dict__),
                tvars=_type_vars((tp,)),
                # Metadata is a tuple so it won't be touched by _replace_args et al.
                args=(tp, metadata),
                origin=self,
            )

        def __call__(self, *args, **kwargs):
            cons = self._get_cons()
            result = cons(*args, **kwargs)
            try:
                result.__orig_class__ = self
            except AttributeError:
                pass
            return result

        def __getattr__(self, attr):
            # For simplicity we just don't relay all dunder names
            if self.__origin__ is not None and not _is_dunder(attr):
                return getattr(self._get_cons(), attr)
            raise AttributeError(attr)

        def __setattr__(self, attr, value):
            if _is_dunder(attr) or attr.startswith('_abc_'):
                super().__setattr__(attr, value)
            elif self.__origin__ is None:
                raise AttributeError(attr)
            else:
                setattr(self._get_cons(), attr, value)

        def __instancecheck__(self, obj):
            raise TypeError("Annotated cannot be used with isinstance().")

        def __subclasscheck__(self, cls):
            raise TypeError("Annotated cannot be used with issubclass().")

    class Annotated(metaclass=AnnotatedMeta):
        """Add context specific metadata to a type.

        Example: Annotated[int, runtime_check.Unsigned] indicates to the
        hypothetical runtime_check module that this type is an unsigned int.
        Every other consumer of this type can ignore this metadata and treat
        this type as int.

        The first argument to Annotated must be a valid type, the remaining
        arguments are kept as a tuple in the __metadata__ field.

        Details:

        - It's an error to call `Annotated` with less than two arguments.
        - Nested Annotated are flattened::

            Annotated[Annotated[T, Ann1, Ann2], Ann3] == Annotated[T, Ann1, Ann2, Ann3]

        - Instantiating an annotated type is equivalent to instantiating the
        underlying type::

            Annotated[C, Ann1](5) == C(5)

        - Annotated can be used as a generic type alias::

            Optimized = Annotated[T, runtime.Optimize()]
            Optimized[int] == Annotated[int, runtime.Optimize()]

            OptimizedList = Annotated[List[T], runtime.Optimize()]
            OptimizedList[int] == Annotated[List[int], runtime.Optimize()]
        """

# Python 3.8 has get_origin() and get_args() but those implementations aren't
# Annotated-aware, so we can't use those. Python 3.9's versions don't support
# ParamSpecArgs and ParamSpecKwargs, so only Python 3.10's versions will do.
if sys.version_info[:2] >= (3, 10):
    get_origin = typing.get_origin
    get_args = typing.get_args
# 3.7-3.9
elif PEP_560:
    try:
        # 3.9+
        from typing import _BaseGenericAlias
    except ImportError:
        _BaseGenericAlias = typing._GenericAlias
    try:
        # 3.9+
        from typing import GenericAlias
    except ImportError:
        GenericAlias = typing._GenericAlias

    def get_origin(tp):
        """Get the unsubscripted version of a type.

        This supports generic types, Callable, Tuple, Union, Literal, Final, ClassVar
        and Annotated. Return None for unsupported types. Examples::

            get_origin(Literal[42]) is Literal
            get_origin(int) is None
            get_origin(ClassVar[int]) is ClassVar
            get_origin(Generic) is Generic
            get_origin(Generic[T]) is Generic
            get_origin(Union[T, int]) is Union
            get_origin(List[Tuple[T, T]][int]) == list
            get_origin(P.args) is P
        """
        if isinstance(tp, _AnnotatedAlias):
            return Annotated
        if isinstance(tp, (typing._GenericAlias, GenericAlias, _BaseGenericAlias,
                           ParamSpecArgs, ParamSpecKwargs)):
            return tp.__origin__
        if tp is typing.Generic:
            return typing.Generic
        return None

    def get_args(tp):
        """Get type arguments with all substitutions performed.

        For unions, basic simplifications used by Union constructor are performed.
        Examples::
            get_args(Dict[str, int]) == (str, int)
            get_args(int) == ()
            get_args(Union[int, Union[T, int], str][int]) == (int, str)
            get_args(Union[int, Tuple[T, int]][str]) == (int, Tuple[str, int])
            get_args(Callable[[], T][int]) == ([], int)
        """
        if isinstance(tp, _AnnotatedAlias):
            return (tp.__origin__,) + tp.__metadata__
        if isinstance(tp, (typing._GenericAlias, GenericAlias)):
            if getattr(tp, "_special", False):
                return ()
            res = tp.__args__
            if get_origin(tp) is collections.abc.Callable and res[0] is not Ellipsis:
                res = (list(res[:-1]), res[-1])
            return res
        return ()


# 3.10+
if hasattr(typing, 'TypeAlias'):
    TypeAlias = typing.TypeAlias
# 3.9
elif sys.version_info[:2] >= (3, 9):
    class _TypeAliasForm(typing._SpecialForm, _root=True):
        def __repr__(self):
            return 'typing_extensions.' + self._name

    @_TypeAliasForm
    def TypeAlias(self, parameters):
        """Special marker indicating that an assignment should
        be recognized as a proper type alias definition by type
        checkers.

        For example::

            Predicate: TypeAlias = Callable[..., bool]

        It's invalid when used anywhere except as in the example above.
        """
        raise TypeError(f"{self} is not subscriptable")
# 3.7-3.8
elif sys.version_info[:2] >= (3, 7):
    class _TypeAliasForm(typing._SpecialForm, _root=True):
        def __repr__(self):
            return 'typing_extensions.' + self._name

    TypeAlias = _TypeAliasForm('TypeAlias',
                               doc="""Special marker indicating that an assignment should
                               be recognized as a proper type alias definition by type
                               checkers.

                               For example::

                                   Predicate: TypeAlias = Callable[..., bool]

                               It's invalid when used anywhere except as in the example
                               above.""")
# 3.6
else:
    class _TypeAliasMeta(typing.TypingMeta):
        """Metaclass for TypeAlias"""

        def __repr__(self):
            return 'typing_extensions.TypeAlias'

    class _TypeAliasBase(typing._FinalTypingBase, metaclass=_TypeAliasMeta, _root=True):
        """Special marker indicating that an assignment should
        be recognized as a proper type alias definition by type
        checkers.

        For example::

            Predicate: TypeAlias = Callable[..., bool]

        It's invalid when used anywhere except as in the example above.
        """
        __slots__ = ()

        def __instancecheck__(self, obj):
            raise TypeError("TypeAlias cannot be used with isinstance().")

        def __subclasscheck__(self, cls):
            raise TypeError("TypeAlias cannot be used with issubclass().")

        def __repr__(self):
            return 'typing_extensions.TypeAlias'

    TypeAlias = _TypeAliasBase(_root=True)


# Python 3.10+ has PEP 612
if hasattr(typing, 'ParamSpecArgs'):
    ParamSpecArgs = typing.ParamSpecArgs
    ParamSpecKwargs = typing.ParamSpecKwargs
# 3.6-3.9
else:
    class _Immutable:
        """Mixin to indicate that object should not be copied."""
        __slots__ = ()

        def __copy__(self):
            return self

        def __deepcopy__(self, memo):
            return self

    class ParamSpecArgs(_Immutable):
        """The args for a ParamSpec object.

        Given a ParamSpec object P, P.args is an instance of ParamSpecArgs.

        ParamSpecArgs objects have a reference back to their ParamSpec:

        P.args.__origin__ is P

        This type is meant for runtime introspection and has no special meaning to
        static type checkers.
        """
        def __init__(self, origin):
            self.__origin__ = origin

        def __repr__(self):
            return f"{self.__origin__.__name__}.args"

        def __eq__(self, other):
            if not isinstance(other, ParamSpecArgs):
                return NotImplemented
            return self.__origin__ == other.__origin__

    class ParamSpecKwargs(_Immutable):
        """The kwargs for a ParamSpec object.

        Given a ParamSpec object P, P.kwargs is an instance of ParamSpecKwargs.

        ParamSpecKwargs objects have a reference back to their ParamSpec:

        P.kwargs.__origin__ is P

        This type is meant for runtime introspection and has no special meaning to
        static type checkers.
        """
        def __init__(self, origin):
            self.__origin__ = origin

        def __repr__(self):
            return f"{self.__origin__.__name__}.kwargs"

        def __eq__(self, other):
            if not isinstance(other, ParamSpecKwargs):
                return NotImplemented
            return self.__origin__ == other.__origin__

# 3.10+
if hasattr(typing, 'ParamSpec'):
    ParamSpec = typing.ParamSpec
# 3.6-3.9
else:

    # Inherits from list as a workaround for Callable checks in Python < 3.9.2.
    class ParamSpec(list):
        """Parameter specification variable.

        Usage::

           P = ParamSpec('P')

        Parameter specification variables exist primarily for the benefit of static
        type checkers.  They are used to forward the parameter types of one
        callable to another callable, a pattern commonly found in higher order
        functions and decorators.  They are only valid when used in ``Concatenate``,
        or s the first argument to ``Callable``. In Python 3.10 and higher,
        they are also supported in user-defined Generics at runtime.
        See class Generic for more information on generic types.  An
        example for annotating a decorator::

           T = TypeVar('T')
           P = ParamSpec('P')

           def add_logging(f: Callable[P, T]) -> Callable[P, T]:
               '''A type-safe decorator to add logging to a function.'''
               def inner(*args: P.args, **kwargs: P.kwargs) -> T:
                   logging.info(f'{f.__name__} was called')
                   return f(*args, **kwargs)
               return inner

           @add_logging
           def add_two(x: float, y: float) -> float:
               '''Add two numbers together.'''
               return x + y

        Parameter specification variables defined with covariant=True or
        contravariant=True can be used to declare covariant or contravariant
        generic types.  These keyword arguments are valid, but their actual semantics
        are yet to be decided.  See PEP 612 for details.

        Parameter specification variables can be introspected. e.g.:

           P.__name__ == 'T'
           P.__bound__ == None
           P.__covariant__ == False
           P.__contravariant__ == False

        Note that only parameter specification variables defined in global scope can
        be pickled.
        """

        # Trick Generic __parameters__.
        __class__ = typing.TypeVar

        @property
        def args(self):
            return ParamSpecArgs(self)

        @property
        def kwargs(self):
            return ParamSpecKwargs(self)

        def __init__(self, name, *, bound=None, covariant=False, contravariant=False):
            super().__init__([self])
            self.__name__ = name
            self.__covariant__ = bool(covariant)
            self.__contravariant__ = bool(contravariant)
            if bound:
                self.__bound__ = typing._type_check(bound, 'Bound must be a type.')
            else:
                self.__bound__ = None

            # for pickling:
            try:
                def_mod = sys._getframe(1).f_globals.get('__name__', '__main__')
            except (AttributeError, ValueError):
                def_mod = None
            if def_mod != 'typing_extensions':
                self.__module__ = def_mod

        def __repr__(self):
            if self.__covariant__:
                prefix = '+'
            elif self.__contravariant__:
                prefix = '-'
            else:
                prefix = '~'
            return prefix + self.__name__

        def __hash__(self):
            return object.__hash__(self)

        def __eq__(self, other):
            return self is other

        def __reduce__(self):
            return self.__name__

        # Hack to get typing._type_check to pass.
        def __call__(self, *args, **kwargs):
            pass

        if not PEP_560:
            # Only needed in 3.6.
            def _get_type_vars(self, tvars):
                if self not in tvars:
                    tvars.append(self)


# 3.6-3.9
if not hasattr(typing, 'Concatenate'):
    # Inherits from list as a workaround for Callable checks in Python < 3.9.2.
    class _ConcatenateGenericAlias(list):

        # Trick Generic into looking into this for __parameters__.
        if PEP_560:
            __class__ = typing._GenericAlias
        else:
            __class__ = typing._TypingBase

        # Flag in 3.8.
        _special = False
        # Attribute in 3.6 and earlier.
        _gorg = typing.Generic

        def __init__(self, origin, args):
            super().__init__(args)
            self.__origin__ = origin
            self.__args__ = args

        def __repr__(self):
            _type_repr = typing._type_repr
            return (f'{_type_repr(self.__origin__)}'
                    f'[{", ".join(_type_repr(arg) for arg in self.__args__)}]')

        def __hash__(self):
            return hash((self.__origin__, self.__args__))

        # Hack to get typing._type_check to pass in Generic.
        def __call__(self, *args, **kwargs):
            pass

        @property
        def __parameters__(self):
            return tuple(
                tp for tp in self.__args__ if isinstance(tp, (typing.TypeVar, ParamSpec))
            )

        if not PEP_560:
            # Only required in 3.6.
            def _get_type_vars(self, tvars):
                if self.__origin__ and self.__parameters__:
                    typing._get_type_vars(self.__parameters__, tvars)


# 3.6-3.9
@typing._tp_cache
def _concatenate_getitem(self, parameters):
    if parameters == ():
        raise TypeError("Cannot take a Concatenate of no types.")
    if not isinstance(parameters, tuple):
        parameters = (parameters,)
    if not isinstance(parameters[-1], ParamSpec):
        raise TypeError("The last parameter to Concatenate should be a "
                        "ParamSpec variable.")
    msg = "Concatenate[arg, ...]: each arg must be a type."
    parameters = tuple(typing._type_check(p, msg) for p in parameters)
    return _ConcatenateGenericAlias(self, parameters)


# 3.10+
if hasattr(typing, 'Concatenate'):
    Concatenate = typing.Concatenate
    _ConcatenateGenericAlias = typing._ConcatenateGenericAlias # noqa
# 3.9
elif sys.version_info[:2] >= (3, 9):
    @_TypeAliasForm
    def Concatenate(self, parameters):
        """Used in conjunction with ``ParamSpec`` and ``Callable`` to represent a
        higher order function which adds, removes or transforms parameters of a
        callable.

        For example::

           Callable[Concatenate[int, P], int]

        See PEP 612 for detailed information.
        """
        return _concatenate_getitem(self, parameters)
# 3.7-8
elif sys.version_info[:2] >= (3, 7):
    class _ConcatenateForm(typing._SpecialForm, _root=True):
        def __repr__(self):
            return 'typing_extensions.' + self._name

        def __getitem__(self, parameters):
            return _concatenate_getitem(self, parameters)

    Concatenate = _ConcatenateForm(
        'Concatenate',
        doc="""Used in conjunction with ``ParamSpec`` and ``Callable`` to represent a
        higher order function which adds, removes or transforms parameters of a
        callable.

        For example::

           Callable[Concatenate[int, P], int]

        See PEP 612 for detailed information.
        """)
# 3.6
else:
    class _ConcatenateAliasMeta(typing.TypingMeta):
        """Metaclass for Concatenate."""

        def __repr__(self):
            return 'typing_extensions.Concatenate'

    class _ConcatenateAliasBase(typing._FinalTypingBase,
                                metaclass=_ConcatenateAliasMeta,
                                _root=True):
        """Used in conjunction with ``ParamSpec`` and ``Callable`` to represent a
        higher order function which adds, removes or transforms parameters of a
        callable.

        For example::

           Callable[Concatenate[int, P], int]

        See PEP 612 for detailed information.
        """
        __slots__ = ()

        def __instancecheck__(self, obj):
            raise TypeError("Concatenate cannot be used with isinstance().")

        def __subclasscheck__(self, cls):
            raise TypeError("Concatenate cannot be used with issubclass().")

        def __repr__(self):
            return 'typing_extensions.Concatenate'

        def __getitem__(self, parameters):
            return _concatenate_getitem(self, parameters)

    Concatenate = _ConcatenateAliasBase(_root=True)

# 3.10+
if hasattr(typing, 'TypeGuard'):
    TypeGuard = typing.TypeGuard
# 3.9
elif sys.version_info[:2] >= (3, 9):
    class _TypeGuardForm(typing._SpecialForm, _root=True):
        def __repr__(self):
            return 'typing_extensions.' + self._name

    @_TypeGuardForm
    def TypeGuard(self, parameters):
        """Special typing form used to annotate the return type of a user-defined
        type guard function.  ``TypeGuard`` only accepts a single type argument.
        At runtime, functions marked this way should return a boolean.

        ``TypeGuard`` aims to benefit *type narrowing* -- a technique used by static
        type checkers to determine a more precise type of an expression within a
        program's code flow.  Usually type narrowing is done by analyzing
        conditional code flow and applying the narrowing to a block of code.  The
        conditional expression here is sometimes referred to as a "type guard".

        Sometimes it would be convenient to use a user-defined boolean function
        as a type guard.  Such a function should use ``TypeGuard[...]`` as its
        return type to alert static type checkers to this intention.

        Using  ``-> TypeGuard`` tells the static type checker that for a given
        function:

        1. The return value is a boolean.
        2. If the return value is ``True``, the type of its argument
        is the type inside ``TypeGuard``.

        For example::

            def is_str(val: Union[str, float]):
                # "isinstance" type guard
                if isinstance(val, str):
                    # Type of ``val`` is narrowed to ``str``
                    ...
                else:
                    # Else, type of ``val`` is narrowed to ``float``.
                    ...

        Strict type narrowing is not enforced -- ``TypeB`` need not be a narrower
        form of ``TypeA`` (it can even be a wider form) and this may lead to
        type-unsafe results.  The main reason is to allow for things like
        narrowing ``List[object]`` to ``List[str]`` even though the latter is not
        a subtype of the former, since ``List`` is invariant.  The responsibility of
        writing type-safe type guards is left to the user.

        ``TypeGuard`` also works with type variables.  For more information, see
        PEP 647 (User-Defined Type Guards).
        """
        item = typing._type_check(parameters, f'{self} accepts only single type.')
        return typing._GenericAlias(self, (item,))
# 3.7-3.8
elif sys.version_info[:2] >= (3, 7):
    class _TypeGuardForm(typing._SpecialForm, _root=True):

        def __repr__(self):
            return 'typing_extensions.' + self._name

        def __getitem__(self, parameters):
            item = typing._type_check(parameters,
                                      f'{self._name} accepts only a single type')
            return typing._GenericAlias(self, (item,))

    TypeGuard = _TypeGuardForm(
        'TypeGuard',
        doc="""Special typing form used to annotate the return type of a user-defined
        type guard function.  ``TypeGuard`` only accepts a single type argument.
        At runtime, functions marked this way should return a boolean.

        ``TypeGuard`` aims to benefit *type narrowing* -- a technique used by static
        type checkers to determine a more precise type of an expression within a
        program's code flow.  Usually type narrowing is done by analyzing
        conditional code flow and applying the narrowing to a block of code.  The
        conditional expression here is sometimes referred to as a "type guard".

        Sometimes it would be convenient to use a user-defined boolean function
        as a type guard.  Such a function should use ``TypeGuard[...]`` as its
        return type to alert static type checkers to this intention.

        Using  ``-> TypeGuard`` tells the static type checker that for a given
        function:

        1. The return value is a boolean.
        2. If the return value is ``True``, the type of its argument
        is the type inside ``TypeGuard``.

        For example::

            def is_str(val: Union[str, float]):
                # "isinstance" type guard
                if isinstance(val, str):
                    # Type of ``val`` is narrowed to ``str``
                    ...
                else:
                    # Else, type of ``val`` is narrowed to ``float``.
                    ...

        Strict type narrowing is not enforced -- ``TypeB`` need not be a narrower
        form of ``TypeA`` (it can even be a wider form) and this may lead to
        type-unsafe results.  The main reason is to allow for things like
        narrowing ``List[object]`` to ``List[str]`` even though the latter is not
        a subtype of the former, since ``List`` is invariant.  The responsibility of
        writing type-safe type guards is left to the user.

        ``TypeGuard`` also works with type variables.  For more information, see
        PEP 647 (User-Defined Type Guards).
        """)
# 3.6
else:
    class _TypeGuard(typing._FinalTypingBase, _root=True):
        """Special typing form used to annotate the return type of a user-defined
        type guard function.  ``TypeGuard`` only accepts a single type argument.
        At runtime, functions marked this way should return a boolean.

        ``TypeGuard`` aims to benefit *type narrowing* -- a technique used by static
        type checkers to determine a more precise type of an expression within a
        program's code flow.  Usually type narrowing is done by analyzing
        conditional code flow and applying the narrowing to a block of code.  The
        conditional expression here is sometimes referred to as a "type guard".

        Sometimes it would be convenient to use a user-defined boolean function
        as a type guard.  Such a function should use ``TypeGuard[...]`` as its
        return type to alert static type checkers to this intention.

        Using  ``-> TypeGuard`` tells the static type checker that for a given
        function:

        1. The return value is a boolean.
        2. If the return value is ``True``, the type of its argument
        is the type inside ``TypeGuard``.

        For example::

            def is_str(val: Union[str, float]):
                # "isinstance" type guard
                if isinstance(val, str):
                    # Type of ``val`` is narrowed to ``str``
                    ...
                else:
                    # Else, type of ``val`` is narrowed to ``float``.
                    ...

        Strict type narrowing is not enforced -- ``TypeB`` need not be a narrower
        form of ``TypeA`` (it can even be a wider form) and this may lead to
        type-unsafe results.  The main reason is to allow for things like
        narrowing ``List[object]`` to ``List[str]`` even though the latter is not
        a subtype of the former, since ``List`` is invariant.  The responsibility of
        writing type-safe type guards is left to the user.

        ``TypeGuard`` also works with type variables.  For more information, see
        PEP 647 (User-Defined Type Guards).
        """

        __slots__ = ('__type__',)

        def __init__(self, tp=None, **kwds):
            self.__type__ = tp

        def __getitem__(self, item):
            cls = type(self)
            if self.__type__ is None:
                return cls(typing._type_check(item,
                           f'{cls.__name__[1:]} accepts only a single type.'),
                           _root=True)
            raise TypeError(f'{cls.__name__[1:]} cannot be further subscripted')

        def _eval_type(self, globalns, localns):
            new_tp = typing._eval_type(self.__type__, globalns, localns)
            if new_tp == self.__type__:
                return self
            return type(self)(new_tp, _root=True)

        def __repr__(self):
            r = super().__repr__()
            if self.__type__ is not None:
                r += f'[{typing._type_repr(self.__type__)}]'
            return r

        def __hash__(self):
            return hash((type(self).__name__, self.__type__))

        def __eq__(self, other):
            if not isinstance(other, _TypeGuard):
                return NotImplemented
            if self.__type__ is not None:
                return self.__type__ == other.__type__
            return self is other

    TypeGuard = _TypeGuard(_root=True)


if sys.version_info[:2] >= (3, 7):
    # Vendored from cpython typing._SpecialFrom
    class _SpecialForm(typing._Final, _root=True):
        __slots__ = ('_name', '__doc__', '_getitem')

        def __init__(self, getitem):
            self._getitem = getitem
            self._name = getitem.__name__
            self.__doc__ = getitem.__doc__

        def __getattr__(self, item):
            if item in {'__name__', '__qualname__'}:
                return self._name

            raise AttributeError(item)

        def __mro_entries__(self, bases):
            raise TypeError(f"Cannot subclass {self!r}")

        def __repr__(self):
            return f'typing_extensions.{self._name}'

        def __reduce__(self):
            return self._name

        def __call__(self, *args, **kwds):
            raise TypeError(f"Cannot instantiate {self!r}")

        def __or__(self, other):
            return typing.Union[self, other]

        def __ror__(self, other):
            return typing.Union[other, self]

        def __instancecheck__(self, obj):
            raise TypeError(f"{self} cannot be used with isinstance()")

        def __subclasscheck__(self, cls):
            raise TypeError(f"{self} cannot be used with issubclass()")

        @typing._tp_cache
        def __getitem__(self, parameters):
            return self._getitem(self, parameters)


if hasattr(typing, "LiteralString"):
    LiteralString = typing.LiteralString
elif sys.version_info[:2] >= (3, 7):
    @_SpecialForm
    def LiteralString(self, params):
        """Represents an arbitrary literal string.

        Example::

          from typing_extensions import LiteralString

          def query(sql: LiteralString) -> ...:
              ...

          query("SELECT * FROM table")  # ok
          query(f"SELECT * FROM {input()}")  # not ok

        See PEP 675 for details.

        """
        raise TypeError(f"{self} is not subscriptable")
else:
    class _LiteralString(typing._FinalTypingBase, _root=True):
        """Represents an arbitrary literal string.

        Example::

          from typing_extensions import LiteralString

          def query(sql: LiteralString) -> ...:
              ...

          query("SELECT * FROM table")  # ok
          query(f"SELECT * FROM {input()}")  # not ok

        See PEP 675 for details.

        """

        __slots__ = ()

        def __instancecheck__(self, obj):
            raise TypeError(f"{self} cannot be used with isinstance().")

        def __subclasscheck__(self, cls):
            raise TypeError(f"{self} cannot be used with issubclass().")

    LiteralString = _LiteralString(_root=True)


if hasattr(typing, "Self"):
    Self = typing.Self
elif sys.version_info[:2] >= (3, 7):
    @_SpecialForm
    def Self(self, params):
        """Used to spell the type of "self" in classes.

        Example::

          from typing import Self

          class ReturnsSelf:
              def parse(self, data: bytes) -> Self:
                  ...
                  return self

        """

        raise TypeError(f"{self} is not subscriptable")
else:
    class _Self(typing._FinalTypingBase, _root=True):
        """Used to spell the type of "self" in classes.

        Example::

          from typing import Self

          class ReturnsSelf:
              def parse(self, data: bytes) -> Self:
                  ...
                  return self

        """

        __slots__ = ()

        def __instancecheck__(self, obj):
            raise TypeError(f"{self} cannot be used with isinstance().")

        def __subclasscheck__(self, cls):
            raise TypeError(f"{self} cannot be used with issubclass().")

    Self = _Self(_root=True)


if hasattr(typing, "Never"):
    Never = typing.Never
elif sys.version_info[:2] >= (3, 7):
    @_SpecialForm
    def Never(self, params):
        """The bottom type, a type that has no members.

        This can be used to define a function that should never be
        called, or a function that never returns::

            from typing_extensions import Never

            def never_call_me(arg: Never) -> None:
                pass

            def int_or_str(arg: int | str) -> None:
                never_call_me(arg)  # type checker error
                match arg:
                    case int():
                        print("It's an int")
                    case str():
                        print("It's a str")
                    case _:
                        never_call_me(arg)  # ok, arg is of type Never

        """

        raise TypeError(f"{self} is not subscriptable")
else:
    class _Never(typing._FinalTypingBase, _root=True):
        """The bottom type, a type that has no members.

        This can be used to define a function that should never be
        called, or a function that never returns::

            from typing_extensions import Never

            def never_call_me(arg: Never) -> None:
                pass

            def int_or_str(arg: int | str) -> None:
                never_call_me(arg)  # type checker error
                match arg:
                    case int():
                        print("It's an int")
                    case str():
                        print("It's a str")
                    case _:
                        never_call_me(arg)  # ok, arg is of type Never

        """

        __slots__ = ()

        def __instancecheck__(self, obj):
            raise TypeError(f"{self} cannot be used with isinstance().")

        def __subclasscheck__(self, cls):
            raise TypeError(f"{self} cannot be used with issubclass().")

    Never = _Never(_root=True)


if hasattr(typing, 'Required'):
    Required = typing.Required
    NotRequired = typing.NotRequired
elif sys.version_info[:2] >= (3, 9):
    class _ExtensionsSpecialForm(typing._SpecialForm, _root=True):
        def __repr__(self):
            return 'typing_extensions.' + self._name

    @_ExtensionsSpecialForm
    def Required(self, parameters):
        """A special typing construct to mark a key of a total=False TypedDict
        as required. For example:

            class Movie(TypedDict, total=False):
                title: Required[str]
                year: int

            m = Movie(
                title='The Matrix',  # typechecker error if key is omitted
                year=1999,
            )

        There is no runtime checking that a required key is actually provided
        when instantiating a related TypedDict.
        """
        item = typing._type_check(parameters, f'{self._name} accepts only single type')
        return typing._GenericAlias(self, (item,))

    @_ExtensionsSpecialForm
    def NotRequired(self, parameters):
        """A special typing construct to mark a key of a TypedDict as
        potentially missing. For example:

            class Movie(TypedDict):
                title: str
                year: NotRequired[int]

            m = Movie(
                title='The Matrix',  # typechecker error if key is omitted
                year=1999,
            )
        """
        item = typing._type_check(parameters, f'{self._name} accepts only single type')
        return typing._GenericAlias(self, (item,))

elif sys.version_info[:2] >= (3, 7):
    class _RequiredForm(typing._SpecialForm, _root=True):
        def __repr__(self):
            return 'typing_extensions.' + self._name

        def __getitem__(self, parameters):
            item = typing._type_check(parameters,
                                      '{} accepts only single type'.format(self._name))
            return typing._GenericAlias(self, (item,))

    Required = _RequiredForm(
        'Required',
        doc="""A special typing construct to mark a key of a total=False TypedDict
        as required. For example:

            class Movie(TypedDict, total=False):
                title: Required[str]
                year: int

            m = Movie(
                title='The Matrix',  # typechecker error if key is omitted
                year=1999,
            )

        There is no runtime checking that a required key is actually provided
        when instantiating a related TypedDict.
        """)
    NotRequired = _RequiredForm(
        'NotRequired',
        doc="""A special typing construct to mark a key of a TypedDict as
        potentially missing. For example:

            class Movie(TypedDict):
                title: str
                year: NotRequired[int]

            m = Movie(
                title='The Matrix',  # typechecker error if key is omitted
                year=1999,
            )
        """)
else:
    # NOTE: Modeled after _Final's implementation when _FinalTypingBase available
    class _MaybeRequired(typing._FinalTypingBase, _root=True):
        __slots__ = ('__type__',)

        def __init__(self, tp=None, **kwds):
            self.__type__ = tp

        def __getitem__(self, item):
            cls = type(self)
            if self.__type__ is None:
                return cls(typing._type_check(item,
                           '{} accepts only single type.'.format(cls.__name__[1:])),
                           _root=True)
            raise TypeError('{} cannot be further subscripted'
                            .format(cls.__name__[1:]))

        def _eval_type(self, globalns, localns):
            new_tp = typing._eval_type(self.__type__, globalns, localns)
            if new_tp == self.__type__:
                return self
            return type(self)(new_tp, _root=True)

        def __repr__(self):
            r = super().__repr__()
            if self.__type__ is not None:
                r += '[{}]'.format(typing._type_repr(self.__type__))
            return r

        def __hash__(self):
            return hash((type(self).__name__, self.__type__))

        def __eq__(self, other):
            if not isinstance(other, type(self)):
                return NotImplemented
            if self.__type__ is not None:
                return self.__type__ == other.__type__
            return self is other

    class _Required(_MaybeRequired, _root=True):
        """A special typing construct to mark a key of a total=False TypedDict
        as required. For example:

            class Movie(TypedDict, total=False):
                title: Required[str]
                year: int

            m = Movie(
                title='The Matrix',  # typechecker error if key is omitted
                year=1999,
            )

        There is no runtime checking that a required key is actually provided
        when instantiating a related TypedDict.
        """

    class _NotRequired(_MaybeRequired, _root=True):
        """A special typing construct to mark a key of a TypedDict as
        potentially missing. For example:

            class Movie(TypedDict):
                title: str
                year: NotRequired[int]

            m = Movie(
                title='The Matrix',  # typechecker error if key is omitted
                year=1999,
            )
        """

    Required = _Required(_root=True)
    NotRequired = _NotRequired(_root=True)


if sys.version_info[:2] >= (3, 9):
    class _UnpackSpecialForm(typing._SpecialForm, _root=True):
        def __repr__(self):
            return 'typing_extensions.' + self._name

    class _UnpackAlias(typing._GenericAlias, _root=True):
        __class__ = typing.TypeVar

    @_UnpackSpecialForm
    def Unpack(self, parameters):
        """A special typing construct to unpack a variadic type. For example:

            Shape = TypeVarTuple('Shape')
            Batch = NewType('Batch', int)

            def add_batch_axis(
                x: Array[Unpack[Shape]]
            ) -> Array[Batch, Unpack[Shape]]: ...

        """
        item = typing._type_check(parameters, f'{self._name} accepts only single type')
        return _UnpackAlias(self, (item,))

    def _is_unpack(obj):
        return isinstance(obj, _UnpackAlias)

elif sys.version_info[:2] >= (3, 7):
    class _UnpackAlias(typing._GenericAlias, _root=True):
        __class__ = typing.TypeVar

    class _UnpackForm(typing._SpecialForm, _root=True):
        def __repr__(self):
            return 'typing_extensions.' + self._name

        def __getitem__(self, parameters):
            item = typing._type_check(parameters,
                                      f'{self._name} accepts only single type')
            return _UnpackAlias(self, (item,))

    Unpack = _UnpackForm(
        'Unpack',
        doc="""A special typing construct to unpack a variadic type. For example:

            Shape = TypeVarTuple('Shape')
            Batch = NewType('Batch', int)

            def add_batch_axis(
                x: Array[Unpack[Shape]]
            ) -> Array[Batch, Unpack[Shape]]: ...

        """)

    def _is_unpack(obj):
        return isinstance(obj, _UnpackAlias)

else:
    # NOTE: Modeled after _Final's implementation when _FinalTypingBase available
    class _Unpack(typing._FinalTypingBase, _root=True):
        """A special typing construct to unpack a variadic type. For example:

            Shape = TypeVarTuple('Shape')
            Batch = NewType('Batch', int)

            def add_batch_axis(
                x: Array[Unpack[Shape]]
            ) -> Array[Batch, Unpack[Shape]]: ...

        """
        __slots__ = ('__type__',)
        __class__ = typing.TypeVar

        def __init__(self, tp=None, **kwds):
            self.__type__ = tp

        def __getitem__(self, item):
            cls = type(self)
            if self.__type__ is None:
                return cls(typing._type_check(item,
                           'Unpack accepts only single type.'),
                           _root=True)
            raise TypeError('Unpack cannot be further subscripted')

        def _eval_type(self, globalns, localns):
            new_tp = typing._eval_type(self.__type__, globalns, localns)
            if new_tp == self.__type__:
                return self
            return type(self)(new_tp, _root=True)

        def __repr__(self):
            r = super().__repr__()
            if self.__type__ is not None:
                r += '[{}]'.format(typing._type_repr(self.__type__))
            return r

        def __hash__(self):
            return hash((type(self).__name__, self.__type__))

        def __eq__(self, other):
            if not isinstance(other, _Unpack):
                return NotImplemented
            if self.__type__ is not None:
                return self.__type__ == other.__type__
            return self is other

        # For 3.6 only
        def _get_type_vars(self, tvars):
            self.__type__._get_type_vars(tvars)

    Unpack = _Unpack(_root=True)

    def _is_unpack(obj):
        return isinstance(obj, _Unpack)


class TypeVarTuple:
    """Type variable tuple.

    Usage::

        Ts = TypeVarTuple('Ts')

    In the same way that a normal type variable is a stand-in for a single
    type such as ``int``, a type variable *tuple* is a stand-in for a *tuple* type such as
    ``Tuple[int, str]``.

    Type variable tuples can be used in ``Generic`` declarations.
    Consider the following example::

        class Array(Generic[*Ts]): ...

    The ``Ts`` type variable tuple here behaves like ``tuple[T1, T2]``,
    where ``T1`` and ``T2`` are type variables. To use these type variables
    as type parameters of ``Array``, we must *unpack* the type variable tuple using
    the star operator: ``*Ts``. The signature of ``Array`` then behaves
    as if we had simply written ``class Array(Generic[T1, T2]): ...``.
    In contrast to ``Generic[T1, T2]``, however, ``Generic[*Shape]`` allows
    us to parameterise the class with an *arbitrary* number of type parameters.

    Type variable tuples can be used anywhere a normal ``TypeVar`` can.
    This includes class definitions, as shown above, as well as function
    signatures and variable annotations::

        class Array(Generic[*Ts]):

            def __init__(self, shape: Tuple[*Ts]):
                self._shape: Tuple[*Ts] = shape

            def get_shape(self) -> Tuple[*Ts]:
                return self._shape

        shape = (Height(480), Width(640))
        x: Array[Height, Width] = Array(shape)
        y = abs(x)  # Inferred type is Array[Height, Width]
        z = x + x   #        ...    is Array[Height, Width]
        x.get_shape()  #     ...    is tuple[Height, Width]

    """

    # Trick Generic __parameters__.
    __class__ = typing.TypeVar

    def __iter__(self):
        yield self.__unpacked__

    def __init__(self, name):
        self.__name__ = name

        # for pickling:
        try:
            def_mod = sys._getframe(1).f_globals.get('__name__', '__main__')
        except (AttributeError, ValueError):
            def_mod = None
        if def_mod != 'typing_extensions':
            self.__module__ = def_mod

        self.__unpacked__ = Unpack[self]

    def __repr__(self):
        return self.__name__

    def __hash__(self):
        return object.__hash__(self)

    def __eq__(self, other):
        return self is other

    def __reduce__(self):
        return self.__name__

    def __init_subclass__(self, *args, **kwds):
        if '_root' not in kwds:
            raise TypeError("Cannot subclass special typing classes")

    if not PEP_560:
        # Only needed in 3.6.
        def _get_type_vars(self, tvars):
            if self not in tvars:
                tvars.append(self)


if hasattr(typing, "reveal_type"):
    reveal_type = typing.reveal_type
else:
    def reveal_type(__obj: T) -> T:
        """Reveal the inferred type of a variable.

        When a static type checker encounters a call to ``reveal_type()``,
        it will emit the inferred type of the argument::

            x: int = 1
            reveal_type(x)

        Running a static type checker (e.g., ``mypy``) on this example
        will produce output similar to 'Revealed type is "builtins.int"'.

        At runtime, the function prints the runtime type of the
        argument and returns it unchanged.

        """
        print(f"Runtime type is {type(__obj).__name__!r}", file=sys.stderr)
        return __obj


if hasattr(typing, "assert_never"):
    assert_never = typing.assert_never
else:
    def assert_never(__arg: Never) -> Never:
        """Assert to the type checker that a line of code is unreachable.

        Example::

            def int_or_str(arg: int | str) -> None:
                match arg:
                    case int():
                        print("It's an int")
                    case str():
                        print("It's a str")
                    case _:
                        assert_never(arg)

        If a type checker finds that a call to assert_never() is
        reachable, it will emit an error.

        At runtime, this throws an exception when called.

        """
        raise AssertionError("Expected code to be unreachable")


if hasattr(typing, 'dataclass_transform'):
    dataclass_transform = typing.dataclass_transform
else:
    def dataclass_transform(
        *,
        eq_default: bool = True,
        order_default: bool = False,
        kw_only_default: bool = False,
        field_descriptors: typing.Tuple[
            typing.Union[typing.Type[typing.Any], typing.Callable[..., typing.Any]],
            ...
        ] = (),
    ) -> typing.Callable[[T], T]:
        """Decorator that marks a function, class, or metaclass as providing
        dataclass-like behavior.

        Example:

            from typing_extensions import dataclass_transform

            _T = TypeVar("_T")

            # Used on a decorator function
            @dataclass_transform()
            def create_model(cls: type[_T]) -> type[_T]:
                ...
                return cls

            @create_model
            class CustomerModel:
                id: int
                name: str

            # Used on a base class
            @dataclass_transform()
            class ModelBase: ...

            class CustomerModel(ModelBase):
                id: int
                name: str

            # Used on a metaclass
            @dataclass_transform()
            class ModelMeta(type): ...

            class ModelBase(metaclass=ModelMeta): ...

            class CustomerModel(ModelBase):
                id: int
                name: str

        Each of the ``CustomerModel`` classes defined in this example will now
        behave similarly to a dataclass created with the ``@dataclasses.dataclass``
        decorator. For example, the type checker will synthesize an ``__init__``
        method.

        The arguments to this decorator can be used to customize this behavior:
        - ``eq_default`` indicates whether the ``eq`` parameter is assumed to be
          True or False if it is omitted by the caller.
        - ``order_default`` indicates whether the ``order`` parameter is
          assumed to be True or False if it is omitted by the caller.
        - ``kw_only_default`` indicates whether the ``kw_only`` parameter is
          assumed to be True or False if it is omitted by the caller.
        - ``field_descriptors`` specifies a static list of supported classes
          or functions, that describe fields, similar to ``dataclasses.field()``.

        At runtime, this decorator records its arguments in the
        ``__dataclass_transform__`` attribute on the decorated object.

        See PEP 681 for details.

        """
        def decorator(cls_or_fn):
            cls_or_fn.__dataclass_transform__ = {
                "eq_default": eq_default,
                "order_default": order_default,
                "kw_only_default": kw_only_default,
                "field_descriptors": field_descriptors,
            }
            return cls_or_fn
        return decorator


# We have to do some monkey patching to deal with the dual nature of
# Unpack/TypeVarTuple:
# - We want Unpack to be a kind of TypeVar so it gets accepted in
#   Generic[Unpack[Ts]]
# - We want it to *not* be treated as a TypeVar for the purposes of
#   counting generic parameters, so that when we subscript a generic,
#   the runtime doesn't try to substitute the Unpack with the subscripted type.
if not hasattr(typing, "TypeVarTuple"):
    typing._collect_type_vars = _collect_type_vars
    typing._check_generic = _check_generic

social_vdgnn/lib/python3.6/site-packages/typing_extensions.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/sre_compile.py
social_vdgnn/lib/python3.6/sre_compile.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/sre_constants.py
social_vdgnn/lib/python3.6/sre_constants.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/sre_parse.py
social_vdgnn/lib/python3.6/sre_parse.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/stat.py
social_vdgnn/lib/python3.6/stat.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/struct.py
social_vdgnn/lib/python3.6/struct.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/tarfile.py
social_vdgnn/lib/python3.6/tarfile.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/tempfile.py
social_vdgnn/lib/python3.6/tempfile.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/token.py
social_vdgnn/lib/python3.6/token.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/tokenize.py
social_vdgnn/lib/python3.6/tokenize.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/types.py
social_vdgnn/lib/python3.6/types.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/warnings.py
social_vdgnn/lib/python3.6/warnings.py
eliatrevisan/IA-MPPI-LBM
IntxLNK/usr/lib/python3.6/weakref.py
social_vdgnn/lib/python3.6/weakref.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import pylab as pl
import sys
if sys.version_info[0] < 3:
  import Trajectory as traj
  import OccupancyGrid as occ
  import AgentData as ad
  import Support as sup
else:
  import src.data_utils.Trajectory as traj
  import src.data_utils.OccupancyGrid as occ
  import src.data_utils.AgentData as ad
  import src.data_utils.Support as sup


class AgentContainer():
  """
  Container for all agents in one demonstration. 
  This data structure provides an interface to easily store and access data of various agents.
  It contains a dictionary of the agent data per agent id and the static occupancy grid.
  """
  def __init__(self):
    self.agent_data = {}
    self.occupancy_grid = occ.OccupancyGrid()
    
  def setOccupancyGrid(self, occupancy_grid):
    self.occupancy_grid = occupancy_grid
  
  def removeAgent(self, id):
    self.agent_data.pop(id)
  
  def addDataSample(self, id, timestamp, pose, vel, goal):
    """
    Add a data sample to the database consisting of agent data per ID.
    """
    # set up new agent if id has not been registered yet
    if id not in self.agent_data.keys():
      self.agent_data[id] = ad.AgentData(id=id)
    
    # add the data to the agent with the right id
    self.agent_data[id].addSample(timestamp, pose, vel, tuple(goal))
  
  
  # Getters and setters
  def getAgentIDs(self):
    """
    Return list of IDs of contained agents.
    """
    return [int(k) for k in self.agent_data.keys()]
  
  def getNumberOfAgents(self):
    """
    Get the number of agents in the container. 
    """
    return len(self.agent_data.keys())
  
  def getAgentTrajectories(self, id):
    """
    Get all trajectories of a certain agent.
    """
    return self.agent_data[id].trajectories
  
  def getAgentData(self, id):
    """
    Return the agent data of agent with ID=id
    """
    return self.agent_data[id]
  
  def getNumberOfTrajectoriesForAgent(self, id):
    """
    Return the number of trajectories (start -> goal) of agent with ID=id.
    """
    return len(self.agent_data[id].trajectories)
  
  def getTrajectorySetForTime(self, query_time):
    """
    Get all trajectories which are available at a certain time.
    """
    traj_set = {}
    for id in self.agent_data.keys():
      t = self.agent_data[id].getTrajectoryForTime(query_time)
      if t != None:
        traj_set[id] = t
    return traj_set
    
  def getAgentPositionsForTimeExclude(self, query_time, agent_id):
    """
    Get all agent positions for certain time.
    query_time: time where the entries should be queried
    agent_id: id to exclude
    """
    pos_array = np.zeros([0, 2])
    cnt = 0
    for id in self.agent_data.keys():
      if id != agent_id:
        t = self.agent_data[id].getTrajectoryForTime(query_time)
        if t != None and len(t) > 1:
          pos_array = np.append(pos_array, np.expand_dims(t.getPoseAtTime(query_time)[0:2], axis=0), axis=0)
          cnt += 1
    return pos_array
  
  def getAgentVelocitiesForTimeExclude(self, query_time, agent_id):
    """
    Get all agent velocities for certain time.
    query_time: time where the entries should be queried
    agent_id: id to exclude
    """
    vel_array = np.zeros([0, 2])
    cnt = 0
    for id in self.agent_data.keys():
      if id != agent_id:
        t = self.agent_data[id].getTrajectoryForTime(query_time)
        if t != None and len(t) > 1:
          vel_array = np.append(vel_array, np.expand_dims(t.getVelAtTime(query_time)[0:2], axis=0), axis=0)
          cnt += 1
    return vel_array
  
  def plot(self, ax, x_scale = 1, y_scale = 1):
    """
    Plots the trajectories and the static occupancy grid of all agents in this container.
    """
    colormap = pl.get_cmap('rainbow')
    c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=len(self.agent_data.keys()))
    scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
    
    for cnt, id in enumerate(self.agent_data.keys()):
      color_value = scalar_color_map.to_rgba(cnt)
      self.agent_data[id].plot(ax, color=color_value,x_scale = 1, y_scale = 1)
      
    #sup.plot_grid(ax, np.array([self.occupancy_grid.center[0], self.occupancy_grid.center[1]]), self.occupancy_grid.gridmap, self.occupancy_grid.resolution, self.occupancy_grid.map_size)
    sup.plot_grid(ax, np.array([0, 0]),self.occupancy_grid.gridmap, self.occupancy_grid.resolution, self.occupancy_grid.map_size)
    ax.set_xlim([-self.occupancy_grid.center[0], self.occupancy_grid.center[0]])
    ax.set_ylim([-self.occupancy_grid.center[1], self.occupancy_grid.center[1]])
    
    ax.set_aspect('equal')
    
  
src/data_utils/AgentContainer.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import sys
if sys.version_info[0] < 3:
  import Trajectory as traj
else:
  import src.data_utils.Trajectory as traj

class AgentData():
  """
  Contains data for one agent, comprising all trajectories the agent traveled within the given map.
  Different trajectories are separated by their different goals.
  """
  
  def __init__(self, id, radius=0.25):
    self.id = id
    self.goals = []
    self.trajectories = []
    self.radius = radius
    self.cached_trajectory_idx = 0
    self.traj_idx = 0
    self.last_goal = 0
    
  def addSample(self, timestamp, pose, vel, goal):
    """
    Automatically appends the input state to the matching trajectory. 
    If this measurement is part of a new trajectory, one will be created.
    """
    # hack to catch situation of changing motion direction
    #if goal[0] < 0:
    #  goal = self.last_goal
    if goal not in self.goals:

        # adding goal identifier
        self.goals.append(goal)
        self.trajectories.append(traj.Trajectory(goal=np.array([[goal[0], goal[1]]])))
        if len(self.goals) > 1:
          self.traj_idx += 1
    else:
      if len(self.goals)>0:
        if not np.all(goal == self.last_goal):
          self.trajectories.append(traj.Trajectory(goal=np.array([[goal[0], goal[1]]])))
          self.traj_idx += 1

    # append data sample to trajectory (timestamp in [ns], pose and vel as np 1x3 arrays)
    self.trajectories[self.traj_idx].addData(timestamp, pose, vel, goal)
    self.last_goal = goal
    
  
  def getTrajectoryForTime(self, query_time):
    """
    Return the matching trajectory for the query time.
    """
    if self.trajectories[self.cached_trajectory_idx].contains(query_time):
      return self.trajectories[self.cached_trajectory_idx]
    elif query_time >= self.trajectories[0].getMinTime() and query_time <= self.trajectories[-1].getMaxTime():
      # Find trajectory that contains the query time
      for idx, t in enumerate(self.trajectories):
        if t.contains(query_time):
          self.cached_trajectory_idx = idx
          return t
    else:
      return None
  
  def plot(self, ax, color='b', x_scale = 1, y_scale = 1):
    for t in self.trajectories:
      ax.plot(t.pose_vec[:,0]*x_scale, t.pose_vec[:,1]*y_scale, color=color)

src/data_utils/AgentData.py
eliatrevisan/IA-MPPI-LBM
#!/usr/bin/env python
import sys
import numpy as np
import matplotlib.pyplot as pl
import pickle as pkl
import tensorflow as tf

if sys.version_info[0] < 3:
	import yaml

	sys.path.append('../src/data_utils')
	sys.path.append('/opt/ros/kinetic/lib/python2.7/dist-packages')
	import Support as sup
	import OccupancyGrid as occ
	import MemoryBuffer
	import rospy
	from geometry_msgs.msg import Pose
	from geometry_msgs.msg import Twist
	from pedsim_msgs.msg import TrackedPersons
	from sensor_msgs.msg import LaserScan
	from nav_msgs.msg import OccupancyGrid
	from nav_msgs.msg import Path
	import rosbag
else:
	sys.path.append('../src/data_utils')
	import src.data_utils.Support as sup
	import src.data_utils.MemoryBuffer as MemoryBuffer
	import src.data_utils.OccupancyGrid as occ
import random
import os

from multiprocessing.pool import ThreadPool
import math
from time import sleep
from parser import *
import argparse

tfk = tf.keras


class DataHandler():

	def __init__(self, args):
		self.args = args
		self.control_command = np.empty([1, 2])
		self.laser_data = np.empty([1, 720])
		self.grid = np.empty([120, 120])
		self.pedestrian = []
		self.trajectory = []
		self.pedestrian_state = []
		self.robot_command_set = []
		self.robot_state = []
		self.robot_pred_cmd_set = []
		self.data = {}
		self.trajectory_set = []
		self.data_idx = 0
		self.dt = 0.1
		self.batch_size = self.args.batch_size
		self.tbpl = self.args.truncated_backprop_length
		self.cmd_vector_dim = self.args.cmd_vector_dim
		self.prev_horizon = self.args.prev_horizon
		self.prediction_horizon = self.args.prediction_horizon
		self.pred_cmd = np.empty([1, self.prediction_horizon * self.cmd_vector_dim])
		self.min_length_trajectory = self.tbpl + 1 + self.prediction_horizon * 2 + args.prev_horizon
		# Batches
		self.batch_cmd = np.zeros((self.batch_size, self.tbpl, self.cmd_vector_dim))
		self.batch_pred_cmd = np.zeros((self.batch_size, self.tbpl, self.prediction_horizon * self.cmd_vector_dim))
		self.batch_robot_state = np.zeros((self.batch_size, self.tbpl, 3))
		self.batch_grid = np.zeros((self.batch_size, self.tbpl, int(self.args.submap_width / self.args.submap_resolution),
		                            int(self.args.submap_height / self.args.submap_resolution)))
		self.batch_x = np.zeros((self.batch_size, self.tbpl, self.args.input_dim*(self.prev_horizon+1)))
		self.batch_vel = np.zeros((self.batch_size, self.tbpl, self.args.input_state_dim*(self.prev_horizon+1)))
		self.batch_pos = np.zeros((self.batch_size, self.tbpl, self.args.input_state_dim*(self.prev_horizon+1)))
		self.batch_goal = np.zeros((self.batch_size, self.tbpl, 2))
		self.other_agents_info = np.zeros((self.batch_size, self.tbpl, self.args.pedestrian_vector_dim))
		self.other_agents_pos = np.zeros((self.batch_size, self.tbpl, self.args.pedestrian_vector_state_dim))
		self.batch_enc = np.zeros((self.batch_size, self.prev_horizon, self.args.input_dim))
		self.batch_dec = np.zeros((self.batch_size, self.prediction_horizon, self.args.output_dim))
		self.batch_target = np.zeros((self.batch_size, self.tbpl, self.args.output_dim * self.prediction_horizon))
		self.batch_sequences = []
		# indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.sequence_reset = np.ones([self.batch_size])
		self.sequence_idx = np.zeros([self.batch_size])+self.args.prev_horizon

		# Replay buffer for oline training
		self.buffer = MemoryBuffer.MemoryBuffer(self.args.min_buffer_size, self.args.max_buffer_size)

		self.max_traj = self.args.max_trajectories

		self.global_fig = pl.figure("occupancy grid ")
		self.ax_grid = pl.subplot()
		"""
		print("Extracting the occupancy grid ...")
		# Occupancy grid data
		self.occupancy_grid = occ.OccupancyGrid()
		temp = self.args.data_path + '/map.npy'
		# Only used to create a map from png
		map_args = {"file_name": self.args.file_name,
								"resolution": self.args.resolution,
								"map_size": np.array([self.args.map_size_x,self.args.map_size_y]), }
		if not os.path.isfile(self.args.data_path + '/map.npy'):
				print("Creating map from png ...")
				sup.create_map_from_png(data_path=self.args.data_path, **map_args)

		map_data = np.load(os.path.join(self.args.data_path, 'map.npy'), encoding='latin1', allow_pickle=True)
		self.occupancy_grid.gridmap = map_data.item(0)['Map']  # occupancy values of cells
		self.occupancy_grid.resolution = map_data.item(0)['Resolution']  # map resolution in [m / cell]
		self.occupancy_grid.map_size = map_data.item(0)['Size']  # map size in [m]
		self.occupancy_grid.center = self.occupancy_grid.map_size / 2.0

		# FIlter all -1 values present on MAP to zero.
		for ii in range(self.occupancy_grid.gridmap.shape[0]):
				for jj in range(self.occupancy_grid.gridmap.shape[1]):
						if self.occupancy_grid.gridmap[ii, jj] == -1:
								self.occupancy_grid.gridmap[ii, jj] = 0.0
		"""
		if sys.version_info[0] < 3:
			config = yaml.load(open(self.args.topics_config))  # read the yaml file and get the required parameters

			self.robot_topic = config['robot_state_topic']
			self.pedestrian_topic = config['pedestrian_topic']
			self.commmand_topic = config['command_topic']
			self.laser_topic = config['laser_topic']
			self.grid_topic = config['grid_topic']
			self.predict_topic = config['predict_topic']
			self.rate = config['rate']
			rospy.Subscriber(self.robot_topic, Pose, self.RobotStateCB)  # subscribe to various topics
			rospy.Subscriber(self.commmand_topic, Twist, self.ControlCommandCB)
			rospy.Subscriber(self.laser_topic, LaserScan, self.LaserDataCB)
			rospy.Subscriber(self.pedestrian_topic, TrackedPersons, self.TrackedPersonsCB)
			rospy.Subscriber(self.grid_topic, OccupancyGrid, self.GridProcessingCB)
			rospy.Subscriber(self.predict_topic, Path, self.PredictedCommandCB)

	def process_simulation_data_(self):
		# load data from csv for offline training

		# read the required csv files
		if os.path.exists(self.args.data_path):

			if os.path.exists(os.path.join(self.args.data_path, self.args.bag_file)):
				bag = rosbag.Bag(os.path.join(self.args.data_path, self.args.bag_file))

				# variables for storing robot control commands
				linear_vel = np.empty([1])
				angular_vel = np.empty([1])

				# get the control command data
				data_idx = 0
				for topic, msg, t in bag.read_messages(topics=['jackal_velocity_controller/cmd_vel']):
					linear_vel[:] = msg.linear.x
					angular_vel[:] = msg.angular.z
					control_command = np.concatenate([linear_vel, angular_vel])  # [1,2]
					self.robot_command_set.append(control_command.copy())
					data_idx += 1

				# convert from list to numpy array
				self.robot_command_set = np.array(self.robot_command_set)

				for topic, msg, t in bag.read_messages(topics=['/predicted_cmd']):
					self.args.prediction_cmd_horizon = len(msg.poses)
					cmd_pred = np.zeros(self.args.prediction_cmd_horizon * 2)
					for n in range(self.args.prediction_cmd_horizon):
						linear_vel = msg.poses[n].pose.position.x
						angular_vel = msg.poses[n].pose.position.y

						if math.isnan(angular_vel):
							angular_vel = 0

						if math.isnan(linear_vel):
							linear_vel = 0
						cmd_pred[2 * n] = linear_vel
						cmd_pred[2 * n + 1] = angular_vel
					self.robot_pred_cmd_set.append(cmd_pred.copy())

				# convert from list to numpy array
				self.robot_pred_cmd_set = np.asanyarray(self.robot_pred_cmd_set)

				# variables storing robot pose data
				position = np.empty([2])
				orientation = np.empty([1])
				self.robot_state = []

				for topic, msg, t in bag.read_messages(topics=['/robot_state']):
					position[:] = [msg.position.x, msg.position.y]
					orientation[:] = msg.orientation.z
					robot_pose = np.concatenate([position, orientation])  # [1,3]
					self.robot_state.append(robot_pose)

				# convert from list to numpy array
				self.robot_state = np.asanyarray(self.robot_state)

				# check the number of times the robot reaches the goal state, hence providing the number of trajectories.
				self.robot_trajectory_number = 0
				for i in range(1, self.robot_state.shape[0]):
					dist = np.linalg.norm(self.robot_state[i, :2] - self.robot_state[i - 1, :2])
					if dist > self.args.distance_threshold:  # trajectory completion condition
						self.robot_trajectory_number += 1

				print("total trajectories: {}".format(self.robot_trajectory_number))

				for topic, msg, t in bag.read_messages(topics=['/pedsim_visualizer/tracked_persons']):
					# number of pedestrians in recorded data
					ped_number = len(msg.tracks)
					pedestrian = []

					for ped_idx in range(ped_number):
						ped_data = {
							"position": [msg.tracks[ped_idx].pose.pose.position.x, msg.tracks[ped_idx].pose.pose.position.y],
							"velocity": [msg.tracks[ped_idx].twist.twist.linear.x, msg.tracks[ped_idx].twist.twist.linear.y]
							}
						pedestrian.append(ped_data)
					self.pedestrian_state.append(pedestrian)

				# convert from list to numpy array
				# self.pedestrian_state = np.asanyarray(self.pedestrian_state)

				# obtain the training data with the necessary data assigned for each trajectory
				traj = []
				for j in range(np.min(
						[self.robot_state.shape[0], self.robot_pred_cmd_set.shape[0], len(self.pedestrian_state),
						 self.robot_command_set.shape[0]])):  # iterate over trajectory steps

					data = {"robot_state": self.robot_state[j, :],
					        "pedestrian_state": self.pedestrian_state[j],
					        "control_command": self.robot_command_set[j, :],
					        "pred_command": self.robot_pred_cmd_set[j, :]
					        }

					traj.append(data)

					# break the loop when ever a goal state is encountered, signalling the end of data storage
					# for current trajectory
					if j > 1:
						dist = np.linalg.norm(self.robot_state[j, :2] - self.robot_state[j - 1, :2])
						if dist > self.args.distance_threshold:  # trajectory completion condition
							self.trajectory_set.append(traj)
							traj = []

				bag.close()
				random.shuffle(self.trajectory_set)  # shuffle the final dataset
			else:
				print("Could not find bag file...")

	def RobotStateCB(self, msg):

		pos_x = msg.position.x  # get the current state of the robot
		pos_y = msg.position.y
		orientation = msg.orientation.z
		self.robot_state = np.array([[pos_x, pos_y, orientation]])

	def ControlCommandCB(self, msg):

		linear_vel = msg.linear.x  # get the output commands from LMPCC
		angular_vel = msg.angular.z
		self.control_command = np.array([[linear_vel, angular_vel]])

	def LaserDataCB(self, msg):

		laser_data = []  # store the LiDAR data
		laser_data[:] = msg.ranges

		for i in range(len(laser_data)):
			if np.isinf(laser_data[i]):
				laser_data[i] = 30.0
		self.laser_data = np.array([laser_data])

	def TrackedPersonsCB(self, msg):
		self.pedestrian_state = []
		pos = np.zeros((1, 2))
		vel = np.zeros((1, 2))
		for person in range(len(msg.tracks)):  # for each detected pedestrian, get the required data
			pose_x = msg.tracks[person].pose.pose.position.x
			pose_y = msg.tracks[person].pose.pose.position.y
			vel_x = msg.tracks[person].twist.twist.linear.x
			vel_y = msg.tracks[person].twist.twist.linear.y
			pos = np.vstack((pos, [pose_x, pose_y]))
			vel = np.vstack((vel, [vel_x, vel_y]))
			ped_data = {"position": pos[1:],
			            "velocity": vel[1:]
			            }
		self.pedestrian = ped_data

	def GridProcessingCB(self, msg):
		data = []
		data[:] = msg.data
		for i in range(len(data)):  # binarize the static grid data
			if data[i] != 0:
				data[i] = 1
			else:
				data[i] = 0
		data = np.array([data])
		data = np.reshape(data, (msg.info.width, msg.info.height))
		self.grid = self.grid_rotate(data)

	def PredictedCommandCB(self, msg):

		pred_cmd = np.zeros(self.args.prediction_cmd_horizon * 2)
		horizon = len(msg.poses)

		for n in range(horizon):
			linear_vel = msg.poses[n].pose.position.x
			angular_vel = msg.poses[n].pose.position.y

			if math.isnan(angular_vel):
				angular_vel = 0
			if math.isnan(linear_vel):
				linear_vel = 0
			pred_cmd[2 * n] = linear_vel
			pred_cmd[2 * n + 1] = angular_vel

		self.pred_cmd = np.array([pred_cmd])

	def collect_data(self):

		if self.control_command.size != 0 and self.pred_cmd.size != 0:
			self.data = {"robot_state": self.robot_state,
			             "pedestrian_state": self.pedestrian,
			             "control_command": self.control_command,
			             "predicted_cmd": self.pred_cmd
			             }

			self.trajectory.append(self.data)

			if len(self.trajectory) > 1:
				dist = np.linalg.norm(self.trajectory[-1]["robot_state"][:, :2] - self.trajectory[-2]["robot_state"][:, :2])
				if dist > self.args.distance_threshold:
					self.trajectory_set.append(self.trajectory)
					self.trajectory = []

		print("Dataset length: " + str(len(self.trajectory_set)))
		print("Trajectory length" + str(len(self.trajectory)))

	def getSeqToSeqBatch(self):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are filled and can be used or need to be updated.
		new_epoch = False
		t_idx = self.prev_horizon
		for ii in range(0, self.batch_size):
			traj = self.trajectory_set[self.data_idx]
			if self.sequence_idx[ii] + self.prediction_horizon + 1 >= len(traj):  # change trajectory
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set))
				if self.data_idx == 0:
					new_epoch = True
				self.sequence_idx[ii] = self.prev_horizon
				self.sequence_reset[ii] = 1
				t_idx = self.prev_horizon
			else:
				# start_index = 0 + self.prev_horizon
				# end_index = len(traj) - self.prediction_horizon
				# for t_idx in range(start_index,end_index):
				idx = 0
				for i in range(t_idx - self.prev_horizon, t_idx + 1, 1):
					robot_pos = traj[i]["robot_state"][:2]
					robot_vel = traj[i]["control_command"][:]
					ped_pos = traj[i]["pedestrian_state"]["position"]
					ped_vel = traj[i]["pedestrian_state"]["velocity"]
					self.batch_enc[ii, idx, :] = np.concatenate((robot_pos, robot_vel, ped_pos, ped_vel))
					idx += 1
				idx = 0
				for i in range(t_idx, t_idx + self.prediction_horizon, 1):
					self.batch_dec[ii, idx, :] = traj[i - 1]["pedestrian_state"]["velocity"]
					self.batch_dec_target[ii, idx, :] = traj[i]["pedestrian_state"]["velocity"]
					idx += 1
				t_idx += 1
				self.sequence_idx[ii] += 1
				self.sequence_reset[ii] = 0

		return self.batch_enc, self.batch_dec, self.batch_dec_target, new_epoch

	def getBatch(self):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are filled and can be used or need to be updated.
		"""
				Get the next batch of training data.
		"""
		# If the dataset if smaller that the batch_size, repeat trajectories
		traj_id = 0
		while len(self.trajectory_set) < self.batch_size:
			self.trajectory_set.append(self.trajectory_set[traj_id])
			traj_id += 1
		# Update sequences
		# If batch sequences are empty and need to be filled
		if len(self.batch_sequences) == 0:
			for b in range(self.batch_size):
				trajectory = self.trajectory_set[self.data_idx]
				self.data_idx += 1
				self.batch_sequences.append(trajectory)

		# If batch sequences are filled and can be used or need to be updated.
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			# if self.sequence_idx[ii] + self.prediction_horizon + self.tbpl + 1 >= len(traj):  # change trajectory
			if self.sequence_idx[ii] + self.prediction_horizon + self.tbpl >= len(traj):  # change trajectory
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set))
				if self.data_idx == 0:
					new_epoch = True
				self.sequence_idx[ii] = self.args.prev_horizon
				self.sequence_reset[ii] = 1
			else:
				self.sequence_reset[ii] = 0
		# fill the batch
		other_agents_pos = []

		for ii in range(self.batch_size):
			traj = self.batch_sequences[ii]
			start_idx = int(self.sequence_idx[ii])
			for tbp_step in range(self.tbpl):
				for prev_step in range(self.args.prev_horizon, -1, -1):
					robot_pos = traj[start_idx + tbp_step - prev_step]["robot_state"][:2]
					robot_vel = traj[start_idx + tbp_step - prev_step]["control_command"][:]
					ped_pos = traj[start_idx + tbp_step - prev_step]["pedestrian_state"]["position"]
					ped_vel = traj[start_idx + tbp_step - prev_step]["pedestrian_state"]["velocity"]

					self.batch_vel[ii, tbp_step, prev_step * self.args.input_state_dim:(prev_step + 1) * self.args.input_state_dim] = ped_vel
					self.batch_pos[ii, tbp_step, prev_step * self.args.input_state_dim:(prev_step + 1) * self.args.input_state_dim] = ped_pos

					self.batch_x[ii, tbp_step, prev_step * self.args.input_dim:(prev_step + 1) * self.args.input_dim] = np.array(
						[ped_pos[0], ped_pos[1], ped_vel[0], ped_vel[1]])

				# self.batch_goal[ii, tbp_step, :] = traj[start_idx + tbp_step]["pedestrian_goal_position"]
				self.batch_goal[ii, tbp_step, :] = traj[start_idx + tbp_step]["goal_position"]

				heading = math.atan2(ped_vel[1], ped_vel[0])

				# expands vector because the first axis represents an individual agent
				other_poses = np.expand_dims(robot_pos, axis=0)
				self.other_agents_pos[ii, tbp_step, :] = other_poses
				other_pos_local_frame = sup.positions_in_local_frame(ped_pos, heading, other_poses)
				radial_pedestrian_grid = sup.compute_radial_distance_vector(self.args.pedestrian_vector_dim,
				                                                            other_pos_local_frame,
				                                                            max_range=self.args.max_range_ped_grid,
				                                                            min_angle=0, max_angle=2 * np.pi,
				                                                            normalize=True)
				rel_pos = robot_pos - ped_pos
				rel_vel = robot_vel - ped_vel

				if (self.args.relative_info):
					if self.args.pedestrian_vector_dim == 3:
						self.other_agents_info[ii, tbp_step, :] = np.concatenate((np.array([np.linalg.norm(rel_pos)]), rel_vel))
					else:
						self.other_agents_info[ii, tbp_step, :] = np.concatenate((rel_pos, rel_vel))
				else:
					self.other_agents_info[ii, tbp_step, :] = radial_pedestrian_grid

				self.batch_grid[ii, tbp_step, :, :] = np.zeros((int(self.args.submap_width / self.args.submap_resolution),
				                                                int(self.args.submap_height / self.args.submap_resolution)))

				if self.args.agents_on_grid:
					self.add_other_agents_to_grid(self.batch_grid, ii, tbp_step, ped_pos, self.other_agents_pos)

				if self.args.rotated_grid:
					grid = sup.rotate_grid_around_center(grid, -heading * 180 / math.pi)  # rotation in degrees

				# Output values
				if self.args.predict_positions:
					data = "position"
				else:
					data = "velocity"
				for pred_step in range(self.prediction_horizon):
					if self.args.normalize_data:
						self.batch_target[ii, tbp_step, pred_step * self.args.output_dim] = \
							(traj[tbp_step + 1 + pred_step]["pedestrian_state"][data][0] - self.min_x) / (
									self.max_x - self.min_x)
						self.batch_target[ii, tbp_step, pred_step * self.args.output_dim + 1] = \
							(traj[tbp_step + 1 + pred_step]["pedestrian_state"][data][1] - self.min_y) / (
									self.max_y - self.min_y)
					else:
						self.batch_target[ii, tbp_step,
						pred_step * self.args.output_dim:pred_step * self.args.output_dim + self.args.output_dim] = \
							traj[start_idx + tbp_step + 1 + pred_step]["pedestrian_state"][data]

			self.sequence_idx[ii] += self.tbpl

		by = self.batch_target.copy()
		bx = self.batch_x.copy()

		# Rotate velocities into agent frame at each time step (x-axis is aligned with agents direction of motion)
		if self.args.rotated_grid:
			for batch_idx in range(self.batch_x.shape[0]):
				for tbp_step in range(self.batch_x.shape[1]):
					heading = math.atan2(bx[batch_idx, tbp_step, 3], bx[batch_idx, tbp_step, 2])
					rot_mat = np.array([[np.cos(-heading), -np.sin(-heading)], [np.sin(-heading), np.cos(-heading)]])
					bx[batch_idx, tbp_step, 2:] = np.dot(rot_mat, bx[batch_idx, tbp_step, 2:])
					for pred_step in range(by.shape[2] / 2):
						by[batch_idx, tbp_step, 2 * pred_step:2 * pred_step + 2] = np.dot(rot_mat,
						                                                                  by[batch_idx, tbp_step,
						                                                                  2 * pred_step:2 * pred_step + 2])

		return bx, self.batch_vel, self.batch_pos, self.batch_goal, self.batch_grid, self.other_agents_info, by, self.other_agents_pos, new_epoch


	def add_other_agents_to_grid(self, batch_grid, seq_index, t_idx, ped_pos, other_agents):
		sigma = 3  # assumed position error
		rows = batch_grid.shape[2]
		cols = batch_grid.shape[3]
		# Todo expand this for more agents
		other_agents_pos = other_agents[seq_index][t_idx]
		for i in range(1):
			center_x = int(
				(other_agents_pos[0] - ped_pos[0] + rows / 2 * self.args.submap_resolution) / self.args.submap_resolution)
			center_y = int(
				(other_agents_pos[1] - ped_pos[1] + cols / 2 * self.args.submap_resolution) / self.args.submap_resolution)
			for idx in range(max(center_x - 5, 0), min(center_x + 5, rows)):
				for idy in range(max(center_y - 5, 0), min(center_y + 5, cols)):
					z = 1.0 * np.square(idx - center_x) / np.square(sigma) + 1.0 * np.square(
						idy - center_y) / np.square(
						sigma)
					pdf = max(0.0, 100.0 / 2.0 / np.pi / np.square(sigma) * np.exp(-z / 2.0))
					batch_grid[seq_index, t_idx, idx, idy] += pdf
					if batch_grid[seq_index, t_idx, idx, idy] > 1.0:
						batch_grid[seq_index, t_idx, idx, idy] = 1.0


	def getBatchTBP(self):
		"""
		Get the next batch of training data.
		"""
		self.batch_enc = np.zeros((self.batch_size, self.tbpl, self.args.input_dim * self.args.prev_horizon))
		self.batch_dec = np.zeros((self.batch_size, self.tbpl, self.args.output_dim))
		self.batch_dec_target = np.zeros((self.batch_size, self.tbpl, self.args.output_dim * self.args.prediction_horizon))
		# Update sequences
		# If batch sequences are empty and need to be filled
		trajectory = []
		if len(self.batch_sequences) == 0:
			for b in range(0, min(self.batch_size, len(self.trajectory_set) - len(self.trajectory_set) % self.batch_size)):
				trajectory = self.trajectory_set[self.data_idx]
				self.data_idx += 1
				self.batch_sequences.append(trajectory)
		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			if self.sequence_idx[ii] + self.tbpl + self.prediction_horizon + 1 >= len(traj):
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set))
				trajectory = self.trajectory_set[self.data_idx]
				if self.data_idx == 0:
					new_epoch = True
				self.batch_sequences[ii] = trajectory
				self.sequence_idx[ii] = self.prev_horizon
				self.sequence_reset[ii] = 1
			else:
				self.sequence_reset[ii] = 0

		# Fill the batch
		for ii in range(0, min(self.batch_size, len(self.trajectory_set) - len(
				self.trajectory_set) % self.batch_size)):  # min condition is used for small datasets
			traj = self.batch_sequences[ii]
			self.fillBatch(ii, int(self.sequence_idx[ii]), self.batch_enc, self.batch_dec, self.batch_dec_target, traj)
			self.sequence_idx[ii] += self.tbpl

		return self.batch_enc, self.batch_dec, self.batch_dec_target, new_epoch


	def fillBatch(self, batch_idx, start_idx, batch_enc, batch_dec, batch_dec_target, traj):
		"""
		Fill the data batches of batch_idx with data for all truncated backpropagation steps.
		"""
		other_agents_pos = []

		for tbp_step in range(self.tbpl):

			# Input values
			idx = 0
			for i in range(-self.prev_horizon, 0, 1):
				robot_pos = traj[start_idx + tbp_step + i]["robot_state"][0, :2]
				robot_vel = traj[start_idx + tbp_step + i]["control_command"][0, :]
				ped_pos = traj[start_idx + tbp_step + i]["pedestrian_state"]["position"]
				ped_vel = traj[start_idx + tbp_step + i]["pedestrian_state"]["velocity"]
				batch_enc[batch_idx, tbp_step, idx:idx + self.args.input_dim] = np.concatenate((robot_pos,
				                                                                                robot_vel,
				                                                                                ped_pos,
				                                                                                ped_vel))
				idx += self.args.input_dim

			ped_vel = traj[start_idx + tbp_step]["pedestrian_state"]["velocity"]

			radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim,
			                                                            other_pos_local_frame,
			                                                            max_range=self.max_range_ped_grid, min_angle=0,
			                                                            max_angle=2 * np.pi,
			                                                            normalize=True)
			pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			batch_dec[batch_idx, tbp_step, :] = ped_vel
			idx = 0
			for i in range(0, self.prediction_horizon, 1):
				ped_vel = traj[start_idx + tbp_step + i]["pedestrian_state"]["velocity"]
				batch_dec_target[batch_idx, tbp_step, idx:idx + self.args.output_dim] = ped_vel
				idx += self.args.output_dim


	def getTrajectoryAsBatchTBP(self):
		"""
		Get the next batch of training data.
		"""
		id = np.random.randint(0, len(self.trajectory_set))
		traj = self.trajectory_set[id]
		start_index = 0 + self.prev_horizon
		end_index = len(traj) - self.prediction_horizon
		batch_enc_seq = []
		batch_dec_seq = []
		for t_idx in range(start_index, end_index):
			enc_seq = np.zeros((1, 8 * self.prev_horizon))
			dec_seq = np.zeros((1, 2))
			idx = 0
			for i in range(t_idx - self.prev_horizon, t_idx, 1):
				robot_pos = traj[i]["robot_state"][0, :2]
				robot_vel = traj[i]["control_command"][0, :]
				ped_pos = traj[i]["pedestrian_state"]["position"]
				ped_vel = traj[i]["pedestrian_state"]["velocity"]
				input_data = np.concatenate((robot_pos, robot_vel, ped_pos, ped_vel))
				enc_seq[0, idx:idx + 8] = input_data
				idx += 8
			idx = 0
			batch_enc_seq.append(enc_seq)

			dec_seq[0, :] = traj[t_idx]["pedestrian_state"]["velocity"]

			batch_dec_seq.append(dec_seq)

		return batch_enc_seq, batch_dec_seq, traj


	def getTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		traj = self.trajectory_set[trajectory_idx]

		sequence_length = min(max_sequence_length, len(traj) - self.prediction_horizon)
		batch_x = np.zeros([1, sequence_length, self.args.input_dim*(self.prev_horizon+1)])  # data fed for training
		batch_pos = np.zeros([1, sequence_length, self.args.input_state_dim*(self.prev_horizon+1)])  # data fed for training
		batch_vel = np.zeros([1, sequence_length, self.args.input_state_dim*(self.prev_horizon+1)])  # data fed for training
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_target = np.zeros([1, sequence_length, self.args.output_dim * self.prediction_horizon])

		other_agents_info = np.zeros([1, sequence_length, self.args.pedestrian_vector_dim])
		batch_grid = np.zeros((1, sequence_length, int(self.args.submap_width / self.args.submap_resolution),
		                       int(self.args.submap_height / self.args.submap_resolution)))
		other_agents_pos = np.zeros((1, sequence_length, self.args.pedestrian_vector_state_dim))
		# fill the batch
		start_idx = self.args.prev_horizon
		for tbp_step in range(sequence_length):
			for prev_step in range(self.args.prev_horizon, -1, -1):
				robot_pos = traj[start_idx + tbp_step - prev_step]["robot_state"][:2]
				robot_vel = traj[start_idx + tbp_step - prev_step]["control_command"][:]
				ped_pos = traj[start_idx + tbp_step - prev_step]["pedestrian_state"]["position"]
				ped_vel = traj[start_idx + tbp_step - prev_step]["pedestrian_state"]["velocity"]

				batch_vel[0, tbp_step,
				prev_step * self.args.input_state_dim:(prev_step + 1) * self.args.input_state_dim] = ped_vel
				batch_pos[0, tbp_step,
				prev_step * self.args.input_state_dim:(prev_step + 1) * self.args.input_state_dim] = ped_pos

				batch_x[0, tbp_step, prev_step * self.args.input_dim:(prev_step + 1) * self.args.input_dim] = np.array(
					[ped_pos[0], ped_pos[1], ped_vel[0], ped_vel[1]])

			# batch_goal[0, tbp_step, :] = traj[tbp_step]["pedestrian_goal_position"]
			batch_goal[0, tbp_step, :] = traj[tbp_step]["goal_position"]

			heading = math.atan2(robot_vel[1], robot_vel[0])

			# expands vector because the first axis represents an individual agent
			other_poses = robot_pos
			other_agents_pos[0, tbp_step, :] = other_poses
			other_pos_local_frame = sup.positions_in_local_frame(ped_pos, heading, np.expand_dims(other_poses, axis=0))
			radial_pedestrian_grid = sup.compute_radial_distance_vector(self.args.pedestrian_vector_dim,
			                                                            other_pos_local_frame,
			                                                            max_range=self.args.max_range_ped_grid,
			                                                            min_angle=0, max_angle=2 * np.pi,
			                                                            normalize=True)
			rel_pos = robot_pos - ped_pos
			rel_vel = robot_vel - ped_vel

			if (self.args.relative_info):
				if self.args.pedestrian_vector_dim == 3:
					other_agents_info[0, tbp_step, :] = np.concatenate((np.array([np.linalg.norm(rel_pos)]), rel_vel))
				else:
					other_agents_info[0, tbp_step, :] = np.concatenate((rel_pos, rel_vel))
			else:
				other_agents_info[0, tbp_step, :] = radial_pedestrian_grid

			# self.add_other_agents_to_grid(batch_grid,0, tbp_step, ped_pos, other_agents_pos)

			# Output values
			for pred_step in range(self.prediction_horizon):
				if self.args.normalize_data:
					batch_target[0, tbp_step, pred_step * self.args.output_dim] = \
						(traj[tbp_step + 1 + pred_step]["pedestrian_state"]["velocity"][0] - self.min_x) / (self.max_x - self.min_x)
					batch_target[0, tbp_step, pred_step * self.args.output_dim + 1] = \
						(traj[tbp_step + 1 + pred_step]["pedestrian_state"]["velocity"][1] - self.min_y) / (self.max_y - self.min_y)
				else:
					batch_target[0, tbp_step,
					pred_step * self.args.output_dim:pred_step * self.args.output_dim + self.args.output_dim] = \
						traj[tbp_step + 1 + pred_step]["pedestrian_state"]["velocity"]

		return batch_x, batch_vel, batch_pos, batch_goal, batch_grid, other_agents_info, batch_target, other_agents_pos, traj


	def getTrajectoryAsBatchSeqToSeq(self):
		"""
		Get the next batch of training data.
		"""
		id = np.random.randint(0, len(self.trajectory_set))
		traj = self.trajectory_set[id]
		start_index = 0 + self.prev_horizon
		end_index = len(traj) - self.prediction_horizon
		batch_enc_seq = []
		batch_dec_seq = []
		for t_idx in range(start_index, end_index):
			enc_seq = np.zeros((self.prev_horizon, 8))
			dec_seq = np.zeros((self.prediction_horizon, 2))
			idx = 0
			for i in range(t_idx - self.prev_horizon, t_idx, 1):
				robot_pos = traj[i]["robot_state"][:2]
				robot_vel = traj[i]["control_command"][:]
				ped_pos = traj[i]["pedestrian_state"]["position"]
				ped_vel = traj[i]["pedestrian_state"]["velocity"]
				input_data = np.concatenate((robot_pos, robot_vel, ped_pos, ped_vel))
				enc_seq[idx, :] = input_data
				idx += 1
			idx = 0
			batch_enc_seq.append(enc_seq)
			for i in range(t_idx, t_idx + self.prediction_horizon, 1):
				dec_seq[idx, :] = traj[i]["pedestrian_state"]["velocity"]
				idx += 1
			batch_dec_seq.append(dec_seq)

		return batch_enc_seq, batch_dec_seq, traj


	def calc_scale(self, keep_ratio=True):
		self.sx = 1 / (self.max_x - self.min_x)
		self.sy = 1 / (self.max_y - self.min_y)
		if keep_ratio:
			if self.sx > self.sy:
				self.sx = self.sy
			else:
				self.sy = self.sx


	def normalize(self, data, shift=True, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_x * shift) * self.sx
			data_copy[1] = (data[1] - self.min_y * shift) * self.sy
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_x * shift) * self.sx
			data_copy[:, 1] = (data[:, 1] - self.min_y * shift) * self.sy
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_x * shift) * self.sx
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_y * shift) * self.sy
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_x * shift) * self.sx
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_y * shift) * self.sy
		else:
			return False
		return data_copy


	def denormalize(self, data, shift=True, inPlace=False):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx + self.min_x * shift
			data_copy[1] = data[1] / self.sy + self.min_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx + self.min_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy + self.min_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx + self.min_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy + self.min_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx + self.min_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy + self.min_y * shift
		else:
			return False

		return data_copy


	def save_data(self):
		file = open(self.args.data_path, 'wb')
		print("Saving data to: '{}'".format(self.args.data_path))
		# self.dataset.pop(0)
		pkl.dump(self.trajectory_set, file)

		file.close()


	def load_data(self):
		print("Loading data from: '{}'".format(self.args.data_path + self.args.dataset))
		self.file = open(self.args.data_path + self.args.dataset, 'rb')
		if sys.version_info[0] < 3:
			tmp_self = pkl.load(self.file)
		else:
			tmp_self = pkl.load(self.file, encoding='latin1')
		self.trajectory_set = tmp_self
		trajectory_set = []
		# Get normalization constants
		self.pedestrian_dataset = []
		self.robot_dataset = []
		self.max_x = -100
		self.min_x = 100
		self.max_y = -100
		self.min_y = 100
		for traj_id in range(len(self.trajectory_set)):
			traj = self.trajectory_set[traj_id]
			self.robot_vel = np.zeros((len(traj), 2))
			self.pedestrian_vel = np.zeros((len(traj), 2))
			sub_traj = []
			# Sub-sample trajectories to make dt = 0.4
			for t_id in range(0, len(traj), 4):
				self.robot_vel[t_id, :] = traj[t_id]["control_command"]
				self.pedestrian_vel[t_id, :] = traj[t_id]["pedestrian_state"][""]
				self.max_x = max(self.max_x, traj[t_id]["pedestrian_state"]["position"][0])
				self.min_x = min(self.min_x, traj[t_id]["pedestrian_state"]["position"][0])
				self.max_y = max(self.max_y, traj[t_id]["pedestrian_state"]["position"][1])
				self.min_y = min(self.min_y, traj[t_id]["pedestrian_state"]["position"][1])
				sub_traj.append(traj[t_id])
			# Append zero velocities to goal position
			if traj:
				goal_data = traj[-1].copy()
				goal_data["pedestrian_state"]["velocity"] = np.array([0, 0])
				for t_id in range(10):
					sub_traj.append(goal_data)
				if len(sub_traj) > self.min_length_trajectory:
					trajectory_set.append(sub_traj)
				self.robot_dataset.append(self.robot_vel)
				self.pedestrian_dataset.append(self.pedestrian_vel)
		self.robot_trajectory_set = np.asarray(self.robot_dataset)
		self.pedestrian_dataset = np.asarray(self.pedestrian_dataset)
		self.trajectory_set = trajectory_set
		self.calc_scale()
		# random.shuffle(self.trajectory_set)


	def load_map(self):
		try:
			file_path = self.args.data_path + '/map.npy'
			grid_map = np.load(file_path)[()]
			print("Map loaded")
			return grid_map
		except IOError:
			print("ERROR: Cannot load map.")
			return False


	def start_node(self):
		rospy.init_node('OnlineDataHandler', anonymous=True)
		# start node to collect data online
		while not rospy.is_shutdown():  # while the node is running, get the complete data, input and output batches
			self.collect_data()
			if len(self.trajectory_set) > self.args.max_trajectories:
				break
			sleep(1.0 / self.rate)
		self.save_data()


if __name__ == '__main__':
	data_prep = DataHandler()
	data_prep.start_node()

src/data_utils/DataHandler.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import os
import random
import sys
import math
import cv2
import pickle as pkl
import json
from pykalman import KalmanFilter
if sys.version_info[0] < 3:
	import Support as sup
	from Trajectory import *
	import AgentContainer as ped_cont
else:
	import src.data_utils.Support as sup
	from src.data_utils.Trajectory import *
	import src.data_utils.AgentContainer as ped_cont
from copy import deepcopy
import matplotlib.pyplot as pl
import matplotlib.animation as animation
from time import sleep
import random
from scipy.stats import multivariate_normal
#import matplotlib.pyplot as plt

from matplotlib.patches import Ellipse

class DataHandlerLSTM():
	"""
	Data handler for training an LSTM pedestrian prediction model
	"""
	def __init__(self, args):

		self.data_path = args.data_path
		self.batch_size = args.batch_size
		self.tbpl = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.output_sequence_length = args.prediction_horizon
		self.prediction_horizon = args.prediction_horizon
		self.input_dim = args.input_dim
		self.input_state_dim = args.input_state_dim
		self.output_state_dim = args.output_dim
		#self.submap_width = args.submap_width
		#self.submap_height = args.submap_height
		self.submap_size = args.submap_size
		self.submap_width_real = args.submap_span_real
		self.submap_height_real = args.submap_span_real
		#self.submap_resolution = args.submap_resolution
		self.centered_grid = args.centered_grid
		self.rotated_grid = args.rotated_grid
		self.multi_pedestrian = True
		self.pedestrian_radius = args.pedestrian_radius
		self.min_length_trajectory = self.tbpl + 1 + self.output_sequence_length + args.prev_horizon
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.max_range_ped_grid = args.max_range_ped_grid
		self.dt = args.dt
		self.n_mixture = args.n_mixtures
		self.output_pred_state_dim = args.output_pred_state_dim
		self.args = args
		self.real_world_data = args.real_world_data
		# Normalization constants
		self.norm_const_x = 1.0
		self.norm_const_y = 1.0
		self.norm_const_heading = 1.0
		self.norm_const_vx = 1.0
		self.norm_const_vy = 1.0
		self.norm_const_omega = 1.0
		self.min_pos_x = 1000
		self.min_pos_y = 1000
		self.max_pos_x = -1000
		self.max_pos_y = -1000
		self.min_vel_x = 1000
		self.min_vel_y = 1000
		self.max_vel_x = -1000
		self.max_vel_y = -1000
		self.train_set = args.train_set
		self.avg_traj_length = 0

		# Default Distribution Parameters
		self.amplitude_ = 1.0
		self.factor_ = 5.0
		self.cutoff_ = 1.0
		self.covar_ = 0.25
		self.var = multivariate_normal(mean=0, cov=1)

		# Data structure containing all the information about agents
		self.agent_container = ped_cont.AgentContainer()
		self.test_container = ped_cont.AgentContainer()

		self.agent_traj_idx = None  # indicates which trajectory index will be the next one per agent
		self.trajectory_set = []
		self.test_trajectory_set = []
		self.test_ids = []
		self.data_idx = 0
		self.val_data_idx = 0

		# Patch for updating the grid around a pedestrian position
		self.pedestrian_patch = None

		# Training variables
		self.sequence_reset = np.ones([self.batch_size])  # indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.sequence_idx = np.zeros([self.batch_size])+self.args.prev_horizon

		# Batch with info about other agents
		if "future" in self.args.others_info:
			self.batch_vel = np.zeros(
				[self.batch_size, self.tbpl, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			self.batch_vel = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		self.batch_x = np.zeros([self.batch_size, self.tbpl, self.input_dim*(self.prev_horizon+1)])  # data fed for training
		self.batch_grid = np.zeros([self.batch_size, self.tbpl,
																int(self.submap_size),
																int(self.submap_size)])
		if self.args.others_info == "relative":
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim*self.args.n_other_agents])
			self.val_pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.pedestrian_vector_dim*self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		elif self.args.others_info == "ped_grid":
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl,
										int(self.submap_size),
										int(self.submap_size)])
		else:
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])
			self.val_pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])
		self.batch_goal = np.zeros([self.batch_size, self.tbpl, 2])
		self.batch_y = np.zeros([self.batch_size, self.tbpl, self.output_state_dim * self.output_sequence_length])
		self.batch_sequences = []  # Sequence stack (tuples of trajectory + grid) where parts of length self.tbpl are taken from until a sequence is finished. As soon as one is finished, a new sequence is put in the stack.
		self.batch_ids = []
		self.batch_pos = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		self.batch_pos_target = np.zeros([self.batch_size, self.tbpl, 2 * self.output_sequence_length])

		# Validation variables
		self.val_sequence_reset = np.ones([self.batch_size])  # indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.val_sequence_idx = np.zeros([self.batch_size]) + self.args.prev_horizon
		# Batch with current query-agent velocities
		self.val_batch_vel = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		# Batch with query-agent current velocities and positions
		self.val_batch_x = np.zeros([self.batch_size, self.tbpl, self.input_dim * (self.prev_horizon + 1)])  # data fed for training
		# Query agent static environment information
		self.val_batch_grid = np.zeros([self.batch_size, self.tbpl,
									int(self.submap_size),
									int(self.submap_size)])

		# Batch of query-agent goal positions
		self.val_batch_goal = np.zeros([self.batch_size, self.tbpl, 2])
		# Batch with quer-agent next velocities
		self.val_batch_y = np.zeros([self.batch_size, self.tbpl, self.output_state_dim * self.output_sequence_length])
		# Sequence stack (tuples of trajectory + grid) where parts of length self.tbpl are taken from until a sequence is finished. As soon as one is finished, a new sequence is put in the stack.
		self.val_batch_sequences = []
		self.val_batch_ids = []
		# Batch of query-agent current position
		self.val_batch_pos = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		# Batch with quer-agent next positions
		self.val_batch_pos_target = np.zeros([self.batch_size, self.tbpl, 2 * self.output_sequence_length])

	def addAgentTrajectoriesToSet(self,agent_container,trajectory_set, id):
		"""
		Goes through all trajectories of agent and adds them to the member set if they fulfill the criteria.
		For all the time steps within the trajectory it also computes the positions of the other agents at that
		timestep in order to make training more efficient.
		"""
		for traj_idx, traj in enumerate(agent_container.getAgentTrajectories(id)):
			traj_with_collision = False
			if len(traj) > self.min_length_trajectory:
				#if traj.getMinTime() < 100:
				traj.updateInterpolators()
				# Find other agent's trajectories which overlap with each time step
				for time_idx in range(traj.time_vec.shape[0]):
					if self.multi_pedestrian:
						query_time = traj.time_vec[time_idx]
						other_agents_positions = agent_container.getAgentPositionsForTimeExclude(query_time, id)
						other_agents_velocities = agent_container.getAgentVelocitiesForTimeExclude(query_time, id)
			
						# Remove ego agent
						traj.other_agents_positions.append(other_agents_positions)
						traj.other_agents_velocities.append(other_agents_velocities)
				# TODO: MAYBE REMOVE IN THE FUTURE
				#if id != -1:
					#if not traj_with_collision:
				trajectory_set.append((id, traj))
	
	def setOtherAgentData(self, agent_container, trajectory_set, id):
		for traj_idx, traj in enumerate(agent_container.getAgentTrajectories(id)):
			traj_with_collision = False
			if len(traj) > self.min_length_trajectory:
				#traj.updateInterpolators()
				for time_idx in range(traj.time_vec.shape[0]):
					if self.multi_pedestrian:
						query_time = traj.time_vec[time_idx]
						other_agents_positions = agent_container.getAgentPositionsForTimeExclude(query_time, id)
						other_agents_velocities = agent_container.getAgentVelocitiesForTimeExclude(query_time, id)
						# Remove ego agent
						traj.other_agents_positions.append(other_agents_positions)
						traj.other_agents_velocities.append(other_agents_velocities)
				trajectory_set[traj_idx][1].other_agents_positions = traj.other_agents_positions
				trajectory_set[traj_idx][1].other_agents_velocities = traj.other_agents_velocities
				

	def processData(self, **kwargs):
		"""
		Processes the simulation or real-world data, depending on the usage.
		"""
		data_pickle = self.args.data_path + self.args.scenario + "/data" + str(self.args.prediction_horizon) + "_" + str(
			self.args.truncated_backprop_length)+ "_" + str(
			self.args.prev_horizon) + "_roboat.pickle"
		#data_pickle = self.args.data_path + self.args.scenario + "/data" + str(self.args.prediction_horizon) + "_" + str(
		#	self.args.truncated_backprop_length)+ "_" + str(
		#	self.args.prev_horizon) + "experiment3agents-norm.pickle"
		if os.path.isfile(data_pickle):
			self.loadTrajectoryData(data_pickle)
		else:
			if "roboat" in data_pickle:
				print("Processing roboat data.")
				self._process_roboat_data_(**kwargs)
			else:
				print("No roboat data present")

			self.saveTrajectoryData(data_pickle)

	def load_map(self,**kwargs):
		print("Extracting the occupancy grid ...")
		# Occupancy grid data

		with open(os.path.join(self.data_path+self.args.scenario, 'map.json')) as json_file:
			map_data = json.load(json_file)

		grid_map = pl.imread(self.data_path+self.args.scenario + '/map.png')
		
		if self.args.scenario == 'open_crossing' or self.args.scenario == "test_open_crossing":
			grid_map = np.round(grid_map)
		else:
			grid_map = np.round(grid_map[:,:,0])

		#map_data = np.load(os.path.join(self.data_path+self.args.scenario, 'map.npy'), encoding='latin1', allow_pickle=True)
		self.agent_container.occupancy_grid.gridmap = grid_map #map_data.item(0)['Map']  # occupancy values of cells
		self.agent_container.occupancy_grid.resolution = map_data['resolution'] #map_data.item(0)['Resolution']  # map resolution in [m / cell]
		self.agent_container.occupancy_grid.map_size = map_data['map_size'] #map_data.item(0)['Size']  # map size in [m]
		self.agent_container.occupancy_grid.center =  map_data['map_center'] #self.agent_container.occupancy_grid.map_size / 2.0

		#FIlter all -1 values present on MAP to zero.
		for ii in range(self.agent_container.occupancy_grid.gridmap.shape[0]):
			for jj in range(self.agent_container.occupancy_grid.gridmap.shape[1]):
				if self.agent_container.occupancy_grid.gridmap[ii, jj] == -1:
					self.agent_container.occupancy_grid.gridmap[ii, jj] = 0.0

	
	def _process_roboat_data_(self, **kwargs):
		"""
		Import the data from the log file stored in the directory of data_path.
		This method brings all the data into a suitable format for training.
		"""
		self.load_map(**kwargs)
		# Pedestrian data
		# [id, timestep (s), timestep (ns), pos x, pos y, yaw, vel x, vel y, omega, goal x, goal y]
		print("Extracting data from csv file")
		pedestrian_data = np.genfromtxt(os.path.join(self.data_path+self.args.scenario, 'total_log.csv'), delimiter=",")[1:, :]
		print("Data extracted")

		# Iterate through the data and fill the register
		for sample_idx in range(pedestrian_data.shape[0]):
			if pedestrian_data[sample_idx, 0] != -1:
				id = pedestrian_data[sample_idx, 0]
				timestamp = np.round(pedestrian_data[sample_idx, 1],1)# + pedestrian_data[sample_idx, 2] * 1e-9  # time in seconds28.35
				pose = np.zeros([1,3])
				vel = np.zeros([1,3])
				pose[:,0:2] = np.true_divide(pedestrian_data[sample_idx, 3:5], np.array([self.norm_const_x, self.norm_const_y]))
				vel[:,0:2] = np.true_divide(pedestrian_data[sample_idx, 5:7], np.array([self.norm_const_vx, self.norm_const_vy]))
				goal = np.true_divide(pedestrian_data[sample_idx, 7:], np.array([self.norm_const_x, self.norm_const_y]))

				self.agent_container.addDataSample(id, timestamp, pose, vel, goal)

		# Set the initial indices for agent trajectories (which trajectory will be returned when queried)
		self.agent_traj_idx = [0] * self.agent_container.getNumberOfAgents()

		# Subsample trajectories (longer discretization time) from dt=0.1 to dt=0.4
		print("Subsampling data..")
		for id in self.agent_container.getAgentIDs():
			for traj in self.agent_container.getAgentTrajectories(id):
				traj.subsample(int(self.args.dt*10))
		print("Sampling done.")

		for id in self.agent_container.getAgentIDs():
			for traj in self.agent_container.getAgentTrajectories(id):
				for t_id in range(1, traj.pose_vec.shape[0]):
					traj.vel_vec[t_id, :2] = (traj.pose_vec[t_id, :2] - traj.pose_vec[t_id-1, :2]) / self.dt

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id in self.agent_container.getAgentIDs():
			for traj_idx, traj in enumerate(self.agent_container.getAgentTrajectories(id)):
				if len(traj) > self.min_length_trajectory:
					traj.updateInterpolators()

		# Put all the trajectories in the trajectory set and randomize
		for id in self.agent_container.getAgentIDs():
			print("Processing agent {} / {}".format(id, self.agent_container.getNumberOfAgents()))
			# Adds trajectory if bigger than a minimum length and maximum size
			self.addAgentTrajectoriesToSet(self.agent_container,self.trajectory_set,id)
		
		self.compute_min_max_values()

		

	def compute_min_max_values(self):
		for traj_id in range(len(self.trajectory_set)):
			for t_id in range(1, self.trajectory_set[traj_id][1].pose_vec.shape[0]):
				self.trajectory_set[traj_id][1].vel_vec[t_id, :2] = (self.trajectory_set[traj_id][1].pose_vec[t_id, :2] -
																	 self.trajectory_set[traj_id][1].pose_vec[t_id - 1,
																	 :2]) / self.dt
				self.min_pos_x = min(self.min_pos_x,self.trajectory_set[traj_id][1].pose_vec[t_id,0])
				self.min_pos_y = min(self.min_pos_y, self.trajectory_set[traj_id][1].pose_vec[t_id, 1])
				self.max_pos_x = max(self.max_pos_x, self.trajectory_set[traj_id][1].pose_vec[t_id, 0])
				self.max_pos_y = max(self.max_pos_y, self.trajectory_set[traj_id][1].pose_vec[t_id, 1])
				self.min_vel_x = min(self.min_vel_x,self.trajectory_set[traj_id][1].vel_vec[t_id,0])
				self.min_vel_y = min(self.min_vel_y, self.trajectory_set[traj_id][1].vel_vec[t_id, 1])
				self.max_vel_x = max(self.max_vel_x, self.trajectory_set[traj_id][1].vel_vec[t_id, 0])
				self.max_vel_y = max(self.max_vel_y, self.trajectory_set[traj_id][1].vel_vec[t_id, 1])

		self.calc_scale()

		self.args.min_pos_x = self.min_pos_x
		self.args.min_pos_y = self.min_pos_y
		self.args.max_pos_x = self.max_pos_x
		self.args.max_pos_y = self.max_pos_y
		self.args.min_vel_x = self.min_vel_x
		self.args.min_vel_y = self.min_vel_y
		self.args.max_vel_x = self.max_vel_x
		self.args.max_vel_y = self.max_vel_y
		self.args.sx_vel = self.sx_vel
		self.args.sy_vel = self.sy_vel
		self.args.sx_pos = self.sx_pos
		self.args.sy_pos = self.sy_pos

	def unit_test_data_(self,map_args):
			"""
			Generates predefined trajectories to analyse the interaction
			"""
			if not os.path.isfile(self.data_path + self.args.scenario + '/unit_map.npy'):
				print("Creating unit map from png ...")
				sup.create_map_from_png(data_path=self.args.data_path + self.args.scenario, **map_args)

			map_data = np.load(os.path.join(self.data_path + self.args.scenario, 'map.npy'), encoding='latin1',
							   allow_pickle=True)
			self.agent_container.occupancy_grid.gridmap = map_data.item(0)['Map']  # occupancy values of cells
			self.agent_container.occupancy_grid.resolution = map_data.item(0)['Resolution']  # map resolution in [m / cell]
			self.agent_container.occupancy_grid.map_size = map_data.item(0)['Size']  # map size in [m]
			self.agent_container.occupancy_grid.center = self.agent_container.occupancy_grid.map_size / 2.0

			# FIlter all -1 values present on MAP to zero.
			for ii in range(self.agent_container.occupancy_grid.gridmap.shape[0]):
				for jj in range(self.agent_container.occupancy_grid.gridmap.shape[1]):
					if self.agent_container.occupancy_grid.gridmap[ii, jj] == -1:
						self.agent_container.occupancy_grid.gridmap[ii, jj] = 0.0

			# Iterate through the data and fill the register
			pose_agent_1 = np.array([[-15.0,0.0,0.0]])
			pose_agent_2 = np.array([[-13.0, 0.0,0.0]])
			vel = np.zeros([1, 3])
			v0 = 1.0
			goal = np.array([15, 0])
			timestamp = 0
			dt =0.4
			for sample_idx in range(int(320/dt/10)):
				id = 0
				timestamp += dt #self.args.dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 1
				pose_agent_2[:, 0] += v0 * dt

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)

			# Second test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 3
				timestamp += dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

			# Third test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 4
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0
				goal = np.array([15,0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 5
				pose_agent_2[:, 0] -= v0 * dt
				vel[:, 0] = -v0

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, -goal)

			# Forth test
			pose_agent_1 = np.array([[15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 6
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] -= v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = -v0
				goal = np.array([-15,0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

			# Fith test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, -1.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 7
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 8
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, -1])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)
			# Forth test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, 1.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 9
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 10
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, -1])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)
			# Forth test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 11
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 12
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)

			# Put all the trajectories in the trajectory set and randomize
			pool_list = []
			for id in self.test_container.getAgentIDs():
				print("Processing agent {} / {}".format(id, self.test_container.getNumberOfAgents()))
				# Adds trajectory if bigger than a minimum length and maximum size
				self.addAgentTrajectoriesToSet(self.test_container,self.test_trajectory_set,id)

			self.min_pos_x = -1.0
			self.min_pos_y = -1.0
			self.max_pos_x = 15.0
			self.max_pos_y = 15.0
			self.min_vel_x = -1.0
			self.min_vel_y = -1.0
			self.max_vel_x = 1.0
			self.max_vel_y = 1.0

			self.calc_scale()


	def calc_scale(self, keep_ratio=True):
		self.sx_vel = 1 / (self.max_vel_x - self.min_vel_x)
		self.sy_vel = 1 / (self.max_vel_y - self.min_vel_y)
		if keep_ratio:
			if self.sx_vel > self.sy_vel:
				self.sx_vel = self.sy_vel
			else:
				self.sy_vel = self.sx_vel

		self.sx_pos = 1 / (self.max_pos_x - self.min_pos_x)
		self.sy_pos = 1 / (self.max_pos_y - self.min_pos_y)
		if keep_ratio:
			if self.sx_pos > self.sy_pos:
				self.sx_pos = self.sy_pos
			else:
				self.sy_pos = self.sx_pos

	def normalize_vel(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[1] = (data[1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, 1] = (data[:, 1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_vel_y * shift) * self.sy_vel
		else:
			return False
		return data_copy

	def normalize_pos(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[1] = (data[1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, 1] = (data[:, 1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_pos_y * shift) * self.sy_pos
		else:
			return False
		return data_copy

	def denormalize_pos(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx_pos + self.min_pos_x * shift
			data_copy[1] = data[1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy_pos + self.min_pos_y * shift
		else:
			return False

		return data_copy

	def denormalize_vel(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx_vel + self.min_vel_x * shift
			data_copy[1] = data[1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy_vel + self.min_vel_y * shift
		else:
			return False

		return data_copy

	def filter_data(self, data):
		# Kalman filter considering a constant velocity model
		Transition_Matrix = [[1, 0, self.dt, 0], [0, 1, 0, self.dt], [0, 0, 1, 0], [0, 0, 0, 1]]
		Observation_Matrix = [[1, 0, 0, 0], [0, 1, 0, 0]]

		xinit = data[0, 0]
		yinit = data[0, 1]
		vxinit = (data[1, 0] - data[0, 0]) / self.dt
		vyinit = (data[1, 1] - data[0, 1]) / self.dt
		initstate = [xinit, yinit, vxinit, vyinit]
		initcovariance = 1.0e-3 * np.eye(4) #  1.0e-3 * np.eye(4)
		transistionCov = 1.0e-4 * np.eye(4) # 1.0e-4 * np.eye(4)
		observationCov = 2.0e-3 * np.eye(2) # 1.0e-2 * np.eye(2)   -6
		kf = KalmanFilter(transition_matrices=Transition_Matrix,
						  observation_matrices=Observation_Matrix,
						  initial_state_mean=initstate,
						  initial_state_covariance=initcovariance,
						  transition_covariance=transistionCov,
						  observation_covariance=observationCov)

		(filtered_state_means, filtered_state_covariances) = kf.filter(data)

		return filtered_state_means

	def saveTrajectoryData(self, save_path):
		print("Saving data to: '{}'".format(save_path))
		if not os.path.isdir(self.args.data_path + self.args.scenario):
			os.makedirs(self.args.data_path + self.args.scenario)

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id, traj in self.trajectory_set:
			traj.updateInterpolators()

		#if "test" not in self.args.scenario:
		random.shuffle(self.trajectory_set)

		data = {
			"trajectories" : self.trajectory_set,
			"agent_container" : self.agent_container,
			"min_pos_x" : self.min_pos_x,
			"min_pos_y" : self.min_pos_y,
			"max_pos_x" : self.max_pos_x,
			"max_pos_y" : self.max_pos_y,
			"min_vel_x" : self.min_vel_x,
			"min_vel_y" : self.min_vel_y,
			"max_vel_x" : self.max_vel_x,
			"max_vel_y" : self.max_vel_y
		}
		pkl.dump(data, open(save_path, 'wb'),protocol=2)

	def loadTrajectoryData(self, load_path):
		print("Loading data from: '{}'".format(load_path))
		self.file = open(load_path, 'rb')
		if sys.version_info[0] < 3:
			tmp_self = pkl.loads(self.file)#,encoding='latin1')
		else:
			tmp_self = pkl.load(self.file , encoding='latin1')
		self.trajectory_set = tmp_self["trajectories"]
		self.agent_container = tmp_self["agent_container"]

		try:
			self.min_pos_x = tmp_self["min_pos_x"]
			self.min_pos_y = tmp_self["min_pos_y"]
			self.max_pos_x = tmp_self["max_pos_x"]
			self.max_pos_y = tmp_self["max_pos_y"]
			self.min_vel_x = tmp_self["min_vel_x"]
			self.min_vel_y = tmp_self["min_vel_y"]
			self.max_vel_x = tmp_self["max_vel_x"]
			self.max_vel_y = tmp_self["max_vel_y"]
		except:
			print("Delete dataset to compute limits")
			self.max_x = 1
			self.max_y = 1
			self.min_x = -1
			self.min_y = -1

		self.calc_scale()

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id, traj in self.trajectory_set:
			traj.updateInterpolators()

	def getAgentTrajectory(self, agent_id):
		"""
		Return the next agent trajectory in the queue for the agent with id agent_id.
		"""
		trajectory = self.agent_container.agent_data[agent_id].trajectories[self.agent_traj_idx[agent_id]]
		self.agent_traj_idx[agent_id] = (self.agent_traj_idx[agent_id] + 1)  % self.agent_container.getNumberOfTrajectoriesForAgent(agent_id)
		return trajectory

	def getRandomAgentTrajectory(self, agent_id):
		"""
		Return a totally random trajectory for the agent with id agent_id.
		"""
		random_traj_idx = np.random.randint(0, len(self.agent_container.agent_data[agent_id].trajectories))
		return self.agent_container.agent_data[agent_id].trajectories[random_traj_idx]

	def fillBatch(self, agent_id, batch_idx, start_idx, truncated_backprop_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, trajectory,batch_pos_target, centered_grid=False, testing = False):
		"""
		Fill the data batches of batch_idx with data for all truncated backpropagation steps.
		"""
		other_agents_pos = []

		for tbp_step in range(truncated_backprop_length):

			# Input values
			query_time = trajectory.time_vec[start_idx + tbp_step]
			for prev_step in range(self.prev_horizon,-1,-1):
				current_pos = np.array([trajectory.pose_vec[start_idx + tbp_step - prev_step, 0], trajectory.pose_vec[
											start_idx + tbp_step - prev_step, 1]])
				#current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step, 0], trajectory.vel_vec[
				#							start_idx + tbp_step - prev_step - prev_step, 1]])
				current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step - prev_step, 0], trajectory.vel_vec[
											start_idx + tbp_step - prev_step, 1]])

				if self.args.normalize_data:
					self.normalize_pos(current_pos)
					self.normalize_vel(current_vel)

				batch_x[batch_idx, tbp_step, prev_step*self.input_dim:(prev_step+1)*self.input_dim] = np.array([current_pos[0],
																										current_pos[1],
																										current_vel[0],
																										current_vel[1]])
				#batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([np.linalg.norm(current_vel),
				#                                                                                                              np.arctan2(current_vel[1],current_vel[0])])
				batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_vel[0],
																											current_vel[1]])
				batch_pos[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_pos[0],
																											current_pos[1]])
				
			heading = math.atan2(current_vel[1], current_vel[0])
			if centered_grid:
				grid_center = current_pos
				grid = self.agent_container.occupancy_grid.getSubmapByCoords(grid_center[0], grid_center[1], self.submap_width_real, self.submap_height_real)
				grid = cv2.resize(grid, (60,60)) # Resize grid to (60,60) to match autoencoder input size
			if self.rotated_grid:
				grid = sup.rotate_grid_around_center(grid, heading * 180 / math.pi)  # rotation in degrees

			batch_grid[batch_idx, tbp_step, :, :] = grid	

			# Find positions of other pedestrians at the current timestep and order them by dstance to query agent
			other_positions = trajectory.other_agents_positions[start_idx + tbp_step]
			n_other_agents = other_positions.shape[0]
			other_velocities = trajectory.other_agents_velocities[start_idx + tbp_step]

			# Compute distance to other agents and order them
			other_poses_ordered = np.zeros((other_positions.shape[0], 6))
			other_poses_ordered[:, :2] = other_positions
			other_poses_ordered[:, 2:4] = other_velocities
			other_agents_pos.append(other_poses_ordered)

			current_pos = np.array(
				[trajectory.pose_vec[start_idx + tbp_step, 0], trajectory.pose_vec[start_idx + tbp_step, 1]])
			current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step, 0], trajectory.vel_vec[start_idx + tbp_step, 1]])
			heading = math.atan2(current_vel[1], current_vel[0])

			for ag_id in range(n_other_agents):
				other_poses_ordered[ag_id,4] = np.linalg.norm(other_poses_ordered[ag_id,:2] - current_pos)
				# ordered ids
				other_poses_ordered[ag_id,5] = ag_id
			other_poses_ordered= other_poses_ordered[other_poses_ordered[:, 4].argsort()]

			# other_poses_ordered = [ x, y, vx, vy, norm distance, id]

			if self.args.others_info == "relative":
				for ag_id in range(min(n_other_agents,self.args.n_other_agents)):
					rel_pos = np.array([other_poses_ordered[ag_id,0] - current_pos[0],other_poses_ordered[ag_id, 1] - current_pos[1]])*\
								  multivariate_normal.pdf(np.linalg.norm(np.array([other_poses_ordered[ag_id,:2] - current_pos])),
																			  mean=0.0,
																			  cov=self.args.relative_covariance)

					rel_vel = np.array([other_poses_ordered[ag_id,2] - current_vel[0],other_poses_ordered[ag_id, 3] - current_vel[1]])
					
					pedestrian_grid[batch_idx, tbp_step, ag_id*4:ag_id*4+4] = np.concatenate([rel_pos, rel_vel])
					#pedestrian_grid[batch_idx, tbp_step, ag_id, 4] = np.linalg.norm(rel_pos)
					#pedestrian_grid[batch_idx, tbp_step, ag_id, 5] = np.arctan2(rel_pos[1], rel_pos[0])
			

			elif self.args.others_info == "angular_grid":
				other_pos_local_frame = sup.positions_in_local_frame(current_pos, heading, other_positions)
				radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																															max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																															normalize=True)
				pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			# Output values
			for pred_step in range(self.output_sequence_length):
				vx = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 0]
				vy = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 1]
				px = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 0]
				py = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 1]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step] = vx #- current_vel[0]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = vy #- current_vel[1]
				if self.args.normalize_data:
					self.normalize_vel(batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step:self.output_state_dim*pred_step+2])
				batch_pos_target[batch_idx, tbp_step, self.output_state_dim*pred_step] = px - trajectory.pose_vec[start_idx + tbp_step, 0]
				batch_pos_target[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = py - trajectory.pose_vec[start_idx + tbp_step, 1]

			current_pos = trajectory.pose_vec[start_idx + tbp_step, :2]
			#delta_pos = trajectory.pose_vec[start_idx + tbp_step + self.prediction_horizon,:2] - current_pos
			delta_pos = trajectory.goal[0] - current_pos
			#batch_goal[batch_idx, tbp_step, 0] = np.linalg.norm(delta_pos)
			#batch_goal[batch_idx, tbp_step, 1] = np.arctan2(delta_pos[1],
			#                                                delta_pos[0])
			# Used for planning
			batch_goal[batch_idx, tbp_step, 0] = delta_pos[0]/np.linalg.norm(delta_pos) #trajectory.goal[0,0]
			batch_goal[batch_idx, tbp_step, 1] = delta_pos[1]/np.linalg.norm(delta_pos) #trajectory.goal[0,1]

		return other_agents_pos

	def getBatch(self):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled
		trajectory=[]
		if len(self.batch_sequences) == 0:
			for b in range(0,min(self.batch_size,int(len(self.trajectory_set)*self.train_set))):
				id, trajectory = self.trajectory_set[self.data_idx]
				self.data_idx += 1
				self.batch_sequences.append(trajectory)
				self.batch_ids.append(id)
		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			if self.sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj): # Make sure trajs are long enough
				id, trajectory = self.trajectory_set[self.data_idx]
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set)*self.train_set)
				if self.data_idx == 0:
					new_epoch = True
				self.batch_sequences[ii] = trajectory
				self.batch_ids[ii] = id
				self.sequence_idx[ii] = self.args.prev_horizon
				self.sequence_reset[ii] = 1
			else:
				self.sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,min(self.batch_size,len(self.trajectory_set)-len(self.trajectory_set)%self.batch_size)):
			traj = self.batch_sequences[ii]
			agent_id = self.batch_ids[ii]
			other_agents_pos.append(self.fillBatch(agent_id, ii, int(self.sequence_idx[ii]), self.tbpl, self.batch_x, self.batch_vel, self.batch_pos,self.batch_grid, self.pedestrian_grid, self.batch_goal, self.batch_y, traj,self.batch_pos_target, centered_grid=self.centered_grid))
			self.sequence_idx[ii] += self.tbpl

		if self.args.rotated_grid:
			_, self.batch_y = sup.rotate_batch_to_local_frame(self.batch_y,self.batch_x)
			self.batch_x, self.batch_pos_target = sup.rotate_batch_to_local_frame(self.batch_pos_target, self.batch_x)

		return deepcopy(self.batch_x), deepcopy(self.batch_vel), deepcopy(self.batch_pos), deepcopy(self.batch_goal), \
						deepcopy(self.batch_grid), deepcopy(self.pedestrian_grid), deepcopy(self.batch_y), deepcopy(self.batch_pos_target), \
						deepcopy(other_agents_pos), new_epoch

	def getTestBatch(self):
		"""
		Get the next batch for model validation
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled

		self.val_data_idx = max(self.val_data_idx,int(len(self.trajectory_set)*self.train_set))
		trajectory=[]
		if len(self.val_batch_sequences) == 0:
			for b in range(0,self.batch_size):
				id, trajectory = self.trajectory_set[self.val_data_idx]
				self.val_data_idx += 1
				self.val_batch_sequences.append(trajectory)
				self.val_batch_ids.append(id)
				if self.val_data_idx == len(self.trajectory_set):
					self.val_data_idx =int (len(self.trajectory_set) * self.train_set)

		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		for ii, traj in enumerate(self.val_batch_sequences):
			if self.val_sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
				id, trajectory = self.trajectory_set[self.val_data_idx]
				self.val_data_idx = (self.val_data_idx + 1) % int(len(self.trajectory_set)*(1-self.train_set)-1) + int(len(self.trajectory_set)*self.train_set)
				self.val_batch_sequences[ii] = trajectory
				self.val_batch_ids[ii] = id
				self.val_sequence_idx[ii] = self.args.prev_horizon
				self.val_sequence_reset[ii] = 1
			else:
				self.val_sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,min(self.batch_size,len(self.trajectory_set)-len(self.trajectory_set)%self.batch_size)):
			traj = self.val_batch_sequences[ii]
			agent_id = self.val_batch_ids[ii]
			other_agents_pos.append(self.fillBatch(agent_id, ii, int(self.val_sequence_idx[ii]), self.tbpl, self.val_batch_x,
									self.val_batch_vel, self.val_batch_pos,self.val_batch_grid, self.val_pedestrian_grid, self.val_batch_goal,
									self.val_batch_y, traj,self.val_batch_pos_target, centered_grid=self.centered_grid))
			self.val_sequence_idx[ii] += self.tbpl

		if self.args.rotated_grid:
			_, self.val_batch_y = sup.rotate_batch_to_local_frame(self.val_batch_y,self.val_batch_x)
			self.val_batch_x, self.val_batch_pos_target = sup.rotate_batch_to_local_frame(self.val_batch_pos_target, self.val_batch_x)

		# Create dictionary to feed into the model
		dict = {"batch_x": self.val_batch_x,
				"batch_vel": self.val_batch_vel,
				"batch_pos": self.val_batch_pos,
				"batch_goal": self.val_batch_goal,
				"batch_grid": self.val_batch_grid,
				"batch_ped_grid": self.val_pedestrian_grid,
				"batch_y": self.val_batch_y,
				"batch_div": self.val_batch_y,
				"batch_pos_target": self.val_batch_pos_target,
				"state_noise": 0.0,
				"grid_noise": 0.0,
				"ped_noise":0.0,
				"other_agents_pos": other_agents_pos
				}

		return dict

	def getTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000, freeze = False):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		id = self.trajectory_set[trajectory_idx][0]
		traj = self.trajectory_set[trajectory_idx][1]
		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])
		batch_pos_target = np.zeros([1, sequence_length, 2*self.args.prediction_horizon])
		if "future" in self.args.others_info:
			batch_vel = np.zeros(
				[1, sequence_length, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													int(self.submap_size),
													int(self.submap_size)])
		
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * (self.output_sequence_length +3)]) #longer prev horizon, tests
		#batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		if self.args.others_info == "relative":
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		elif self.args.others_info == "ped_grid":
			pedestrian_grid = np.zeros([1, sequence_length,
								   int(self.submap_size),
								   int(self.submap_size)])
		else:
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
		other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj,batch_pos_target, centered_grid=self.centered_grid)

		if freeze:
			pedestrian_grid = pedestrian_grid*0.0

		return batch_x, batch_vel, batch_pos, batch_goal, batch_grid, pedestrian_grid, batch_y, batch_pos_target, other_agents_pos, traj

	def getGroupOfTrajectoriesAsBatch(self, trajectory_idx, max_sequence_length=1000, n_trajs=1):
		"""
		Get a group of trajectories out o the trajectory set in the same format as for the standard training data
		(e.g. for group prediction).
		"""

		batch_x_list = []
		batch_vel_list = []
		batch_pos_list = []
		batch_goal_list = []
		batch_grid_list = []
		pedestrian_grid_list = []
		batch_y_list = []
		batch_pos_final_list = []
		other_agents_pos_list = []
		traj_list = []

		agent_id = self.trajectory_set[trajectory_idx][0]
		traj = self.trajectory_set[trajectory_idx][1]
		grid = self.agent_container.occupancy_grid.getSubmapByCoords(traj.pose_vec[0,0] * self.norm_const_x,
																																 traj.pose_vec[0,1] * self.norm_const_y,
																																 self.submap_size, self.submap_size)
		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])
		batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
		batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													 int(self.submap_size),
													 int(self.submap_size)])
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])

		other_agents_pos_list.append(self.fillBatch(agent_id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel, batch_pos,
										  batch_grid, pedestrian_grid, batch_goal, batch_y, traj, batch_pos_final,
										  centered_grid=self.centered_grid))

		batch_x_list.append(batch_x)
		batch_vel_list.append(batch_vel)
		batch_pos_list.append(batch_pos)
		batch_goal_list.append(batch_goal)
		batch_grid_list.append(batch_grid)
		pedestrian_grid_list.append(pedestrian_grid)
		batch_y_list.append(batch_y)
		batch_pos_final_list.append(batch_pos_final)
		traj_list.append(traj)

		for id in self.agent_container.agent_data.keys():
			if id != agent_id:
				t = self.agent_container.agent_data[id].getTrajectoryForTime(traj.time_vec[0])
				sequence_length = min(max_sequence_length,
									  t.pose_vec.shape[0] - self.output_sequence_length - self.prev_horizon)
				batch_x = np.zeros([1, sequence_length, (self.prev_horizon + 1) * self.input_dim])
				batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
				batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
				batch_grid = np.zeros([1, sequence_length,
									   int(self.submap_size),
									   int(self.submap_size)])
				batch_goal = np.zeros([1, sequence_length, 2])
				batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
				batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
				pedestrian_grid = np.zeros(
					[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
				other_agents_pos_list.append(self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, t,batch_pos_final, centered_grid=self.centered_grid))
				batch_x_list.append(batch_x)
				batch_vel_list.append(batch_vel)
				batch_pos_list.append(batch_pos)
				batch_goal_list.append(batch_goal)
				batch_grid_list.append(batch_grid)
				pedestrian_grid_list.append(pedestrian_grid)
				batch_y_list.append(batch_y)
				batch_pos_final_list.append(batch_pos_final)
				traj_list.append(t)

		return batch_x_list, batch_vel_list, batch_pos_list,batch_goal_list, batch_grid_list, pedestrian_grid_list, batch_y_list, batch_pos_final_list, other_agents_pos_list, traj_list

	def getTestTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000,freeze=False):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		print(trajectory_idx)
		id = self.test_trajectory_set[trajectory_idx][0]
		traj = self.test_trajectory_set[trajectory_idx][1]

		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])  # data fed for training
		batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
		if "future" in self.args.others_info:
			batch_vel = np.zeros(
				[1, sequence_length, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													 int(self.submap_size),
													 int(self.submap_size)])
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		if self.args.others_info == "relative":
			pedestrian_grid = np.zeros([1,sequence_length, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		else:
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
		other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj,batch_pos_final, centered_grid=self.centered_grid,testing=True)

		if freeze:
			pedestrian_grid = pedestrian_grid*0.0

		return batch_x, batch_vel, batch_pos,batch_goal, batch_grid, pedestrian_grid, batch_y, other_agents_pos, traj

	def getAgentDataAtTimestep(self, query_time):
		"""
		Get data from all agents which are active at a certain time.
		"""
		agent_dict = {}  # key is agent id
		for id in self.agent_container.agent_data.keys():
			traj_for_time = self.agent_container.agent_data[id].getTrajectoryForTime(query_time)
			if traj_for_time != None:
				pose = traj_for_time.getPoseAtTime(query_time)
				vel = traj_for_time.getVelAtTime(query_time)
				heading = math.atan2(vel[1], vel[0])
				grid = self.agent_container.occupancy_grid.getSubmapByCoords(pose[0], pose[1], self.submap_size, self.submap_size)
				rotated_grid = sup.rotate_grid_around_center(grid, -heading * 180 / math.pi)
				radial_pedestrian_grid = np.zeros([self.pedestrian_vector_dim])  # Initialize pedestriangrid with 0, will be evaluated later
				goal = traj_for_time.goal
				future_pos_gt = np.zeros([0, 2])
				future_vel_gt = np.zeros([0, 2])
				for future_step in range(1, self.output_sequence_length+1):
					t_future = query_time + future_step * self.dt
					if t_future <= traj_for_time.getMaxTime():
						pose_step = traj_for_time.getPoseAtTime(t_future)
						vel_step = traj_for_time.getVelAtTime(t_future)
						future_pos_gt = np.append(future_pos_gt, np.expand_dims(pose_step[:2], axis=0), axis=0)
						future_vel_gt = np.append(future_vel_gt, np.expand_dims(vel_step[:2], axis=0), axis=0)

				agent_dict[id] = [pose, vel, rotated_grid, radial_pedestrian_grid, goal, future_pos_gt, future_vel_gt]

		# For each pedestrian, find other agents position for this time and compute pedestrian grid
		for id in agent_dict.keys():
			other_agents_pos = np.zeros([0,2])
			for other_id in agent_dict.keys():
				if other_id != id:
					other_agents_pos = np.append(other_agents_pos, np.expand_dims(agent_dict[other_id][0][:2], axis=0), axis=0)
			pos = agent_dict[id][0][:2]
			vel = agent_dict[id][1][:2]
			heading = math.atan2(vel[1], vel[0])
			other_pos_local_frame = sup.positions_in_local_frame(pos, heading, other_agents_pos)
			agent_dict[id][3] = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																													max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																													normalize=True)

		return agent_dict

	def plot_global_scenario(self,batch_grid,batch_x,batch_agent_real_traj,batch_goal,other_agents_pos,batch_agent_predicted_traj,t,seq_index):
		"""
			inputs:
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_x: batch of initial states [x,y, vx, vy] on the global frame
				batch_agent_real_traj: agent ground truth velocity in global frame also know as batch_y
				batch_goal: batch with the trajectory goal positons
				other_agents_pos: other agent position in the global frame
				batch_agent_predicted_traj: agent predicted velocities in global frame
			"""
		global_grid = self.agent_container.occupancy_grid
		fig = pl.figure("Global Trajectory Predictions")
		ax_in = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)

		ax_in.clear()
		#plot scenario grid
		sup.plot_grid(ax_in, np.array([0.0, 0.0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
		ax_in.set_xlim([-global_grid.center[0], global_grid.center[0]])
		ax_in.set_ylim([-global_grid.center[1], global_grid.center[1]])
		#ax_in.set_aspect("equal")

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index, t, 1]*self.norm_const_y
		ax_in.plot(x_init_global_frame, y_init_global_frame, color='g', marker='o', label='Agent initial pos')
		ax_in.plot(batch_goal[seq_index,t, 0]*self.norm_const_x, batch_goal[seq_index,t, 1]*self.norm_const_y,
							 color='purple', marker='o', label='Agent goal')

		#plot predicted trajectory global frame
		vel_pred = np.zeros((self.args.prediction_horizon, 2))
		for mix_idx in range(self.args.n_mixtures):
			# plot predicted trajectory global frame
			for pred_step in range(self.args.prediction_horizon):
				idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
				idy = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures
				mu_x = batch_agent_predicted_traj[seq_index, t, idx]
				mu_y = batch_agent_predicted_traj[seq_index, t, idy]
				vel_pred[pred_step, :] = [mu_x, mu_y]

			traj_pred = sup.path_from_vel(initial_pos=np.array([batch_x[seq_index,t, 0]*self.norm_const_x, batch_x[seq_index,t, 1]*self.norm_const_y]),
																		pred_vel=vel_pred, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
		x_pred = traj_pred[:, 0] #* self.norm_const_x
		y_pred = traj_pred[:, 1] #* self.norm_const_y
		ax_in.plot(x_pred, y_pred, color='r',label='Predicted trajectory')

		#plot real trajectory global frame
		color_value = scalar_color_map.to_rgba(r)
		real_traj = np.zeros((self.args.prediction_horizon, 2))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = batch_agent_real_traj[seq_index, t, idx]
			mu_y = batch_agent_real_traj[seq_index, t, idy]
			real_traj[i, :] = [mu_x, mu_y]
		real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([batch_x[seq_index, t, 0]*self.norm_const_x, batch_x[seq_index, t, 1]*self.norm_const_y]),
												pred_vel=real_traj, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
		x = real_traj_global_frame[:, 0] #* self.norm_const_x
		y = real_traj_global_frame[:, 1] #* self.norm_const_y
		ax_in.plot(x, y, marker='x',color='b',label='Real trajectory')

		#Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			try:
				ax_in.plot(other_pos[0, 0], other_pos[0, 1], color='r', marker='o', label='Other agents')
				for jj in range(1,other_pos.shape[0]):
					ax_in.plot(other_pos[jj, 0], other_pos[jj, 1], color='r', marker='o') #other agent pos positon
			except IndexError:
				print("Oops!  That was no valid number.  Try again...")
		ax_in.legend()
		fig.canvas.draw()
		pl.show(block=False)

	def plot_local_scenario(self,batch_grid,batch_x,batch_agent_real_traj,batch_goal,other_agents_pos,batch_agent_predicted_traj,t,seq_index):
		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth trajectory in global frame
				batch_agent_predicted_traj: agent predicted trajectory in global frame
			"""
		#pl.close('all')
		global_grid = self.agent_container.occupancy_grid

		fig = pl.figure("Local Trajectory Predictions")
		ax_out = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)
		color_value = scalar_color_map.to_rgba(r)
		#Convert initial position to local frame
		#ini_pos_local_frame = sup.rotate_batch_x_to_local_frame(batch_x)
		ax_out.clear()

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index, t, 1]*self.norm_const_y

		# plot initial and goal pose
		ax_out.plot(0, 0, color='g', marker='o', label='Agent initial pos')

		x_lim = [-self.submap_size/2, self.submap_size/2]
		y_lim = [-self.submap_size/2, self.submap_size/2]
		ax_out.set_xlim(x_lim)
		ax_out.set_ylim(y_lim)

		# plot predicted trajectory local frame
		vel_pred = np.zeros((self.args.prediction_horizon, 2))
		for mix_idx in range(self.args.n_mixtures):
			# plot predicted trajectory global frame
			for pred_step in range(self.args.prediction_horizon):
				idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
				idy = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures
				mu_x = batch_agent_predicted_traj[seq_index, t, idx]
				mu_y = batch_agent_predicted_traj[seq_index, t, idy]
				vel_pred[pred_step, :] = [mu_x, mu_y]

			traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),
																		pred_vel=vel_pred, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
			x_pred = traj_pred[:, 0] #* self.norm_const_x
			y_pred = traj_pred[:, 1] #* self.norm_const_y
			ax_out.plot(x_pred, y_pred, color='r',label='Predicted trajectory')

		# plot real trajectory local frame does it make sense? isnt batch_y in the robot frame
		real_traj = np.zeros((self.args.prediction_horizon, 2))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = batch_agent_real_traj[seq_index, t, idx]
			mu_y = batch_agent_real_traj[seq_index, t, idy]
			real_traj[i, :] = [mu_x, mu_y]
		real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]),
																								 pred_vel=real_traj,
																							 dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)

		# Query Agent Local frame plot
		if real_traj[0, 0] > 0.1:
			sup.plot_grid(ax_out, np.array([3.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))
		elif real_traj[0, 0] < -0.1:
			sup.plot_grid(ax_out, np.array([-3.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))
		else:
			sup.plot_grid(ax_out, np.array([0.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))

		x = real_traj_global_frame[:, 0]
		y = real_traj_global_frame[:, 1]
		ax_out.plot(x, y, marker='x', color='b', label='Agent real trajectory')
		#ax_out.plot(batch_x[seq_index,:, 0] * self.norm_const_x-x_init_global_frame*np.ones_like(batch_x[seq_index,:, 0]),
		#            batch_x[seq_index, :, 1] * self.norm_const_y-y_init_global_frame*np.ones_like(batch_x[seq_index,:, 0]), color='k')

		# Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			heading = math.atan2(batch_x[seq_index, t, 3], batch_x[seq_index, t, 2])

			try:
				ax_out.plot(other_pos[0, 0]-x_init_global_frame, other_pos[0, 1]-y_init_global_frame, color='r', marker='o',label='Other agents')
				for jj in range(1,other_pos.shape[0]):
					ax_out.plot(other_pos[jj, 0]-x_init_global_frame, other_pos[jj, 1]-y_init_global_frame, color='r', marker='o')#,markersize=23)  # other agent pos positon
					#ax_out.annotate(str(jj), xy=(other_pos[jj, 0]-x_init_global_frame , other_pos[jj, 1]-y_init_global_frame))
			except IndexError:
				print("Oops!  That was no valid number.  Try again...")

		circle = pl.Circle((0, 0), self.max_range_ped_grid, color='blue',
											 fill=False)
		ax_out.add_artist(circle)
		ax_out.set_aspect("equal")
		ax_out.legend()
		fig.canvas.draw()
				#sleep(0.5)  # Time in seconds.
		pl. show(block=False)

	def normalize_angle_positive(self,angle):
		""" Normalizes the angle to be 0 to 2*pi
				It takes and returns radians. """
		return angle % (2.0 * np.pi)

	def normalize_angle(self,angle):
		""" Normalizes the angle to be -pi to +pi
				It takes and returns radians."""
		a = self.normalize_angle_positive(angle)
		if a > np.pi:
			a -= 2.0 * np.pi
		return a

	def shortest_angular_distance(self,from_angle, to_angle):
		""" Given 2 angles, this returns the shortest angular
				difference.  The inputs and ouputs are of course radians.

				The result would always be -pi <= result <= pi. Adding the result
				to "from" will always get you an equivelent angle to "to".
		"""
		return self.normalize_angle(to_angle - from_angle)

	def get_radius(self,cutoff, A,var):

		return np.sqrt(-2 * var * np.log(cutoff / A))

	def getGaussian(self,x,y,x0,y0,varx,vary,skew):

		dx = x - x0
		dy = y - y0

		h = np.linalg.norm(np.array([dx,dy]))
		angle = np.arctan2(dy, dx)
		mx = np.cos(angle - skew) * h
		my = np.sin(angle - skew) * h
		f1 = np.square(mx) / (2.0 * varx)
		f2 = np.square(my) / (2.0 * vary)

		return self.amplitude_ * np.exp(-(f1 + f2))

	def plot_angular_grid(self, batch_ped_grid,t,seq_index):
		if self.multi_pedestrian:
			fig = pl.figure("Angular grid")
			ax_ped_grid = pl.subplot()
			ax_ped_grid.clear()
			n_sectors = batch_ped_grid.shape[2]
			grid = batch_ped_grid[seq_index,t,:]
			grid = np.expand_dims(grid, axis=1)
			grid_flipped = np.zeros_like(grid)
			grid_flipped[0:int(self.pedestrian_vector_dim/2)] = grid[-int(self.pedestrian_vector_dim/2):]
			grid_flipped[int(self.pedestrian_vector_dim / 2):] = grid[0:int(self.pedestrian_vector_dim / 2)]
			sup.plot_radial_distance_vector(ax_ped_grid, grid_flipped, max_range=1.0, min_angle=0.0,
																			max_angle=2 * np.pi)
			ax_ped_grid.plot(30, 30, color='r', marker='o', markersize=4)
			ax_ped_grid.arrow(0, 0, 1, 0, head_width=0.1,head_length=self.max_range_ped_grid) #agent poiting direction
		# x- and y-range only need to be [-1, 1] since the pedestrian grid is normalized
			ax_ped_grid.set_xlim([-self.max_range_ped_grid-1, self.max_range_ped_grid+1])
			ax_ped_grid.set_ylim([-self.max_range_ped_grid-1, self.max_range_ped_grid+1])

			fig.canvas.draw()
			# sleep(0.5)  # Time in seconds.
			pl.show(block=False)
			sleep(0.5)  # Time in seconds.

	def query_agent_velocity_direction(self, batch_grid, batch_x, batch_agent_real_traj,other_agents_pos,t,seq_index):
		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth velocity in global frame
				batch_agent_predicted_traj: agent predicted trajectory in global frame
			"""
		#pl.close('all')
		global_grid = self.agent_container.occupancy_grid
		fig = pl.figure("Query agent direction")
		ax_in = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)

		number_samples =1
		ax_in.clear()
		#clear plots
		#plot scenario grid
		sup.plot_grid(ax_in, np.array([0.0, 0.0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
		ax_in.set_xlim([-global_grid.center[0], global_grid.center[0]])
		ax_in.set_ylim([-global_grid.center[1], global_grid.center[1]])

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index,t, 1]*self.norm_const_y
		ax_in.plot(x_init_global_frame, y_init_global_frame, color='g', marker='o', label='start')

		#plot real trajectory global frame
		color_value = scalar_color_map.to_rgba(r)

		#Plot velocity direction
		ax_in.arrow(x_init_global_frame, y_init_global_frame, batch_x[seq_index,t, 2]*self.norm_const_vx, batch_x[seq_index,t, 3]*self.norm_const_vy
				, head_width=0.1, head_length=self.max_range_ped_grid+0.5,alpha=0.4,label='Agent velocity direction') # angular grid detection distance
		circle = pl.Circle((x_init_global_frame,y_init_global_frame),self.max_range_ped_grid,color='blue',fill=False,label='Detection range')
		ax_in.add_artist(circle)
		#Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			for jj in range(other_pos.shape[0]):
				ax_in.plot(other_pos[jj, 0], other_pos[jj, 1], color='r', marker='o') #other agent pos positon

		ax_in.set_aspect("equal")
		fig.canvas.draw()
		#sleep(0.5)  # Time in seconds.
		pl. show(block=False)

src/data_utils/DataHandlerLSTM.py
eliatrevisan/IA-MPPI-LBM
#!/usr/bin/env python
import sys
import numpy as np
# from scipy.linalg import hankel

if sys.version_info[0] < 3:
    sys.path.append('../src/data_utils')
    import Support as sup
    from DataHandler import *
else:
    import src.data_utils.Support as sup
    from src.data_utils.DataHandler import *
    
class DataHandlerLSTM_Keras(DataHandler):
    def __init__(self, args):
        super().__init__(args)
        # self.tbpl = self.prev_horizon+1
    
    def getTestTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000):
        """
        Get a trajectory out of the trajectory set in the same format as for the standard training data
        (e.g. for validation purposes).
        """
        print(trajectory_idx)
        id = self.test_trajectory_set[trajectory_idx][0]
        traj = self.test_trajectory_set[trajectory_idx][1]

        sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
        batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])  # data fed for training
        batch_vel = np.zeros([1, sequence_length, (self.prev_horizon + 1) * self.input_state_dim])
        batch_grid = np.zeros([1, sequence_length, int(np.ceil(self.submap_width / self.args.submap_resolution)), int(np.ceil(self.submap_height / self.args.submap_resolution))])
        batch_goal = np.zeros([1, sequence_length, 2])
        batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
        batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
        pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
        other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj, centered_grid=self.centered_grid,testing=True)

        vel_seq = 
        other_agents_info_seq = 

        # return batch_x, batch_vel, batch_pos,batch_goal, batch_grid, pedestrian_grid, batch_y, other_agents_pos, traj

        batch = {
            'x': batch_x,
            'vel': batch_vel,
            'pos': batch_pos,
            'goal': batch_goal,
            'grid': batch_grid,
            'other_agents_info': other_agents_info,
            'y': batch_target,
            'other_agents_pos': other_agents_pos,
            'other_agents_vel': other_agents_vel,
            'traj': traj,
            'vel_seq': batch_vel_seq,
            'other_agents_info_seq': other_agents_info_seq
        }

        return batch
		
    
    def fillBatch(self, agent_id, batch_idx, start_idx, truncated_backprop_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, trajectory, centered_grid=False,testing=False):
		"""
		Fill the data batches of batch_idx with data for all truncated backpropagation steps.
		"""
		other_agents_pos = []

		for tbp_step in range(truncated_backprop_length):

			# Input values
			query_time = trajectory.time_vec[start_idx + tbp_step]
			for prev_step in range(self.prev_horizon,-1,-1):
				current_pos = np.array([trajectory.pose_vec[start_idx + tbp_step - prev_step, 0], trajectory.pose_vec[
					                        start_idx + tbp_step - prev_step, 1]])
				current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step - prev_step, 0] * self.norm_const_vx, trajectory.vel_vec[
					                        start_idx + tbp_step - prev_step, 1] * self.norm_const_vy])

				if self.args.normalize_data:
					self.normalize_pos(current_pos)
					self.normalize_vel(current_vel)

				batch_x[batch_idx, tbp_step, prev_step*self.input_dim:(prev_step+1)*self.input_dim] = np.array([current_pos[0],
																										current_pos[1],
																										current_vel[0],
																										current_vel[1]])
				batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_vel[0],
																											current_vel[1]])

			heading = math.atan2(current_vel[1], current_vel[0])
			if not testing:
				if centered_grid:
					grid_center = current_pos
					grid = self.agent_container.occupancy_grid.getSubmapByCoords(grid_center[0],
																																			 grid_center[1],
																																			 self.submap_width, self.submap_height)
				else:
					grid = self.batch_sequences[batch_idx][1]
					grid_center = trajectory.pose_vec[0, 0:2]

				if self.rotated_grid:
					grid = sup.rotate_grid_around_center(grid, heading * 180 / math.pi)  # rotation in degrees
				else:
					heading = 0
					grid = sup.rotate_grid_around_center(grid, heading * 180 / math.pi)  # rotation in degrees

				batch_grid[batch_idx, tbp_step, :, :] = grid

			batch_goal[batch_idx, tbp_step, :] = trajectory.goal
			# Find positions of other pedestrians at the current timestep

			other_poses = trajectory.other_agents_positions[start_idx + tbp_step]
			n_other_agents = other_poses.shape[0]
			other_velocities = trajectory.other_agents_velocities[start_idx + tbp_step]
			other_agents_pos.append(other_poses)
			other_pos_local_frame = sup.positions_in_local_frame(current_pos, heading, other_poses)
			try:
				if self.args.relative_info:
					radial_pedestrian_grid = np.zeros([self.pedestrian_vector_dim])
					for ag_id in range(n_other_agents):
						radial_pedestrian_grid[ag_id*4:ag_id*4+4] = np.array([other_poses[ag_id,0] - current_pos[0],
						                                                      other_poses[ag_id, 1] - current_pos[1],
						                                                      other_velocities[ag_id, 0] - current_vel[0],
						                                                      other_velocities[ag_id, 1] - current_vel[1]])
					for ag_id in range(n_other_agents,int(self.args.pedestrian_vector_dim/4)):
						radial_pedestrian_grid[ag_id*4:ag_id*4+4] = np.array([100,100,0,0])

				else:
					radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																																max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																																normalize=True)
			except:
				radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
				                                                            max_range=self.max_range_ped_grid, min_angle=0,
				                                                            max_angle=2 * np.pi,
				                                                            normalize=True)

			pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			try:
				rel_pos = np.array([other_poses[0,0] - current_pos[0],other_poses[0, 1] - current_pos[1]])
				rel_vel = np.array([other_velocities[0, 0] - current_vel[0],
							              other_velocities[0, 1] - current_vel[1]])
			except:
				rel_pos = np.array([-100,-100])
				rel_vel = np.array([0,0])

			if (self.args.pedestrian_vector_dim == 4):
				pedestrian_grid[batch_idx, tbp_step, :] = np.concatenate((rel_pos, rel_vel), axis=0)
			elif (self.args.pedestrian_vector_dim == 3):
				pedestrian_grid[batch_idx, tbp_step, :] = np.concatenate((np.linalg.norm(rel_pos), rel_vel), axis=0)
			else:
				pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			# Output values
			for pred_step in range(self.output_sequence_length):
				vx = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 0]
				vy = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 1]
				px = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 0]
				py = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 1]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step] = vx
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = vy
				if self.args.normalize_data:
					self.normalize_vel(batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step:self.output_state_dim*pred_step+2])
				batch_pos[batch_idx, tbp_step, self.output_state_dim*pred_step] = px
				batch_pos[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = py


		return other_agents_pos
src/data_utils/DataHandlerLSTM_Keras.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import os
import random
import sys
import math
import cv2
import pickle as pkl
import json
from pykalman import KalmanFilter
if sys.version_info[0] < 3:
	import Support as sup
	from Trajectory import *
	import AgentContainer as ped_cont
else:
	import src.data_utils.Support as sup
	from src.data_utils.Trajectory import *
	import src.data_utils.AgentContainer as ped_cont
from copy import deepcopy
import matplotlib.pyplot as pl
import matplotlib.animation as animation
from time import sleep
import random
from scipy.stats import multivariate_normal
#import matplotlib.pyplot as plt

from matplotlib.patches import Ellipse

class DataHandlerRoboat():
	"""
	Data handler for training an LSTM pedestrian prediction model
	"""
	def __init__(self, args):

		self.data_path = args.data_path
		self.batch_size = args.batch_size
		self.tbpl = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.output_sequence_length = args.prediction_horizon
		self.prediction_horizon = args.prediction_horizon
		self.input_dim = args.input_dim
		self.input_state_dim = args.input_state_dim
		self.output_state_dim = args.output_dim
		self.submap_size = args.submap_size
		self.submap_width_real = args.submap_span_real
		self.submap_height_real = args.submap_span_real
		self.centered_grid = args.centered_grid
		self.rotated_grid = args.rotated_grid
		self.multi_pedestrian = True
		self.pedestrian_radius = args.pedestrian_radius
		self.min_length_trajectory = self.tbpl + 1 + self.output_sequence_length + args.prev_horizon
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.max_range_ped_grid = args.max_range_ped_grid
		self.dt = args.dt
		self.n_mixture = args.n_mixtures
		self.output_pred_state_dim = args.output_pred_state_dim
		self.args = args
		self.real_world_data = args.real_world_data
		# Normalization constants
		self.norm_const_x = 1.0
		self.norm_const_y = 1.0
		self.norm_const_heading = 1.0
		self.norm_const_vx = 1.0
		self.norm_const_vy = 1.0
		self.norm_const_omega = 1.0
		self.min_pos_x = 1000
		self.min_pos_y = 1000
		self.max_pos_x = -1000
		self.max_pos_y = -1000
		self.min_vel_x = 1000
		self.min_vel_y = 1000
		self.max_vel_x = -1000
		self.max_vel_y = -1000
		self.train_set = args.train_set
		self.avg_traj_length = 0

		# Default Distribution Parameters
		self.amplitude_ = 1.0
		self.factor_ = 5.0
		self.cutoff_ = 1.0
		self.covar_ = 0.25
		self.var = multivariate_normal(mean=0, cov=1)

		# Data structure containing all the information about agents
		self.agent_container = ped_cont.AgentContainer()
		self.test_container = ped_cont.AgentContainer()

		self.agent_traj_idx = None  # indicates which trajectory index will be the next one per agent
		self.trajectory_set = []
		self.test_trajectory_set = []
		self.test_ids = []
		# HARDCODED, 4 datasets
		#self.data_idx = [0,0,0,0]
		#self.val_data_idx = [0,0,0,0]
		self.data_idx = 0
		self.val_data_idx = 0

		self.data_idx_log = []
		self.data_val_log = []
		self.datahandlers = []
		self.traj_set_idx = []
		self.val_traj_idx = []

		# Patch for updating the grid around a pedestrian position
		self.pedestrian_patch = None

		# Training variables
		self.sequence_reset = np.ones([self.batch_size])  # indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.sequence_idx = np.zeros([self.batch_size])+self.args.prev_horizon

		# Batch with info about other agents
		if "future" in self.args.others_info:
			self.batch_vel = np.zeros(
				[self.batch_size, self.tbpl, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			self.batch_vel = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		self.batch_x = np.zeros([self.batch_size, self.tbpl, self.input_dim*(self.prev_horizon+1)])  # data fed for training
		self.batch_grid = np.zeros([self.batch_size, self.tbpl,
																int(self.submap_size),
																int(self.submap_size)])
		if self.args.others_info == "relative":
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim*self.args.n_other_agents])
			self.val_pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.pedestrian_vector_dim*self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			self.pedestrian_grid = np.zeros(
				[self.batch_size, self.tbpl, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		elif self.args.others_info == "ped_grid":
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl,
										int(self.submap_size),
										int(self.submap_size)])
		else:
			self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])
			self.val_pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])
		self.batch_goal = np.zeros([self.batch_size, self.tbpl, 2])
		self.batch_y = np.zeros([self.batch_size, self.tbpl, self.output_state_dim * self.output_sequence_length])
		self.batch_sequences = []  # Sequence stack (tuples of trajectory + grid) where parts of length self.tbpl are taken from until a sequence is finished. As soon as one is finished, a new sequence is put in the stack.
		self.batch_ids = []
		self.batch_pos = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		self.batch_pos_target = np.zeros([self.batch_size, self.tbpl, 2 * self.output_sequence_length])

		# Validation variables
		self.val_sequence_reset = np.ones([self.batch_size])  # indicates whether sequences are reset (hidden state of LSTMs needs to be reset accordingly)
		self.val_sequence_idx = np.zeros([self.batch_size]) + self.args.prev_horizon
		# Batch with current query-agent velocities
		self.val_batch_vel = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		# Batch with query-agent current velocities and positions
		self.val_batch_x = np.zeros([self.batch_size, self.tbpl, self.input_dim * (self.prev_horizon + 1)])  # data fed for training
		# Query agent static environment information
		self.val_batch_grid = np.zeros([self.batch_size, self.tbpl,
									int(self.submap_size),
									int(self.submap_size)])

		# Batch of query-agent goal positions
		self.val_batch_goal = np.zeros([self.batch_size, self.tbpl, 2])
		# Batch with quer-agent next velocities
		self.val_batch_y = np.zeros([self.batch_size, self.tbpl, self.output_state_dim * self.output_sequence_length])
		# Sequence stack (tuples of trajectory + grid) where parts of length self.tbpl are taken from until a sequence is finished. As soon as one is finished, a new sequence is put in the stack.
		self.val_batch_sequences = []
		self.val_batch_ids = []
		# Batch of query-agent current position
		self.val_batch_pos = np.zeros([self.batch_size, self.tbpl, self.input_state_dim * (self.prev_horizon + 1)])
		# Batch with quer-agent next positions
		self.val_batch_pos_target = np.zeros([self.batch_size, self.tbpl, 2 * self.output_sequence_length])

	def addAgentTrajectoriesToSet(self,agent_container,trajectory_set, id):
		"""
		Goes through all trajectories of agent and adds them to the member set if they fulfill the criteria.
		For all the time steps within the trajectory it also computes the positions of the other agents at that
		timestep in order to make training more efficient.
		"""
		for traj_idx, traj in enumerate(agent_container.getAgentTrajectories(id)):
			traj_with_collision = False
			if len(traj) > self.min_length_trajectory:
				#if traj.getMinTime() < 100:
				traj.updateInterpolators()
				# Find other agent's trajectories which overlap with each time step
				for time_idx in range(traj.time_vec.shape[0]):
					if self.multi_pedestrian:
						query_time = traj.time_vec[time_idx]
						other_agents_positions = agent_container.getAgentPositionsForTimeExclude(query_time, id)
						other_agents_velocities = agent_container.getAgentVelocitiesForTimeExclude(query_time, id)
			
						# Remove ego agent
						traj.other_agents_positions.append(other_agents_positions)
						traj.other_agents_velocities.append(other_agents_velocities)
				# TODO: MAYBE REMOVE IN THE FUTURE
				#if id != -1:
					#if not traj_with_collision:
				trajectory_set.append((id, traj))
	
	def setOtherAgentData(self, agent_container, trajectory_set, id):
		for traj_idx, traj in enumerate(agent_container.getAgentTrajectories(id)):
			traj_with_collision = False
			if len(traj) > self.min_length_trajectory:
				#traj.updateInterpolators()
				for time_idx in range(traj.time_vec.shape[0]):
					if self.multi_pedestrian:
						query_time = traj.time_vec[time_idx]
						other_agents_positions = agent_container.getAgentPositionsForTimeExclude(query_time, id)
						other_agents_velocities = agent_container.getAgentVelocitiesForTimeExclude(query_time, id)
						# Remove ego agent
						traj.other_agents_positions.append(other_agents_positions)
						traj.other_agents_velocities.append(other_agents_velocities)
				trajectory_set[traj_idx][1].other_agents_positions = traj.other_agents_positions
				trajectory_set[traj_idx][1].other_agents_velocities = traj.other_agents_velocities
				

	def processData(self, **kwargs):
		"""
		Processes the simulation or real-world data, depending on the usage.
		"""
		data_pickle = self.args.data_path + self.args.scenario + "/data" + str(self.args.prediction_horizon) + "_" + str(
			self.args.truncated_backprop_length)+ "_" + str(
			self.args.prev_horizon) + "_roboat.pickle"
		#data_pickle = self.args.data_path + self.args.scenario + "/data" + str(self.args.prediction_horizon) + "_" + str(
		#	self.args.truncated_backprop_length)+ "_" + str(
		#	self.args.prev_horizon) + "experiment3agents-norm.pickle"
		if os.path.isfile(data_pickle):
			self.loadTrajectoryData(data_pickle)
		else:
			if "roboat" in data_pickle:
				print("Processing roboat data.")
				self._process_roboat_data_(**kwargs)
			else:
				print("No roboat data present")

			self.saveTrajectoryData(data_pickle)

	def load_map(self,**kwargs):
		print("Extracting the occupancy grid ...")
		# Occupancy grid data

		with open(os.path.join(self.data_path+self.args.scenario, 'map.json')) as json_file:
			map_data = json.load(json_file)

		grid_map = pl.imread(self.data_path+self.args.scenario + '/map.png')
		
		if self.args.scenario == 'open_crossing' or self.args.scenario == "test_open_crossing":
			grid_map = np.round(grid_map)
		else:
			grid_map = np.round(grid_map[:,:,0])

		#map_data = np.load(os.path.join(self.data_path+self.args.scenario, 'map.npy'), encoding='latin1', allow_pickle=True)
		self.agent_container.occupancy_grid.gridmap = grid_map #map_data.item(0)['Map']  # occupancy values of cells
		self.agent_container.occupancy_grid.resolution = map_data['resolution'] #map_data.item(0)['Resolution']  # map resolution in [m / cell]
		self.agent_container.occupancy_grid.map_size = map_data['map_size'] #map_data.item(0)['Size']  # map size in [m]
		self.agent_container.occupancy_grid.center =  map_data['map_center'] #self.agent_container.occupancy_grid.map_size / 2.0

		#FIlter all -1 values present on MAP to zero.
		for ii in range(self.agent_container.occupancy_grid.gridmap.shape[0]):
			for jj in range(self.agent_container.occupancy_grid.gridmap.shape[1]):
				if self.agent_container.occupancy_grid.gridmap[ii, jj] == -1:
					self.agent_container.occupancy_grid.gridmap[ii, jj] = 0.0

	
	def _process_roboat_data_(self, **kwargs):
		"""
		Import the data from the log file stored in the directory of data_path.
		This method brings all the data into a suitable format for training.
		"""
		self.load_map(**kwargs)
		# Pedestrian data
		# [id, timestep (s), timestep (ns), pos x, pos y, yaw, vel x, vel y, omega, goal x, goal y]
		print("Extracting data from csv file")
		pedestrian_data = np.genfromtxt(os.path.join(self.data_path+self.args.scenario, 'total_log.csv'), delimiter=",")[1:, :]
		print("Data extracted")

		# Iterate through the data and fill the register
		for sample_idx in range(pedestrian_data.shape[0]):
			if pedestrian_data[sample_idx, 0] != -1:
				id = pedestrian_data[sample_idx, 0]
				timestamp = np.round(pedestrian_data[sample_idx, 1],1)# + pedestrian_data[sample_idx, 2] * 1e-9  # time in seconds28.35
				pose = np.zeros([1,3])
				vel = np.zeros([1,3])
				pose[:,0:2] = np.true_divide(pedestrian_data[sample_idx, 3:5], np.array([self.norm_const_x, self.norm_const_y]))
				vel[:,0:2] = np.true_divide(pedestrian_data[sample_idx, 5:7], np.array([self.norm_const_vx, self.norm_const_vy]))
				goal = np.true_divide(pedestrian_data[sample_idx, 7:], np.array([self.norm_const_x, self.norm_const_y]))
				self.agent_container.addDataSample(id, timestamp, pose, vel, goal)

		# Set the initial indices for agent trajectories (which trajectory will be returned when queried)
		self.agent_traj_idx = [0] * self.agent_container.getNumberOfAgents()

		# Subsample trajectories (longer discretization time) from dt=0.1 to dt=0.4
		print("Subsampling data..")
		for id in self.agent_container.getAgentIDs():
			for traj in self.agent_container.getAgentTrajectories(id):
				traj.subsample(int(self.args.dt*10))
		print("Sampling done.")

		for id in self.agent_container.getAgentIDs():
			for traj in self.agent_container.getAgentTrajectories(id):
				for t_id in range(1, traj.pose_vec.shape[0]):
					traj.vel_vec[t_id, :2] = (traj.pose_vec[t_id, :2] - traj.pose_vec[t_id-1, :2]) / self.dt

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id in self.agent_container.getAgentIDs():
			for traj_idx, traj in enumerate(self.agent_container.getAgentTrajectories(id)):
				if len(traj) > self.min_length_trajectory:
					traj.updateInterpolators()

		# Put all the trajectories in the trajectory set and randomize
		for id in self.agent_container.getAgentIDs():
			print("Processing agent {} / {}".format(id, self.agent_container.getNumberOfAgents()))
			# Adds trajectory if bigger than a minimum length and maximum size
			self.addAgentTrajectoriesToSet(self.agent_container,self.trajectory_set,id)
		
		self.compute_min_max_values()

		

	def compute_min_max_values(self):
		for traj_id in range(len(self.trajectory_set)):
			for t_id in range(1, self.trajectory_set[traj_id][1].pose_vec.shape[0]):
				self.trajectory_set[traj_id][1].vel_vec[t_id, :2] = (self.trajectory_set[traj_id][1].pose_vec[t_id, :2] -
																	 self.trajectory_set[traj_id][1].pose_vec[t_id - 1,
																	 :2]) / self.dt
				self.min_pos_x = min(self.min_pos_x,self.trajectory_set[traj_id][1].pose_vec[t_id,0])
				self.min_pos_y = min(self.min_pos_y, self.trajectory_set[traj_id][1].pose_vec[t_id, 1])
				self.max_pos_x = max(self.max_pos_x, self.trajectory_set[traj_id][1].pose_vec[t_id, 0])
				self.max_pos_y = max(self.max_pos_y, self.trajectory_set[traj_id][1].pose_vec[t_id, 1])
				self.min_vel_x = min(self.min_vel_x,self.trajectory_set[traj_id][1].vel_vec[t_id,0])
				self.min_vel_y = min(self.min_vel_y, self.trajectory_set[traj_id][1].vel_vec[t_id, 1])
				self.max_vel_x = max(self.max_vel_x, self.trajectory_set[traj_id][1].vel_vec[t_id, 0])
				self.max_vel_y = max(self.max_vel_y, self.trajectory_set[traj_id][1].vel_vec[t_id, 1])

		self.calc_scale()

		self.args.min_pos_x = self.min_pos_x
		self.args.min_pos_y = self.min_pos_y
		self.args.max_pos_x = self.max_pos_x
		self.args.max_pos_y = self.max_pos_y
		self.args.min_vel_x = self.min_vel_x
		self.args.min_vel_y = self.min_vel_y
		self.args.max_vel_x = self.max_vel_x
		self.args.max_vel_y = self.max_vel_y
		self.args.sx_vel = self.sx_vel
		self.args.sy_vel = self.sy_vel
		self.args.sx_pos = self.sx_pos
		self.args.sy_pos = self.sy_pos

	def unit_test_data_(self,map_args):
			"""
			Generates predefined trajectories to analyse the interaction
			"""
			if not os.path.isfile(self.data_path + self.args.scenario + '/unit_map.npy'):
				print("Creating unit map from png ...")
				sup.create_map_from_png(data_path=self.args.data_path + self.args.scenario, **map_args)

			map_data = np.load(os.path.join(self.data_path + self.args.scenario, 'map.npy'), encoding='latin1',
							   allow_pickle=True)
			self.agent_container.occupancy_grid.gridmap = map_data.item(0)['Map']  # occupancy values of cells
			self.agent_container.occupancy_grid.resolution = map_data.item(0)['Resolution']  # map resolution in [m / cell]
			self.agent_container.occupancy_grid.map_size = map_data.item(0)['Size']  # map size in [m]
			self.agent_container.occupancy_grid.center = self.agent_container.occupancy_grid.map_size / 2.0

			# FIlter all -1 values present on MAP to zero.
			for ii in range(self.agent_container.occupancy_grid.gridmap.shape[0]):
				for jj in range(self.agent_container.occupancy_grid.gridmap.shape[1]):
					if self.agent_container.occupancy_grid.gridmap[ii, jj] == -1:
						self.agent_container.occupancy_grid.gridmap[ii, jj] = 0.0

			# Iterate through the data and fill the register
			pose_agent_1 = np.array([[-15.0,0.0,0.0]])
			pose_agent_2 = np.array([[-13.0, 0.0,0.0]])
			vel = np.zeros([1, 3])
			v0 = 1.0
			goal = np.array([15, 0])
			timestamp = 0
			dt =0.4
			for sample_idx in range(int(320/dt/10)):
				id = 0
				timestamp += dt #self.args.dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 1
				pose_agent_2[:, 0] += v0 * dt

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)

			# Second test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 3
				timestamp += dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

			# Third test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 4
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 0] = v0
				goal = np.array([15,0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 5
				pose_agent_2[:, 0] -= v0 * dt
				vel[:, 0] = -v0

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, -goal)

			# Forth test
			pose_agent_1 = np.array([[15.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 6
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] -= v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = -v0
				goal = np.array([-15,0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

			# Fith test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, -1.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 7
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 8
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, -1])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)
			# Forth test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, 1.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 9
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 10
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, -1])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)
			# Forth test
			pose_agent_1 = np.array([[-15.0, 0.0, 0.0]])
			pose_agent_2 = np.array([[0.0, 0.0, 0.0]])
			for sample_idx in range(int(320/dt/10)):
				id = 11
				timestamp += (sample_idx + 1) * dt
				pose_agent_1[:, 0] += v0 * dt
				vel[:, 1] = 0
				vel[:, 0] = v0
				goal = np.array([15, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_1, vel, goal)

				id = 12
				vel[:, 0] = 0
				vel[:, 1] = 0
				goal = np.array([0, 0])

				self.test_container.addDataSample(id, timestamp, pose_agent_2, vel, goal)

			# Put all the trajectories in the trajectory set and randomize
			pool_list = []
			for id in self.test_container.getAgentIDs():
				print("Processing agent {} / {}".format(id, self.test_container.getNumberOfAgents()))
				# Adds trajectory if bigger than a minimum length and maximum size
				self.addAgentTrajectoriesToSet(self.test_container,self.test_trajectory_set,id)

			self.min_pos_x = -1.0
			self.min_pos_y = -1.0
			self.max_pos_x = 15.0
			self.max_pos_y = 15.0
			self.min_vel_x = -1.0
			self.min_vel_y = -1.0
			self.max_vel_x = 1.0
			self.max_vel_y = 1.0

			self.calc_scale()


	def calc_scale(self, keep_ratio=True):
		self.sx_vel = 1 / (self.max_vel_x - self.min_vel_x)
		self.sy_vel = 1 / (self.max_vel_y - self.min_vel_y)
		if keep_ratio:
			if self.sx_vel > self.sy_vel:
				self.sx_vel = self.sy_vel
			else:
				self.sy_vel = self.sx_vel

		self.sx_pos = 1 / (self.max_pos_x - self.min_pos_x)
		self.sy_pos = 1 / (self.max_pos_y - self.min_pos_y)
		if keep_ratio:
			if self.sx_pos > self.sy_pos:
				self.sx_pos = self.sy_pos
			else:
				self.sy_pos = self.sx_pos

	def normalize_vel(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[1] = (data[1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, 1] = (data[:, 1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_vel_y * shift) * self.sy_vel
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_vel_x * shift) * self.sx_vel
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_vel_y * shift) * self.sy_vel
		else:
			return False
		return data_copy

	def normalize_pos(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		if data.ndim == 1:
			data_copy[0] = (data[0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[1] = (data[1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 2:
			data_copy[:, 0] = (data[:, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, 1] = (data[:, 1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 3:
			data_copy[:, :, 0] = (data[:, :, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, :, 1] = (data[:, :, 1] - self.min_pos_y * shift) * self.sy_pos
		elif data.ndim == 4:
			data_copy[:, :, :, 0] = (data[:, :, :, 0] - self.min_pos_x * shift) * self.sx_pos
			data_copy[:, :, :, 1] = (data[:, :, :, 1] - self.min_pos_y * shift) * self.sy_pos
		else:
			return False
		return data_copy

	def denormalize_pos(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx_pos + self.min_pos_x * shift
			data_copy[1] = data[1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy_pos + self.min_pos_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx_pos + self.min_pos_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy_pos + self.min_pos_y * shift
		else:
			return False

		return data_copy

	def denormalize_vel(self, data, shift=False, inPlace=True):
		if inPlace:
			data_copy = data
		else:
			data_copy = np.copy(data)

		ndim = data.ndim
		if ndim == 1:
			data_copy[0] = data[0] / self.sx_vel + self.min_vel_x * shift
			data_copy[1] = data[1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 2:
			data_copy[:, 0] = data[:, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, 1] = data[:, 1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 3:
			data_copy[:, :, 0] = data[:, :, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, :, 1] = data[:, :, 1] / self.sy_vel + self.min_vel_y * shift
		elif ndim == 4:
			data_copy[:, :, :, 0] = data[:, :, :, 0] / self.sx_vel + self.min_vel_x * shift
			data_copy[:, :, :, 1] = data[:, :, :, 1] / self.sy_vel + self.min_vel_y * shift
		else:
			return False

		return data_copy

	def filter_data(self, data):
		# Kalman filter considering a constant velocity model
		Transition_Matrix = [[1, 0, self.dt, 0], [0, 1, 0, self.dt], [0, 0, 1, 0], [0, 0, 0, 1]]
		Observation_Matrix = [[1, 0, 0, 0], [0, 1, 0, 0]]

		xinit = data[0, 0]
		yinit = data[0, 1]
		vxinit = (data[1, 0] - data[0, 0]) / self.dt
		vyinit = (data[1, 1] - data[0, 1]) / self.dt
		initstate = [xinit, yinit, vxinit, vyinit]
		initcovariance = 1.0e-3 * np.eye(4) #  1.0e-3 * np.eye(4)
		transistionCov = 1.0e-4 * np.eye(4) # 1.0e-4 * np.eye(4)
		observationCov = 2.0e-3 * np.eye(2) # 1.0e-2 * np.eye(2)   -6
		kf = KalmanFilter(transition_matrices=Transition_Matrix,
						  observation_matrices=Observation_Matrix,
						  initial_state_mean=initstate,
						  initial_state_covariance=initcovariance,
						  transition_covariance=transistionCov,
						  observation_covariance=observationCov)

		(filtered_state_means, filtered_state_covariances) = kf.filter(data)

		return filtered_state_means

	def saveTrajectoryData(self, save_path):
		print("Saving data to: '{}'".format(save_path))
		if not os.path.isdir(self.args.data_path + self.args.scenario):
			os.makedirs(self.args.data_path + self.args.scenario)

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id, traj in self.trajectory_set:
			traj.updateInterpolators()

		#if "test" not in self.args.scenario:
		random.shuffle(self.trajectory_set)

		data = {
			"trajectories" : self.trajectory_set,
			"agent_container" : self.agent_container,
			"min_pos_x" : self.min_pos_x,
			"min_pos_y" : self.min_pos_y,
			"max_pos_x" : self.max_pos_x,
			"max_pos_y" : self.max_pos_y,
			"min_vel_x" : self.min_vel_x,
			"min_vel_y" : self.min_vel_y,
			"max_vel_x" : self.max_vel_x,
			"max_vel_y" : self.max_vel_y
		}
		pkl.dump(data, open(save_path, 'wb'),protocol=2)

	def loadTrajectoryData(self, load_path):
		print("Loading data from: '{}'".format(load_path))
		self.file = open(load_path, 'rb')
		if sys.version_info[0] < 3:
			tmp_self = pkl.loads(self.file)#,encoding='latin1')
		else:
			tmp_self = pkl.load(self.file , encoding='latin1')
		self.trajectory_set = tmp_self["trajectories"]
		self.agent_container = tmp_self["agent_container"]

		try:
			self.min_pos_x = tmp_self["min_pos_x"]
			self.min_pos_y = tmp_self["min_pos_y"]
			self.max_pos_x = tmp_self["max_pos_x"]
			self.max_pos_y = tmp_self["max_pos_y"]
			self.min_vel_x = tmp_self["min_vel_x"]
			self.min_vel_y = tmp_self["min_vel_y"]
			self.max_vel_x = tmp_self["max_vel_x"]
			self.max_vel_y = tmp_self["max_vel_y"]
		except:
			print("Delete dataset to compute limits")
			self.max_x = 1
			self.max_y = 1
			self.min_x = -1
			self.min_y = -1

		self.calc_scale()

		# Reconstruct interpolators since they were not pickled with the rest of the trajectory
		for id, traj in self.trajectory_set:
			traj.updateInterpolators()

	def getAgentTrajectory(self, agent_id):
		"""
		Return the next agent trajectory in the queue for the agent with id agent_id.
		"""
		trajectory = self.agent_container.agent_data[agent_id].trajectories[self.agent_traj_idx[agent_id]]
		self.agent_traj_idx[agent_id] = (self.agent_traj_idx[agent_id] + 1)  % self.agent_container.getNumberOfTrajectoriesForAgent(agent_id)
		return trajectory

	def getRandomAgentTrajectory(self, agent_id):
		"""
		Return a totally random trajectory for the agent with id agent_id.
		"""
		random_traj_idx = np.random.randint(0, len(self.agent_container.agent_data[agent_id].trajectories))
		return self.agent_container.agent_data[agent_id].trajectories[random_traj_idx]

	def fillBatch(self, dh, agent_id, batch_idx, start_idx, truncated_backprop_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, trajectory,batch_pos_target, centered_grid=False, testing = False):
		"""
		Fill the data batches of batch_idx with data for all truncated backpropagation steps.
		"""
		other_agents_pos = []

		for tbp_step in range(truncated_backprop_length):

			# Input values
			query_time = trajectory.time_vec[start_idx + tbp_step]
			for prev_step in range(self.prev_horizon,-1,-1):
				current_pos = np.array([trajectory.pose_vec[start_idx + tbp_step - prev_step, 0], trajectory.pose_vec[
											start_idx + tbp_step - prev_step, 1]])
				#current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step, 0], trajectory.vel_vec[
				#							start_idx + tbp_step - prev_step - prev_step, 1]])
				current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step - prev_step, 0], trajectory.vel_vec[
											start_idx + tbp_step - prev_step, 1]])

				if self.args.normalize_data:
					self.normalize_pos(current_pos)
					self.normalize_vel(current_vel)

				batch_x[batch_idx, tbp_step, prev_step*self.input_dim:(prev_step+1)*self.input_dim] = np.array([current_pos[0],
																										current_pos[1],
																										current_vel[0],
																										current_vel[1]])
				#batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([np.linalg.norm(current_vel),
				#                                                                                                              np.arctan2(current_vel[1],current_vel[0])])
				batch_vel[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_vel[0],
																											current_vel[1]])
				batch_pos[batch_idx, tbp_step, prev_step*self.input_state_dim:(prev_step+1)*self.input_state_dim] = np.array([current_pos[0],
																											current_pos[1]])
				
			heading = math.atan2(current_vel[1], current_vel[0])
			if centered_grid:
				grid_center = current_pos
				grid = dh.agent_container.occupancy_grid.getSubmapByCoords(grid_center[0], grid_center[1], self.submap_width_real, self.submap_height_real)
				grid = cv2.resize(grid, (60,60)) # Resize grid to (60,60) to match autoencoder input size
			if self.rotated_grid:
				grid = sup.rotate_grid_around_center(grid, heading * 180 / math.pi)  # rotation in degrees

			batch_grid[batch_idx, tbp_step, :, :] = grid	

			# Find positions of other pedestrians at the current timestep and order them by dstance to query agent
			other_positions = trajectory.other_agents_positions[start_idx + tbp_step]
			n_other_agents = other_positions.shape[0]
			other_velocities = trajectory.other_agents_velocities[start_idx + tbp_step]

			# Compute distance to other agents and order them
			other_poses_ordered = np.zeros((other_positions.shape[0], 6))
			other_poses_ordered[:, :2] = other_positions
			other_poses_ordered[:, 2:4] = other_velocities
			other_agents_pos.append(other_poses_ordered)

			current_pos = np.array(
				[trajectory.pose_vec[start_idx + tbp_step, 0], trajectory.pose_vec[start_idx + tbp_step, 1]])
			current_vel = np.array([trajectory.vel_vec[start_idx + tbp_step, 0], trajectory.vel_vec[start_idx + tbp_step, 1]])
			heading = math.atan2(current_vel[1], current_vel[0])

			for ag_id in range(n_other_agents):
				other_poses_ordered[ag_id,4] = np.linalg.norm(other_poses_ordered[ag_id,:2] - current_pos)
				# ordered ids
				other_poses_ordered[ag_id,5] = ag_id
			other_poses_ordered= other_poses_ordered[other_poses_ordered[:, 4].argsort()]

			# other_poses_ordered = [ x, y, vx, vy, norm distance, id]

			if self.args.others_info == "relative":
				for ag_id in range(min(n_other_agents,self.args.n_other_agents)):
					rel_pos = np.array([other_poses_ordered[ag_id,0] - current_pos[0],other_poses_ordered[ag_id, 1] - current_pos[1]])*\
								  multivariate_normal.pdf(np.linalg.norm(np.array([other_poses_ordered[ag_id,:2] - current_pos])),
																			  mean=0.0,
																			  cov=self.args.relative_covariance)

					rel_vel = np.array([other_poses_ordered[ag_id,2] - current_vel[0],other_poses_ordered[ag_id, 3] - current_vel[1]])
					
					pedestrian_grid[batch_idx, tbp_step, ag_id*4:ag_id*4+4] = np.concatenate([rel_pos, rel_vel])
					#pedestrian_grid[batch_idx, tbp_step, ag_id, 4] = np.linalg.norm(rel_pos)
					#pedestrian_grid[batch_idx, tbp_step, ag_id, 5] = np.arctan2(rel_pos[1], rel_pos[0])
			

			elif self.args.others_info == "angular_grid":
				other_pos_local_frame = sup.positions_in_local_frame(current_pos, heading, other_positions)
				radial_pedestrian_grid = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																															max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																															normalize=True)
				pedestrian_grid[batch_idx, tbp_step, :] = radial_pedestrian_grid

			# Output values
			for pred_step in range(self.output_sequence_length):
				vx = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 0]
				vy = trajectory.vel_vec[start_idx + tbp_step + 1 + pred_step, 1]
				px = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 0]
				py = trajectory.pose_vec[start_idx + tbp_step + 1 + pred_step, 1]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step] = vx #- current_vel[0]
				batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = vy #- current_vel[1]
				if self.args.normalize_data:
					self.normalize_vel(batch_y[batch_idx, tbp_step, self.output_state_dim*pred_step:self.output_state_dim*pred_step+2])
				batch_pos_target[batch_idx, tbp_step, self.output_state_dim*pred_step] = px - trajectory.pose_vec[start_idx + tbp_step, 0]
				batch_pos_target[batch_idx, tbp_step, self.output_state_dim*pred_step + 1] = py - trajectory.pose_vec[start_idx + tbp_step, 1]

			current_pos = trajectory.pose_vec[start_idx + tbp_step, :2]
			#delta_pos = trajectory.pose_vec[start_idx + tbp_step + self.prediction_horizon,:2] - current_pos
			delta_pos = trajectory.goal[0] - current_pos
			#batch_goal[batch_idx, tbp_step, 0] = np.linalg.norm(delta_pos)
			#batch_goal[batch_idx, tbp_step, 1] = np.arctan2(delta_pos[1],
			#                                                delta_pos[0])
			# Used for planning
			batch_goal[batch_idx, tbp_step, 0] = delta_pos[0]/np.linalg.norm(delta_pos) #trajectory.goal[0,0]
			batch_goal[batch_idx, tbp_step, 1] = delta_pos[1]/np.linalg.norm(delta_pos) #trajectory.goal[0,1]

		return other_agents_pos

	def CombineTrajectorySets(self, dpreps):
		print("Combining data into Roboat DataHandler...")
		n = len(dpreps)
		for dh in range(n):
			for t in range(len(dpreps[dh].trajectory_set)):
				self.trajectory_set.append( dpreps[dh].trajectory_set[t])
				self.datahandlers.append(dh)
		# Shuffle
		c = list(zip(self.trajectory_set, self.datahandlers))
		random.shuffle(c)
		self.trajectory_set, self.datahandlers = zip(*c)
		print("Sucessfully combined datasets!")

	def RemoveStops(self):
		print("Remove stopping trajectories...")
		for t in range(0, len(self.trajectory_set)):
			id, traj = self.trajectory_set[t]
			idx = 1
			while idx < len(traj):	# In other words, check for stopping, but only in the last 10 time steps 
				if np.linalg.norm(traj.pose_vec[-1-idx,:] - traj.pose_vec[-1,:]) < 1.0:
					idx += 1
				else:
					if idx == 1:
						break
					else: 
						traj.pose_vec = traj.pose_vec[:-1-(idx-1)]
						traj.other_agents_positions = traj.other_agents_positions[:-1-(idx-1)]
						traj.other_agents_velocities = traj.other_agents_velocities[:-1-(idx-1)]
						traj.time_vec = traj.time_vec[:-1-(idx-1)]
						traj.vel_vec = traj.vel_vec[:-1-(idx-1)]
						
						#if len(traj.pose_vec) < (self.tbpl + self.output_sequence_length + self.prev_horizon):
							# Remove the trajectory if its too short as a result of deleting time steps
							#print(self.tbpl + self.output_sequence_length + self.prev_horizon, len(traj.pose_vec))
							#self.trajectory_set = self.trajectory_set[:t] + self.trajectory_set[t+1:]

						break
		self.trajectory_set = tuple(item for item in self.trajectory_set if len(item[1].pose_vec) >= (self.tbpl + self.output_sequence_length + self.prev_horizon))
		print("Done!")
	
	def getRoboatBatch(self, dpreps):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled
		trajectory=[]
		if len(self.batch_sequences) == 0:
			for b in range(0,min(self.batch_size,int(len(self.trajectory_set)*self.train_set))):
				id, trajectory = self.trajectory_set[self.data_idx]
				self.traj_set_idx.append(self.data_idx)
				self.data_idx += 1
				self.batch_sequences.append(trajectory)
				self.batch_ids.append(id)
				 
		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			if self.sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj): # Make sure trajs are long enough
				id, trajectory = self.trajectory_set[self.data_idx]
				self.traj_set_idx[ii] = self.data_idx
				self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set)*self.train_set)
				if self.data_idx == 0:
					new_epoch = True
				self.batch_sequences[ii] = trajectory
				self.batch_ids[ii] = id
				self.sequence_idx[ii] = self.args.prev_horizon
				self.sequence_reset[ii] = 1
			else:
				self.sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		dhs = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,min(self.batch_size,len(self.trajectory_set)-len(self.trajectory_set)%self.batch_size)):
			traj = self.batch_sequences[ii]
			agent_id = self.batch_ids[ii]
			dh = self.datahandlers[self.traj_set_idx[ii]]
			other_agents_pos.append(self.fillBatch(dpreps[dh], agent_id, ii, int(self.sequence_idx[ii]), self.tbpl, self.batch_x, self.batch_vel, self.batch_pos,self.batch_grid, self.pedestrian_grid, self.batch_goal, self.batch_y, traj,self.batch_pos_target, centered_grid=self.centered_grid))
			self.sequence_idx[ii] += self.tbpl
			dhs.append(dh)
		
		
		if self.args.rotated_grid:
			_, self.batch_y = sup.rotate_batch_to_local_frame(self.batch_y,self.batch_x)
			self.batch_x, self.batch_pos_target = sup.rotate_batch_to_local_frame(self.batch_pos_target, self.batch_x)

		return deepcopy(self.batch_x), deepcopy(self.batch_vel), deepcopy(self.batch_pos), deepcopy(self.batch_goal), \
						deepcopy(self.batch_grid), deepcopy(self.pedestrian_grid), deepcopy(self.batch_y), deepcopy(self.batch_pos_target), \
						deepcopy(other_agents_pos), new_epoch, [self.batch_ids, dh]

	def getRoboatValidationBatch(self, dpreps):
		"""
		Get the next batch for model validation
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled

		self.val_data_idx = max(self.val_data_idx,int(len(self.trajectory_set)*self.train_set))
		trajectory=[]
		if len(self.val_batch_sequences) == 0:
			for b in range(0,self.batch_size):
				id, trajectory = self.trajectory_set[self.val_data_idx]
				self.val_traj_idx.append(self.val_data_idx)
				self.val_data_idx += 1
				self.val_batch_sequences.append(trajectory)
				self.val_batch_ids.append(id)
				if self.val_data_idx == len(self.trajectory_set):
					self.val_data_idx =int (len(self.trajectory_set) * self.train_set)

		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		for ii, traj in enumerate(self.val_batch_sequences):
			if self.val_sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
				id, trajectory = self.trajectory_set[self.val_data_idx]
				self.val_traj_idx[ii] = self.val_data_idx
				self.val_data_idx = (self.val_data_idx + 1) % int(len(self.trajectory_set)*(1-self.train_set)-1) + int(len(self.trajectory_set)*self.train_set)
				self.val_batch_sequences[ii] = trajectory
				self.val_batch_ids[ii] = id
				self.val_sequence_idx[ii] = self.args.prev_horizon
				self.val_sequence_reset[ii] = 1
			else:
				self.val_sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		dhs = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,min(self.batch_size,len(self.trajectory_set)-len(self.trajectory_set)%self.batch_size)):
			traj = self.val_batch_sequences[ii]
			agent_id = self.val_batch_ids[ii]
			dh = self.datahandlers[self.val_traj_idx[ii]]
			other_agents_pos.append(self.fillBatch(dpreps[dh], agent_id, ii, int(self.val_sequence_idx[ii]), self.tbpl, self.val_batch_x,
									self.val_batch_vel, self.val_batch_pos,self.val_batch_grid, self.val_pedestrian_grid, self.val_batch_goal,
									self.val_batch_y, traj,self.val_batch_pos_target, centered_grid=self.centered_grid))
			self.val_sequence_idx[ii] += self.tbpl
			dhs.append(dh)

		if self.args.rotated_grid:
			_, self.val_batch_y = sup.rotate_batch_to_local_frame(self.val_batch_y,self.val_batch_x)
			self.val_batch_x, self.val_batch_pos_target = sup.rotate_batch_to_local_frame(self.val_batch_pos_target, self.val_batch_x)

		# Create dictionary to feed into the model
		dict = {"batch_x": self.val_batch_x,
				"batch_vel": self.val_batch_vel,
				"batch_pos": self.val_batch_pos,
				"batch_goal": self.val_batch_goal,
				"batch_grid": self.val_batch_grid,
				"batch_ped_grid": self.val_pedestrian_grid,
				"batch_y": self.val_batch_y,
				"batch_div": self.val_batch_y,
				"batch_pos_target": self.val_batch_pos_target,
				"state_noise": 0.0,
				"grid_noise": 0.0,
				"ped_noise":0.0,
				"other_agents_pos": other_agents_pos
				}

		return dict, [self.val_batch_ids, dhs]	


	def getBatch(self, dpreps):
		"""
		Get the next batch of training data.
		"""
		# Update sequences
		# If batch sequences are empty and need to be filled
		trajectory=[]
		
		# Number of datahandlers
		n_dhs = len(dpreps)

		# Determine the (random) order of dataset to get traj from. Re
		dhs = np.random.choice(n_dhs, self.batch_size)
		
		# Select traj from random dataset and fill the trajectory_set
		if len(self.batch_sequences) == 0:
			for b in range(0,self.batch_size):
				# Choose dataset to pick traj from
				id, trajectory = dpreps[dhs[b]].trajectory_set[self.data_idx[dhs[b]]]
				self.data_idx_log.append(self.data_idx[dhs[b]])
				self.data_idx[dhs[b]] += 1
				self.batch_sequences.append(trajectory)
				self.batch_ids.append(id)

		# Right now, new epoch when one of the datasets has reached end of trajectory_set.
		# TODO: make sure new epoch when ALL data is taken

		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		new_epoch = False
		for ii, traj in enumerate(self.batch_sequences):
			if self.sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj): 
				id, trajectory =  dpreps[dhs[ii]].trajectory_set[self.data_idx[dhs[ii]]]
				self.data_idx_log[ii] = self.data_idx[dhs[ii]]
				self.data_idx[dhs[ii]] = (self.data_idx[dhs[ii]] + 1) % int(len(dpreps[dhs[ii]].trajectory_set)*self.train_set)
				if self.data_idx[dhs[ii]] == 0:
					# Make sure they're all set to zero when one dataset was fully passed through network
					for id in range(n_dhs):
						self.data_idx[id] = 0
					new_epoch = True
				self.batch_sequences[ii] = trajectory
				self.batch_ids[ii] = id
				self.sequence_idx[ii] = self.args.prev_horizon
				self.sequence_reset[ii] = 1
				
			else:
				self.sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0,self.batch_size):
			traj = self.batch_sequences[ii]
			agent_id = self.batch_ids[ii]
			other_agents_pos.append(self.fillBatch(dpreps[dhs[ii]], agent_id, ii, int(self.sequence_idx[ii]), self.tbpl, self.batch_x, self.batch_vel, self.batch_pos,self.batch_grid, self.pedestrian_grid, self.batch_goal, self.batch_y, traj,self.batch_pos_target, centered_grid=self.centered_grid))
			self.sequence_idx[ii] += self.tbpl

		if self.args.rotated_grid:
			_, self.batch_y = sup.rotate_batch_to_local_frame(self.batch_y,self.batch_x)
			self.batch_x, self.batch_pos_target = sup.rotate_batch_to_local_frame(self.batch_pos_target, self.batch_x)

		return deepcopy(self.batch_x), deepcopy(self.batch_vel), deepcopy(self.batch_pos), deepcopy(self.batch_goal), \
						deepcopy(self.batch_grid), deepcopy(self.pedestrian_grid), deepcopy(self.batch_y), deepcopy(self.batch_pos_target), \
						deepcopy(other_agents_pos), new_epoch, [dhs, self.data_idx_log]

	def getTestBatch(self, dpreps):
		"""
		Get the next batch for model validation
		"""
		# Number of datahandlers
		n_dhs = len(dpreps)

		# Determine the (random) order of dataset to get traj from. Re
		dhs = np.random.choice(n_dhs, self.batch_size)

		## HARDCODED, amount of trajectories of smallest dataset.
		# Should later make sure ALL data is used
		self.total_len_small = int(len(dpreps[0].trajectory_set)*self.train_set)

		# Update sequences
		# If batch sequences are empty and need to be filled

		#self.val_data_idx = max(self.val_data_idx, self.total_len_small)
		
		#self.data_len = []
		#for i in range(n_dhs):
		#	self.data_len.append(int(len(dpreps[i].trajectory_set)*self.train_set))

		for i in range(n_dhs):
			self.val_data_idx[i] = max( self.val_data_idx[i], int(len(dpreps[i].trajectory_set)*self.train_set) ) 
		
		
		trajectory=[]
		if len(self.val_batch_sequences) == 0:
			for b in range(0,self.batch_size):
				id, trajectory = dpreps[dhs[b]].trajectory_set[self.val_data_idx[dhs[b]]]
				self.data_val_log.append(self.val_data_idx[dhs[b]])
				self.val_data_idx[dhs[b]] += 1
				self.val_batch_sequences.append(trajectory)
				self.val_batch_ids.append(id)
				if self.val_data_idx[dhs[b]] == len(dpreps[dhs[b]].trajectory_set):
					self.val_data_idx[dhs[b]] = int (len(dpreps[dhs[b]].trajectory_set) * self.train_set)

		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		for ii, traj in enumerate(self.val_batch_sequences):
			if self.val_sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
				id, trajectory = dpreps[dhs[ii]].trajectory_set[self.val_data_idx[dhs[ii]]]
				self.data_val_log[ii] = self.val_data_idx[dhs[ii]]
				self.val_data_idx[dhs[ii]] = (self.val_data_idx[dhs[ii]] + 1) % int(len(dpreps[dhs[ii]].trajectory_set)*(1-self.train_set)-1) + int(len(dpreps[dhs[ii]].trajectory_set)*self.train_set) 
				self.val_batch_sequences[ii] = trajectory
				self.val_batch_ids[ii] = id
				self.val_sequence_idx[ii] = self.args.prev_horizon
				self.val_sequence_reset[ii] = 1

			else:
				self.val_sequence_reset[ii] = 0

		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0, self.batch_size):
			traj = self.val_batch_sequences[ii]
			agent_id = self.val_batch_ids[ii]
			other_agents_pos.append(self.fillBatch(dpreps[dhs[ii]], agent_id, ii, int(self.val_sequence_idx[ii]), self.tbpl, self.val_batch_x,
									self.val_batch_vel, self.val_batch_pos,self.val_batch_grid, self.val_pedestrian_grid, self.val_batch_goal,
									self.val_batch_y, traj,self.val_batch_pos_target, centered_grid=self.centered_grid))
			self.val_sequence_idx[ii] += self.tbpl

		if self.args.rotated_grid:
			_, self.val_batch_y = sup.rotate_batch_to_local_frame(self.val_batch_y,self.val_batch_x)
			self.val_batch_x, self.val_batch_pos_target = sup.rotate_batch_to_local_frame(self.val_batch_pos_target, self.val_batch_x)

		# Create dictionary to feed into the model
		dict = {"batch_x": self.val_batch_x,
				"batch_vel": self.val_batch_vel,
				"batch_pos": self.val_batch_pos,
				"batch_goal": self.val_batch_goal,
				"batch_grid": self.val_batch_grid,
				"batch_ped_grid": self.val_pedestrian_grid,
				"batch_y": self.val_batch_y,
				"batch_div": self.val_batch_y,
				"batch_pos_target": self.val_batch_pos_target,
				"state_noise": 0.0,
				"grid_noise": 0.0,
				"ped_noise":0.0,
				"other_agents_pos": other_agents_pos
				}

		return dict, [dhs, self.data_val_log]

	def getPredefinedBatch(self, dpreps, log):

		dset = log[0]
		traj_idx = log[1]

		self.val_batch_sequences = []

		trajectory=[]
		if len(self.val_batch_sequences) == 0:
			for b in range(0,self.batch_size):
				id, trajectory = dpreps[dset[b]].trajectory_set[traj_idx[dset[b]]]
				self.val_batch_sequences.append(trajectory)
				self.val_batch_ids.append(id)
		"""
		# If batch sequences are filled and can be used or need to be updated.
		other_agents_pos = []
		for ii, traj in enumerate(self.val_batch_sequences):
			if self.val_sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
				id, trajectory = dpreps[dhs[ii]].trajectory_set[self.val_data_idx[dhs[ii]]]
				self.val_data_idx[dhs[ii]] = (self.val_data_idx[dhs[ii]] + 1) % int(len(dpreps[dhs[ii]].trajectory_set)*(1-self.train_set)-1) + int(len(dpreps[dhs[ii]].trajectory_set)*self.train_set) 
				self.val_batch_sequences[ii] = trajectory
				self.val_batch_ids[ii] = id
				self.val_sequence_idx[ii] = self.args.prev_horizon
				self.val_sequence_reset[ii] = 1
			else:
				self.val_sequence_reset[ii] = 0
		"""
		# Fill the batch
		other_agents_pos = []
		# Th second argument is needed such that when the dataset is too small the code does not break
		for ii in range(0, self.batch_size):
			traj = self.val_batch_sequences[ii]
			agent_id = self.val_batch_ids[ii]
			other_agents_pos.append(self.fillBatch(dpreps[dset[ii]], agent_id, ii, int(self.val_sequence_idx[ii]), self.tbpl, self.val_batch_x,
									self.val_batch_vel, self.val_batch_pos,self.val_batch_grid, self.val_pedestrian_grid, self.val_batch_goal,
									self.val_batch_y, traj,self.val_batch_pos_target, centered_grid=self.centered_grid))
			self.val_sequence_idx[ii] += self.tbpl

		# Create dictionary
		dict = {"batch_x": self.val_batch_x,
				"batch_vel": self.val_batch_vel,
				"batch_pos": self.val_batch_pos,
				"batch_goal": self.val_batch_goal,
				"batch_grid": self.val_batch_grid,
				"batch_ped_grid": self.val_pedestrian_grid,
				"batch_y": self.val_batch_y,
				"batch_div": self.val_batch_y,
				"batch_pos_target": self.val_batch_pos_target,
				"state_noise": 0.0,
				"grid_noise": 0.0,
				"ped_noise":0.0,
				"other_agents_pos": other_agents_pos
				}

		return dict


	def getTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000, freeze = False):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		id = self.trajectory_set[trajectory_idx][0]
		traj = self.trajectory_set[trajectory_idx][1]
		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])
		batch_pos_target = np.zeros([1, sequence_length, 2*self.args.prediction_horizon])
		if "future" in self.args.others_info:
			batch_vel = np.zeros(
				[1, sequence_length, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													int(self.submap_size),
													int(self.submap_size)])
		
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * (self.output_sequence_length +3)]) #longer prev horizon, tests
		#batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		if self.args.others_info == "relative":
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		elif self.args.others_info == "ped_grid":
			pedestrian_grid = np.zeros([1, sequence_length,
								   int(self.submap_size),
								   int(self.submap_size)])
		else:
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
		other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj,batch_pos_target, centered_grid=self.centered_grid)

		if freeze:
			pedestrian_grid = pedestrian_grid*0.0

		return batch_x, batch_vel, batch_pos, batch_goal, batch_grid, pedestrian_grid, batch_y, batch_pos_target, other_agents_pos, traj

	def getGroupOfTrajectoriesAsBatch(self, trajectory_idx, max_sequence_length=1000, n_trajs=1):
		"""
		Get a group of trajectories out o the trajectory set in the same format as for the standard training data
		(e.g. for group prediction).
		"""

		batch_x_list = []
		batch_vel_list = []
		batch_pos_list = []
		batch_goal_list = []
		batch_grid_list = []
		pedestrian_grid_list = []
		batch_y_list = []
		batch_pos_final_list = []
		other_agents_pos_list = []
		traj_list = []

		agent_id = self.trajectory_set[trajectory_idx][0]
		traj = self.trajectory_set[trajectory_idx][1]
		grid = self.agent_container.occupancy_grid.getSubmapByCoords(traj.pose_vec[0,0] * self.norm_const_x,
																																 traj.pose_vec[0,1] * self.norm_const_y,
																																 self.submap_size, self.submap_size)
		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])
		batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
		batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													 int(self.submap_size),
													 int(self.submap_size)])
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])

		other_agents_pos_list.append(self.fillBatch(agent_id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel, batch_pos,
										  batch_grid, pedestrian_grid, batch_goal, batch_y, traj, batch_pos_final,
										  centered_grid=self.centered_grid))

		batch_x_list.append(batch_x)
		batch_vel_list.append(batch_vel)
		batch_pos_list.append(batch_pos)
		batch_goal_list.append(batch_goal)
		batch_grid_list.append(batch_grid)
		pedestrian_grid_list.append(pedestrian_grid)
		batch_y_list.append(batch_y)
		batch_pos_final_list.append(batch_pos_final)
		traj_list.append(traj)

		for id in self.agent_container.agent_data.keys():
			if id != agent_id:
				t = self.agent_container.agent_data[id].getTrajectoryForTime(traj.time_vec[0])
				sequence_length = min(max_sequence_length,
									  t.pose_vec.shape[0] - self.output_sequence_length - self.prev_horizon)
				batch_x = np.zeros([1, sequence_length, (self.prev_horizon + 1) * self.input_dim])
				batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
				batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
				batch_grid = np.zeros([1, sequence_length,
									   int(self.submap_size),
									   int(self.submap_size)])
				batch_goal = np.zeros([1, sequence_length, 2])
				batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
				batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
				pedestrian_grid = np.zeros(
					[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
				other_agents_pos_list.append(self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, t,batch_pos_final, centered_grid=self.centered_grid))
				batch_x_list.append(batch_x)
				batch_vel_list.append(batch_vel)
				batch_pos_list.append(batch_pos)
				batch_goal_list.append(batch_goal)
				batch_grid_list.append(batch_grid)
				pedestrian_grid_list.append(pedestrian_grid)
				batch_y_list.append(batch_y)
				batch_pos_final_list.append(batch_pos_final)
				traj_list.append(t)

		return batch_x_list, batch_vel_list, batch_pos_list,batch_goal_list, batch_grid_list, pedestrian_grid_list, batch_y_list, batch_pos_final_list, other_agents_pos_list, traj_list

	def getTestTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000,freeze=False):
		"""
		Get a trajectory out of the trajectory set in the same format as for the standard training data
		(e.g. for validation purposes).
		"""
		print(trajectory_idx)
		id = self.test_trajectory_set[trajectory_idx][0]
		traj = self.test_trajectory_set[trajectory_idx][1]

		sequence_length = min(max_sequence_length, traj.pose_vec.shape[0] - self.output_sequence_length-self.prev_horizon)
		batch_x = np.zeros([1, sequence_length, (self.prev_horizon+1)*self.input_dim])  # data fed for training
		batch_pos_final = np.zeros([1, sequence_length, self.args.output_dim])
		if "future" in self.args.others_info:
			batch_vel = np.zeros(
				[1, sequence_length, self.input_state_dim * self.prediction_horizon])  # data fed for training
		else:
			batch_vel = np.zeros([1, sequence_length, self.input_state_dim * (self.prev_horizon + 1)])
		batch_grid = np.zeros([1, sequence_length,
													 int(self.submap_size),
													 int(self.submap_size)])
		batch_goal = np.zeros([1, sequence_length, 2])
		batch_y = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		batch_pos = np.zeros([1, sequence_length, self.output_state_dim * self.output_sequence_length])
		if self.args.others_info == "relative":
			pedestrian_grid = np.zeros([1,sequence_length, self.pedestrian_vector_dim*self.args.n_other_agents])
		elif "sequence" in self.args.others_info:
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim * self.prediction_horizon])
		elif self.args.others_info == "prev_sequence":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.pedestrian_vector_dim*(self.prev_horizon + 1)])
		elif self.args.others_info == "sequence2":
			pedestrian_grid = np.zeros(
				[1, sequence_length, self.args.n_other_agents, self.prediction_horizon,self.pedestrian_vector_dim])
		else:
			pedestrian_grid = np.zeros([1, sequence_length, self.pedestrian_vector_dim])
		other_agents_pos = self.fillBatch(id, 0, self.prev_horizon, sequence_length, batch_x, batch_vel,batch_pos,batch_grid, pedestrian_grid, batch_goal, batch_y, traj,batch_pos_final, centered_grid=self.centered_grid,testing=True)

		if freeze:
			pedestrian_grid = pedestrian_grid*0.0

		return batch_x, batch_vel, batch_pos,batch_goal, batch_grid, pedestrian_grid, batch_y, other_agents_pos, traj

	def getAgentDataAtTimestep(self, query_time):
		"""
		Get data from all agents which are active at a certain time.
		"""
		agent_dict = {}  # key is agent id
		for id in self.agent_container.agent_data.keys():
			traj_for_time = self.agent_container.agent_data[id].getTrajectoryForTime(query_time)
			if traj_for_time != None:
				pose = traj_for_time.getPoseAtTime(query_time)
				vel = traj_for_time.getVelAtTime(query_time)
				heading = math.atan2(vel[1], vel[0])
				grid = self.agent_container.occupancy_grid.getSubmapByCoords(pose[0], pose[1], self.submap_size, self.submap_size)
				rotated_grid = sup.rotate_grid_around_center(grid, -heading * 180 / math.pi)
				radial_pedestrian_grid = np.zeros([self.pedestrian_vector_dim])  # Initialize pedestriangrid with 0, will be evaluated later
				goal = traj_for_time.goal
				future_pos_gt = np.zeros([0, 2])
				future_vel_gt = np.zeros([0, 2])
				for future_step in range(1, self.output_sequence_length+1):
					t_future = query_time + future_step * self.dt
					if t_future <= traj_for_time.getMaxTime():
						pose_step = traj_for_time.getPoseAtTime(t_future)
						vel_step = traj_for_time.getVelAtTime(t_future)
						future_pos_gt = np.append(future_pos_gt, np.expand_dims(pose_step[:2], axis=0), axis=0)
						future_vel_gt = np.append(future_vel_gt, np.expand_dims(vel_step[:2], axis=0), axis=0)

				agent_dict[id] = [pose, vel, rotated_grid, radial_pedestrian_grid, goal, future_pos_gt, future_vel_gt]

		# For each pedestrian, find other agents position for this time and compute pedestrian grid
		for id in agent_dict.keys():
			other_agents_pos = np.zeros([0,2])
			for other_id in agent_dict.keys():
				if other_id != id:
					other_agents_pos = np.append(other_agents_pos, np.expand_dims(agent_dict[other_id][0][:2], axis=0), axis=0)
			pos = agent_dict[id][0][:2]
			vel = agent_dict[id][1][:2]
			heading = math.atan2(vel[1], vel[0])
			other_pos_local_frame = sup.positions_in_local_frame(pos, heading, other_agents_pos)
			agent_dict[id][3] = sup.compute_radial_distance_vector(self.pedestrian_vector_dim, other_pos_local_frame,
																													max_range=self.max_range_ped_grid, min_angle=0, max_angle=2*np.pi,
																													normalize=True)

		return agent_dict

	def plot_global_scenario(self,batch_grid,batch_x,batch_agent_real_traj,batch_goal,other_agents_pos,batch_agent_predicted_traj,t,seq_index):
		"""
			inputs:
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_x: batch of initial states [x,y, vx, vy] on the global frame
				batch_agent_real_traj: agent ground truth velocity in global frame also know as batch_y
				batch_goal: batch with the trajectory goal positons
				other_agents_pos: other agent position in the global frame
				batch_agent_predicted_traj: agent predicted velocities in global frame
			"""
		global_grid = self.agent_container.occupancy_grid
		fig = pl.figure("Global Trajectory Predictions")
		ax_in = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)

		ax_in.clear()
		#plot scenario grid
		sup.plot_grid(ax_in, np.array([0.0, 0.0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
		ax_in.set_xlim([-global_grid.center[0], global_grid.center[0]])
		ax_in.set_ylim([-global_grid.center[1], global_grid.center[1]])
		#ax_in.set_aspect("equal")

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index, t, 1]*self.norm_const_y
		ax_in.plot(x_init_global_frame, y_init_global_frame, color='g', marker='o', label='Agent initial pos')
		ax_in.plot(batch_goal[seq_index,t, 0]*self.norm_const_x, batch_goal[seq_index,t, 1]*self.norm_const_y,
							 color='purple', marker='o', label='Agent goal')

		#plot predicted trajectory global frame
		vel_pred = np.zeros((self.args.prediction_horizon, 2))
		for mix_idx in range(self.args.n_mixtures):
			# plot predicted trajectory global frame
			for pred_step in range(self.args.prediction_horizon):
				idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
				idy = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures
				mu_x = batch_agent_predicted_traj[seq_index, t, idx]
				mu_y = batch_agent_predicted_traj[seq_index, t, idy]
				vel_pred[pred_step, :] = [mu_x, mu_y]

			traj_pred = sup.path_from_vel(initial_pos=np.array([batch_x[seq_index,t, 0]*self.norm_const_x, batch_x[seq_index,t, 1]*self.norm_const_y]),
																		pred_vel=vel_pred, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
		x_pred = traj_pred[:, 0] #* self.norm_const_x
		y_pred = traj_pred[:, 1] #* self.norm_const_y
		ax_in.plot(x_pred, y_pred, color='r',label='Predicted trajectory')

		#plot real trajectory global frame
		color_value = scalar_color_map.to_rgba(r)
		real_traj = np.zeros((self.args.prediction_horizon, 2))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = batch_agent_real_traj[seq_index, t, idx]
			mu_y = batch_agent_real_traj[seq_index, t, idy]
			real_traj[i, :] = [mu_x, mu_y]
		real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([batch_x[seq_index, t, 0]*self.norm_const_x, batch_x[seq_index, t, 1]*self.norm_const_y]),
												pred_vel=real_traj, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
		x = real_traj_global_frame[:, 0] #* self.norm_const_x
		y = real_traj_global_frame[:, 1] #* self.norm_const_y
		ax_in.plot(x, y, marker='x',color='b',label='Real trajectory')

		#Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			try:
				ax_in.plot(other_pos[0, 0], other_pos[0, 1], color='r', marker='o', label='Other agents')
				for jj in range(1,other_pos.shape[0]):
					ax_in.plot(other_pos[jj, 0], other_pos[jj, 1], color='r', marker='o') #other agent pos positon
			except IndexError:
				print("Oops!  That was no valid number.  Try again...")
		ax_in.legend()
		fig.canvas.draw()
		pl.show(block=False)

	def plot_local_scenario(self,batch_grid,batch_x,batch_agent_real_traj,batch_goal,other_agents_pos,batch_agent_predicted_traj,t,seq_index):
		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth trajectory in global frame
				batch_agent_predicted_traj: agent predicted trajectory in global frame
			"""
		#pl.close('all')
		global_grid = self.agent_container.occupancy_grid

		fig = pl.figure("Local Trajectory Predictions")
		ax_out = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)
		color_value = scalar_color_map.to_rgba(r)
		#Convert initial position to local frame
		#ini_pos_local_frame = sup.rotate_batch_x_to_local_frame(batch_x)
		ax_out.clear()

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index, t, 1]*self.norm_const_y

		# plot initial and goal pose
		ax_out.plot(0, 0, color='g', marker='o', label='Agent initial pos')

		x_lim = [-self.submap_size/2, self.submap_size/2]
		y_lim = [-self.submap_size/2, self.submap_size/2]
		ax_out.set_xlim(x_lim)
		ax_out.set_ylim(y_lim)

		# plot predicted trajectory local frame
		vel_pred = np.zeros((self.args.prediction_horizon, 2))
		for mix_idx in range(self.args.n_mixtures):
			# plot predicted trajectory global frame
			for pred_step in range(self.args.prediction_horizon):
				idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
				idy = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures
				mu_x = batch_agent_predicted_traj[seq_index, t, idx]
				mu_y = batch_agent_predicted_traj[seq_index, t, idy]
				vel_pred[pred_step, :] = [mu_x, mu_y]

			traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),
																		pred_vel=vel_pred, dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)
			x_pred = traj_pred[:, 0] #* self.norm_const_x
			y_pred = traj_pred[:, 1] #* self.norm_const_y
			ax_out.plot(x_pred, y_pred, color='r',label='Predicted trajectory')

		# plot real trajectory local frame does it make sense? isnt batch_y in the robot frame
		real_traj = np.zeros((self.args.prediction_horizon, 2))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = batch_agent_real_traj[seq_index, t, idx]
			mu_y = batch_agent_real_traj[seq_index, t, idy]
			real_traj[i, :] = [mu_x, mu_y]
		real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]),
																								 pred_vel=real_traj,
																							 dt=self.dt,n_vx=self.norm_const_vx,n_vy=self.norm_const_vy)

		# Query Agent Local frame plot
		if real_traj[0, 0] > 0.1:
			sup.plot_grid(ax_out, np.array([3.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))
		elif real_traj[0, 0] < -0.1:
			sup.plot_grid(ax_out, np.array([-3.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))
		else:
			sup.plot_grid(ax_out, np.array([0.0, 0.0]), batch_grid[seq_index,t,:,:], global_grid.resolution,
						  np.array([self.args.submap_size, self.args.submap_size]))

		x = real_traj_global_frame[:, 0]
		y = real_traj_global_frame[:, 1]
		ax_out.plot(x, y, marker='x', color='b', label='Agent real trajectory')
		#ax_out.plot(batch_x[seq_index,:, 0] * self.norm_const_x-x_init_global_frame*np.ones_like(batch_x[seq_index,:, 0]),
		#            batch_x[seq_index, :, 1] * self.norm_const_y-y_init_global_frame*np.ones_like(batch_x[seq_index,:, 0]), color='k')

		# Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			heading = math.atan2(batch_x[seq_index, t, 3], batch_x[seq_index, t, 2])

			try:
				ax_out.plot(other_pos[0, 0]-x_init_global_frame, other_pos[0, 1]-y_init_global_frame, color='r', marker='o',label='Other agents')
				for jj in range(1,other_pos.shape[0]):
					ax_out.plot(other_pos[jj, 0]-x_init_global_frame, other_pos[jj, 1]-y_init_global_frame, color='r', marker='o')#,markersize=23)  # other agent pos positon
					#ax_out.annotate(str(jj), xy=(other_pos[jj, 0]-x_init_global_frame , other_pos[jj, 1]-y_init_global_frame))
			except IndexError:
				print("Oops!  That was no valid number.  Try again...")

		circle = pl.Circle((0, 0), self.max_range_ped_grid, color='blue',
											 fill=False)
		ax_out.add_artist(circle)
		ax_out.set_aspect("equal")
		ax_out.legend()
		fig.canvas.draw()
				#sleep(0.5)  # Time in seconds.
		pl. show(block=False)

	def normalize_angle_positive(self,angle):
		""" Normalizes the angle to be 0 to 2*pi
				It takes and returns radians. """
		return angle % (2.0 * np.pi)

	def normalize_angle(self,angle):
		""" Normalizes the angle to be -pi to +pi
				It takes and returns radians."""
		a = self.normalize_angle_positive(angle)
		if a > np.pi:
			a -= 2.0 * np.pi
		return a

	def shortest_angular_distance(self,from_angle, to_angle):
		""" Given 2 angles, this returns the shortest angular
				difference.  The inputs and ouputs are of course radians.

				The result would always be -pi <= result <= pi. Adding the result
				to "from" will always get you an equivelent angle to "to".
		"""
		return self.normalize_angle(to_angle - from_angle)

	def get_radius(self,cutoff, A,var):

		return np.sqrt(-2 * var * np.log(cutoff / A))

	def getGaussian(self,x,y,x0,y0,varx,vary,skew):

		dx = x - x0
		dy = y - y0

		h = np.linalg.norm(np.array([dx,dy]))
		angle = np.arctan2(dy, dx)
		mx = np.cos(angle - skew) * h
		my = np.sin(angle - skew) * h
		f1 = np.square(mx) / (2.0 * varx)
		f2 = np.square(my) / (2.0 * vary)

		return self.amplitude_ * np.exp(-(f1 + f2))

	def plot_angular_grid(self, batch_ped_grid,t,seq_index):
		if self.multi_pedestrian:
			fig = pl.figure("Angular grid")
			ax_ped_grid = pl.subplot()
			ax_ped_grid.clear()
			n_sectors = batch_ped_grid.shape[2]
			grid = batch_ped_grid[seq_index,t,:]
			grid = np.expand_dims(grid, axis=1)
			grid_flipped = np.zeros_like(grid)
			grid_flipped[0:int(self.pedestrian_vector_dim/2)] = grid[-int(self.pedestrian_vector_dim/2):]
			grid_flipped[int(self.pedestrian_vector_dim / 2):] = grid[0:int(self.pedestrian_vector_dim / 2)]
			sup.plot_radial_distance_vector(ax_ped_grid, grid_flipped, max_range=1.0, min_angle=0.0,
																			max_angle=2 * np.pi)
			ax_ped_grid.plot(30, 30, color='r', marker='o', markersize=4)
			ax_ped_grid.arrow(0, 0, 1, 0, head_width=0.1,head_length=self.max_range_ped_grid) #agent poiting direction
		# x- and y-range only need to be [-1, 1] since the pedestrian grid is normalized
			ax_ped_grid.set_xlim([-self.max_range_ped_grid-1, self.max_range_ped_grid+1])
			ax_ped_grid.set_ylim([-self.max_range_ped_grid-1, self.max_range_ped_grid+1])

			fig.canvas.draw()
			# sleep(0.5)  # Time in seconds.
			pl.show(block=False)
			sleep(0.5)  # Time in seconds.

	def query_agent_velocity_direction(self, batch_grid, batch_x, batch_agent_real_traj,other_agents_pos,t,seq_index):
		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth velocity in global frame
				batch_agent_predicted_traj: agent predicted trajectory in global frame
			"""
		#pl.close('all')
		global_grid = self.agent_container.occupancy_grid
		fig = pl.figure("Query agent direction")
		ax_in = pl.subplot()

		#color definitions
		colormap = pl.get_cmap('rainbow')
		c_norm = pl.matplotlib.colors.Normalize(vmin=0, vmax=100)
		scalar_color_map = pl.cm.ScalarMappable(norm=c_norm, cmap=colormap)
		r= random.randint(0,100)

		number_samples =1
		ax_in.clear()
		#clear plots
		#plot scenario grid
		sup.plot_grid(ax_in, np.array([0.0, 0.0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
		ax_in.set_xlim([-global_grid.center[0], global_grid.center[0]])
		ax_in.set_ylim([-global_grid.center[1], global_grid.center[1]])

		#plot initial and goal pose
		x_init_global_frame = batch_x[seq_index,t, 0]*self.norm_const_x
		y_init_global_frame = batch_x[seq_index,t, 1]*self.norm_const_y
		ax_in.plot(x_init_global_frame, y_init_global_frame, color='g', marker='o', label='start')

		#plot real trajectory global frame
		color_value = scalar_color_map.to_rgba(r)

		#Plot velocity direction
		ax_in.arrow(x_init_global_frame, y_init_global_frame, batch_x[seq_index,t, 2]*self.norm_const_vx, batch_x[seq_index,t, 3]*self.norm_const_vy
				, head_width=0.1, head_length=self.max_range_ped_grid+0.5,alpha=0.4,label='Agent velocity direction') # angular grid detection distance
		circle = pl.Circle((x_init_global_frame,y_init_global_frame),self.max_range_ped_grid,color='blue',fill=False,label='Detection range')
		ax_in.add_artist(circle)
		#Plot other agent positions
		if self.multi_pedestrian:
			other_pos = other_agents_pos[seq_index][t]
			for jj in range(other_pos.shape[0]):
				ax_in.plot(other_pos[jj, 0], other_pos[jj, 1], color='r', marker='o') #other agent pos positon

		ax_in.set_aspect("equal")
		fig.canvas.draw()
		#sleep(0.5)  # Time in seconds.
		pl. show(block=False)

src/data_utils/DataHandlerRoboat.py
eliatrevisan/IA-MPPI-LBM
#!/usr/bin/env python
import sys
import numpy as np
from scipy.linalg import hankel

if sys.version_info[0] < 3:
    sys.path.append('../src/data_utils')
    import Support as sup
    from DataHandler import *
else:
    import src.data_utils.Support as sup
    from src.data_utils.DataHandlerLSTM import *

class DataHandler_Keras(DataHandlerLSTM):
    def __init__(self, args):
        super().__init__(args)
        self.tbpl = self.prev_horizon+1
        
        # Batches size redefinition
        self.batch_grid = np.zeros((self.batch_size, self.prev_horizon+1, int(self.args.submap_width / self.args.submap_resolution),
                                    int(self.args.submap_height / self.args.submap_resolution)))
        self.batch_x = np.zeros((self.batch_size, self.prev_horizon+1, self.args.input_dim))
        self.batch_vel = np.zeros((self.batch_size, self.prev_horizon+1, self.args.input_state_dim))
        self.batch_pos = np.zeros((self.batch_size, self.prev_horizon+1, self.args.input_state_dim))
        self.batch_goal = np.zeros((self.batch_size, self.prev_horizon+1, 2))
        self.other_agents_info = np.zeros((self.batch_size, self.prev_horizon+1, self.args.pedestrian_vector_dim))
        self.other_agents_pos = np.zeros((self.batch_size, self.prev_horizon+1, self.args.pedestrian_vector_state_dim))
        self.other_agents_vel = np.zeros((self.batch_size, self.prev_horizon+1, self.args.pedestrian_vector_state_dim))
        self.batch_y = np.zeros((self.batch_size, self.prediction_horizon, self.args.output_dim))
        self.pedestrian_grid = np.zeros([self.batch_size, self.tbpl, self.pedestrian_vector_dim])

    def getBatch(self):
        """
        		Get the next batch of training data.
        		"""
        # Update sequences
        # If batch sequences are empty and need to be filled
        trajectory = []
        if len(self.batch_sequences) == 0:
            for b in range(0, min(self.batch_size, len(self.trajectory_set))):
                id, trajectory = self.trajectory_set[self.data_idx]
                self.data_idx += 1
                self.batch_sequences.append(trajectory)
                self.batch_ids.append(id)
        # If batch sequences are filled and can be used or need to be updated.
        other_agents_pos = []
        new_epoch = False
        for ii, traj in enumerate(self.batch_sequences):
            if self.sequence_idx[ii] + self.tbpl + self.output_sequence_length + 1 >= len(traj):
                id, trajectory = self.trajectory_set[self.data_idx]
                self.data_idx = (self.data_idx + 1) % int(len(self.trajectory_set) * self.train_set)
                if self.data_idx == 0:
                    new_epoch = True
                self.batch_sequences[ii] = trajectory
                self.batch_ids[ii] = id
                self.sequence_idx[ii] = self.args.prev_horizon
                self.sequence_reset[ii] = 1
            else:
                self.sequence_reset[ii] = 0

        # Fill the batch
        other_agents_pos = []
        for ii in range(0, min(self.batch_size, len(self.trajectory_set) - len(self.trajectory_set) % self.batch_size)):
            traj = self.batch_sequences[ii]
            agent_id = self.batch_ids[ii]
            other_agents_pos.append(
                self.fillBatch(agent_id, ii, int(self.sequence_idx[ii]), self.tbpl, self.batch_x, self.batch_vel,
                               self.batch_pos, self.batch_grid, self.pedestrian_grid, self.batch_goal, self.batch_y,
                               traj, centered_grid=self.centered_grid))
            self.sequence_idx[ii] += self.tbpl
        
        batch = {
            'x': self.batch_x,
            'vel': self.batch_vel,
            'pos': self.batch_pos,
            'goal': self.batch_goal,
            'grid': self.batch_grid,
            'other_agents_info': self.other_agents_info,
            'y': self.batch_y,
            'other_agents_pos': self.other_agents_pos,
            'other_agents_vel': self.other_agents_vel,
            'new_epoch': new_epoch
        }
        return batch

    def fillBatch(self, agent_id, batch_idx, start_idx, truncated_backprop_length, batch_x, batch_vel, batch_pos,
                  batch_grid, pedestrian_grid, batch_goal, batch_y, trajectory, centered_grid=False):
        """
				Fill the data batches of batch_idx with data for all truncated backpropagation steps.
				"""

        for prev_step in range(self.prev_horizon,-1,-1):

            # Input values
            current_pos = np.array([trajectory.pose_vec[start_idx - prev_step, 0], trajectory.pose_vec[start_idx - prev_step, 1]])
            current_vel = np.array([trajectory.vel_vec[start_idx - prev_step, 0] , trajectory.vel_vec[start_idx - prev_step, 1]])

            if self.args.normalize_data:
                self.normalize_pos(current_pos)
                self.normalize_vel(current_vel)

            batch_x[batch_idx, prev_step, :] = np.array([current_pos[0],current_pos[1],current_vel[0],current_vel[1]])
            batch_vel[batch_idx, prev_step, :] = np.array([current_vel[0],current_vel[1]])

            heading = math.atan2(current_vel[1], current_vel[0])

            # Find positions of other pedestrians at the current timestep
            other_poses = trajectory.other_agents_positions[start_idx - prev_step]
            other_agents_pos = other_poses
            n_other_agents = other_poses.shape[0]
            if n_other_agents>0:
                other_velocities = trajectory.other_agents_velocities[start_idx - prev_step]
                other_pos_local_frame = sup.positions_in_local_frame(current_pos, heading, other_poses)

                # TODO: it only works for one agent now
                rel_pos = other_poses - current_pos
                rel_vel = other_velocities - current_vel
                distance = np.linalg.norm(rel_pos)
                pedestrian_grid[batch_idx, prev_step, :] = np.concatenate(
                  (np.array([np.linalg.norm(rel_pos)]), rel_vel[0]))

        # Output values
        for pred_step in range(self.output_sequence_length):
            vx = trajectory.vel_vec[start_idx +1+ pred_step, 0]
            vy = trajectory.vel_vec[start_idx +1+ pred_step, 1]
            px = trajectory.pose_vec[start_idx +1+ pred_step, 0]
            py = trajectory.pose_vec[start_idx +1+ pred_step, 1]
            batch_y[batch_idx, pred_step, 0] = vx
            batch_y[batch_idx, pred_step, 1] = vy
            if self.args.normalize_data:
                self.normalize_vel(batch_y[batch_idx, pred_step,:])
            batch_pos[batch_idx, pred_step, 0] = px
            batch_pos[batch_idx, pred_step, 1] = py

        return other_agents_pos
    
    def getTrajectoryAsBatch(self, trajectory_idx, max_sequence_length=1000,unit_testing=False):
        """
        Get a trajectory out of the trajectory set in the same format as for the standard training data
        (e.g. for validation purposes).
        """
        if unit_testing:
            traj = self.test_trajectory_set[trajectory_idx][1]
        else:
            traj = self.trajectory_set[trajectory_idx][1]

        sequence_length = min(max_sequence_length, len(traj) - self.prediction_horizon) - self.prev_horizon
        
        # Old data structures, used for plotting
        batch_x = np.zeros([sequence_length,self.prev_horizon+1, self.args.input_dim])  # data fed for training
        batch_pos = np.zeros([sequence_length,self.prev_horizon+1, self.args.input_state_dim])  # data fed for training
        batch_vel = np.zeros([sequence_length,self.prev_horizon+1, self.args.input_state_dim])  # data fed for training
        batch_goal = np.zeros([sequence_length, 2])
        batch_target = np.zeros([sequence_length, self.prediction_horizon,self.args.output_dim ])

        other_agents_info = np.zeros([sequence_length, self.prev_horizon+1, self.args.pedestrian_vector_dim])
        batch_grid = np.zeros((sequence_length, int(self.args.submap_width / self.args.submap_resolution),
                            int(self.args.submap_height / self.args.submap_resolution)))
        n_other_agents =  traj.other_agents_positions[0].shape[0]
        other_agents_pos = np.zeros((sequence_length, n_other_agents,self.args.pedestrian_vector_state_dim))
        other_agents_vel = np.zeros((sequence_length, self.args.pedestrian_vector_state_dim))
        
        # New data structures, fed to the network
        batch_vel_seq = np.zeros([sequence_length, self.prev_horizon+1, self.args.input_state_dim])
        other_agents_info_seq = np.zeros([sequence_length, self.prev_horizon+1, self.args.pedestrian_vector_dim])

        for batch_idx in range(sequence_length):
            other_agents_pos[batch_idx,:,:] = self.fillBatch(id, batch_idx, self.prev_horizon+batch_idx, sequence_length, batch_x, batch_vel, batch_pos,
                                              batch_grid, other_agents_info, batch_goal, batch_target, traj,
                                              centered_grid=self.centered_grid)

        batch = {
            'x': batch_x,
            'vel': batch_vel,
            'pos': batch_pos,
            'goal': batch_goal,
            'grid': batch_grid,
            'other_agents_info': other_agents_info,
            'y': batch_target,
            'other_agents_pos': other_agents_pos,
            'other_agents_vel': other_agents_vel,
            'traj': traj,
            'vel_seq': batch_vel_seq,
            'other_agents_info_seq': other_agents_info_seq
        }
        return batch
        
        
def expand_sequence(sequence_array, horizon): 
    # Sequence array has shape [features, time steps]
    # Expanded sequence will have shape [time steps, horizon, features]
    expanded_sequence = np.zeros((sequence_array.shape[1]-horizon+1,\
                                  horizon,\
                                  sequence_array.shape[0]))
    
    for i in range(sequence_array.shape[0]): # For each feature
        sequence = sequence_array[i, :]
        expanded_sequence[:, :, i] = hankel(sequence[0:horizon],\
                                            sequence[horizon-1:]).transpose()
    
    return expanded_sequence


def reduce_sequence(hankel_matrix):
    aux = []
    for feature_idx in range(hankel_matrix.shape[2]):
        aux.append( np.concatenate([hankel_matrix[0, :, feature_idx], hankel_matrix[1:, -1, feature_idx]], axis = 0) )
    return np.stack(aux, axis = 1)

src/data_utils/DataHandler_Keras.py
eliatrevisan/IA-MPPI-LBM
"""
New independent Keras DataHandler
Requires Python 3 and TensorFlow 2

List of admitted input 1 types:
> 'vel'
> 'polarvel'

List of admitted input 2 types (both strings for position and velocity components need to be contained in input2_type):
> For the position component:
>> 'posabs'
>> 'posrel'
>> 'posinvdist'
>> 'polarposabs'
>> 'polarposrel'
>> 'polarposinvdist'
> For the velocity component:
>> 'velabs'
>> 'velrel'
>> 'polarvelabs'
>> 'polarvelreldiff'
>> 'polarvelrelcos'
"""
#!/usr/bin/env python
import sys
import numpy as np
import pickle as pkl
from numpy.linalg import norm
from scipy.linalg import hankel
from scipy.spatial.distance import pdist
from math import atan2, cos, sin

class DataHandler():
    def __init__(self, args):
        self.args = args
        self.tbpl = self.args.truncated_backprop_length # unused
        self.prev_horizon = self.args.prev_horizon
        self.prediction_horizon = self.args.prediction_horizon
        
        self.input1_type = self.args.input1_type
        self.input2_type = self.args.input2_type

        self.shuffle = True
        self.batch_traj_idx = 0
        self.batch_size = self.args.batch_size
        self.input1_batch = np.zeros( [self.batch_size, self.prev_horizon, self.args.input_dim] )
        self.input2_batch = np.zeros( [self.batch_size, self.prev_horizon, self.args.pedestrian_vector_dim] )
        self.target_batch = np.zeros( [self.batch_size, self.prediction_horizon, self.args.output_pred_state_dim] )
        
        self.trajectory_set = []
        """ 
        trajectory_set is a list of trajectories
        > Each trajectory is a list of dictionary variables representing each data point
        >> Each datapoint contains the follwing keys:
        >>> 'control_command': Command of the robot
        >>> 'future_positions': 
        >>> 'goal_position': Goal position of the robot
        >>> 'pedestrian_goal_position': Goal position of the pedestrian
        >>> 'pedestrian_state': Dictionary which contains the current 'position' and 'velocity' of the pedestrian
        >>> 'predicted_cmd': Not used
        >>> 'robot_state': Current position and orientation of the robot
        >>> 'robot_velocity': Current speed of the robot
        """
        
        self.file = None

    def load_data(self):
        min_length_trajectory = self.prev_horizon + self.prediction_horizon
        self.file = open(self.args.data_path + self.args.dataset, 'rb')
        tmp_trajectory_set = pkl.load(self.file, encoding='latin1')
        trajectory_set = []
        
        for traj in tmp_trajectory_set:
            sub_traj = []
            # Sub-sample trajectories to make dt = 0.4
            for t_id in range(0, len(traj), 4):
            # for t_id in range(0, len(traj), 2):
                sub_traj.append(traj[t_id])
            # Append zero velocities to goal position
            if traj:
                goal_data = traj[-1].copy()
                goal_data["pedestrian_state"]["velocity"] = np.array([0, 0])
                for t_id in range(10):
                    sub_traj.append(goal_data)
                if len(sub_traj) > min_length_trajectory:
                    trajectory_set.append(sub_traj)
        self.trajectory_set = trajectory_set

        input1_arrays = []
        input2_arrays = []
        Y_arrays = []
        
        for trajectory_idx in range(len(self.trajectory_set)):
            feedArrays = self.getTrajectoryAsBatch(trajectory_idx)
            input1_arrays.append( feedArrays['input'][0] )
            input2_arrays.append( feedArrays['input'][1] )
            Y_arrays.append( feedArrays['target'] )
            
        self.input1 = np.concatenate(input1_arrays, axis = 0)
        self.input2 = np.concatenate(input2_arrays, axis = 0)
        self.target = np.concatenate(Y_arrays, axis = 0)
        self.n_data_samples = self.target.shape[0]
        
        if self.shuffle:
            idxs = np.random.permutation(self.n_data_samples)
            self.input1 = self.input1[idxs]
            self.input2 = self.input2[idxs]
            self.target = self.target[idxs]    

    def getBatch(self):
        if self.batch_traj_idx + self.batch_size <= self.n_data_samples: # Regular batches
            self.input1_batch = self.input1[self.batch_traj_idx:self.batch_traj_idx+self.batch_size,:,:]
            self.input2_batch = self.input2[self.batch_traj_idx:self.batch_traj_idx+self.batch_size,:,:]
            self.target_batch = self.target[self.batch_traj_idx:self.batch_traj_idx+self.batch_size,:,:]
            self.batch_traj_idx += self.batch_size
            new_epoch = False
            
        else: # Last batch, fill with random data points
            remaining_samples = self.n_data_samples - self.batch_traj_idx
            self.input1_batch[0:remaining_samples, :, :] = self.input1[self.batch_traj_idx:, :, :]
            self.input2_batch[0:remaining_samples, :, :] = self.input2[self.batch_traj_idx:, :, :]
            self.target_batch[0:remaining_samples, :, :] = self.target[self.batch_traj_idx:, :, :]
            
            # Take random data points outside of the remaining ones in order to completely fill up the batches
            idxs = np.random.permutation(self.batch_traj_idx)[0:self.batch_size-remaining_samples] 
            self.input1_batch[remaining_samples:, :, :] = self.input1[idxs, :, :]
            self.input2_batch[remaining_samples:, :, :] = self.input2[idxs, :, :]
            self.target_batch[remaining_samples:, :, :] = self.target[idxs, :, :]
            
            self.batch_traj_idx = 0
            new_epoch = True

        return {"input": [self.input1_batch, self.input2_batch],
                "target": self.target_batch,
                "new_epoch": new_epoch}

    def getTrajectoryAsBatch(self, trajectory_idx):
        trajectory = self.trajectory_set[trajectory_idx]
        input1_seq_list = []
        input2_seq_list = []

        for sample in trajectory:
            input1_data = self.getInput1Data(sample)
            input2_data = self.getInput2Data(sample)

            input1_seq_list.append(input1_data)
            input2_seq_list.append(input2_data)
            
        vel_seq =np.stack(input1_seq_list, axis = 1) # (features, timesteps)
        other_seq = np.stack(input2_seq_list, axis = 1) # (features, timesteps)
        
        input1_feedArray = expand_sequence(vel_seq[:, 0:-self.prediction_horizon], self.prev_horizon)
        input2_feedArray = expand_sequence(other_seq[:, 0:-self.prediction_horizon], self.prev_horizon)
        target_feedArray = expand_sequence(vel_seq[:, self.prev_horizon:], self.prediction_horizon)
        
        return {
            'input': [input1_feedArray, input2_feedArray],
            'target': target_feedArray,
            'trajectory': trajectory
        }
        
    def getPaddedTrajectoryAsBatch(self, trajectory_idx):
        trajectory = self.trajectory_set[trajectory_idx]
        input1_seq_list = []
        input2_seq_list = []

        for sample in trajectory:            
            input1_data = self.getInput1Data(sample)
            input2_data = self.getInput2Data(sample)

            input1_seq_list.append(input1_data)
            input2_seq_list.append( input2_data )
        vel_seq =np.stack(input1_seq_list, axis = 1) # (features, timesteps)
        other_seq = np.stack(input2_seq_list, axis = 1) # (features, timesteps)

        padded_vel_seq = np.concatenate( [np.zeros((2, self.prev_horizon-1)), vel_seq], axis = 1 )
        padded_other_seq = np.concatenate( [np.zeros((4, self.prev_horizon-1)), other_seq], axis = 1 )

        input1_feedArray = expand_sequence(padded_vel_seq, self.prev_horizon)
        input2_feedArray = expand_sequence(padded_other_seq, self.prev_horizon)

        return {
            'input': [input1_feedArray, input2_feedArray],
            'trajectory': trajectory
        }
        
    def getInput1Data(self, sample):
        ped_vel = sample['pedestrian_state']['velocity']
        
        if self.input1_type == "vel": # Velocity in cartesian coordinates
            # Velocity as data
            data = ped_vel
        elif self.input1_type == "polarvel": # Velocity in polar coordinates
            # Velocity modulus and direction as data
            modulus = norm(ped_vel)
            direction = atan2(ped_vel[1], ped_vel[0])
            data = np.array([modulus, direction])
        else:
            raise Exception("No valid type for Input 1")
        
        return data
    
    def getInput2Data(self, sample):
        ped_pos = sample['pedestrian_state']['position']
        ped_vel = sample['pedestrian_state']['velocity']
        robot_pos = sample['robot_state'][0:2]
        robot_vel = sample['robot_velocity']
        
        # Position component
        if "polarpos" in self.input2_type: # Position in polar coordinates
            if "posabs" in self.input2_type:
                modulus = norm(robot_pos)
                direction = atan2(robot_pos[1], robot_pos[0])
                pos = np.array([modulus, direction])
            elif "posrel" in self.input2_type:
                distance = robot_pos - ped_pos
                modulus = norm(distance)
                direction = atan2(distance[1], distance[0])
                pos = np.array([modulus, direction])
            elif "posinvdist" in self.input2_type:
                distance = robot_pos - ped_pos
                invdist = np.zeros(2)
                invdist[0] = 1/(distance[0] + 1)
                invdist[1] = 1/(distance[1] + 1)
                modulus = norm(invdist)
                direction = atan2(invdist[1], invdist[0])
                pos = np.array([modulus, direction])
            else:
                raise Exception("No valid position-component type for Input 2")
            
        else: # Position in cartesian coordinates
            if "posabs" in self.input2_type:
                pos = robot_pos
            elif "posrel" in self.input2_type:
                pos = robot_pos - ped_pos
            elif "posinvdist" in self.input2_type:
                distance = robot_pos - ped_pos
                invdist = np.zeros(2)
                invdist[0] = 1/(distance[0] + 1)
                invdist[1] = 1/(distance[1] + 1)
                pos = invdist
            else:
                raise Exception("No valid position-component type for Input 2")
        
        # Velocity component
        if "polarvel" in self.input2_type: # Velocity in polar coordinates
            if "velabs" in self.input2_type: # Polar absolute velocity
                modulus = norm(robot_vel)
                direction = atan2(robot_vel[1], robot_vel[0])
                vel = np.array([modulus, direction])
            elif "velreldiff" in self.input2_type: # Polar relative velocity  
                modulus = norm(robot_vel - ped_vel)
                direction_rob = atan2(robot_vel[1], robot_vel[0])
                direction_ped = atan2(robot_vel[1], robot_vel[0])
                angular_distance = direction_rob - direction_ped
                if angular_distance > 180:
                    angular_distance -= 360
                elif angular_distance < -180:
                    angular_distance += 360 
                vel = np.array([modulus, angular_distance])
            elif "velrelcos" in self.input2_type: # Polar relative velocity using cosine distance
                modulus = norm(robot_vel - ped_vel)
                distance = pdist([robot_vel, ped_vel], 'cosine')[0]
                if np.isnan(distance):
                    distance = 0
                vel = np.array([modulus, distance])
            else:
                raise Exception("No valid velocity-component type for Input 2")
                
        else: # Velocity in cartesian coordinates
            if "velabs" in self.input2_type:
                vel = robot_vel
            elif "velrel" in self.input2_type:
                vel = robot_vel - ped_vel
            else:
                raise Exception("No valid velocity-component type for Input 2")
        
        data =  np.concatenate( [pos, vel] )
        
        return data


def prediction_to_trajectory(initial_position, prediction, args):
    dt = args.dt
    output_type = args.input1_type
    
    trajectory_list = [ initial_position ]

    for pred_step in prediction:
        if output_type == "vel": # Velocity in cartesian coordinates
            new_point = trajectory_list[-1] + pred_step * dt
            
        elif output_type == "polarvel": # Velocity in polar coordinates
            modulus = pred_step[0]
            direction = pred_step[1]
            new_point = trajectory_list[-1] + modulus * np.array((cos(direction), sin(direction))) * dt
        else:
            raise Exception("Unimplemented output type")
        
        trajectory_list.append(new_point)
    trajectory = np.stack(trajectory_list, axis = 0)
    return trajectory

def expand_sequence(sequence_array, horizon): 
    # Sequence array has shape [features, time steps]
    # Expanded sequence will have shape [time steps, horizon, features]
    expanded_sequence = np.zeros((sequence_array.shape[1]-horizon+1,\
                                  horizon,\
                                  sequence_array.shape[0]))
    
    for i in range(sequence_array.shape[0]): # For each feature
        sequence = sequence_array[i, :]
        expanded_sequence[:, :, i] = hankel(sequence[0:horizon],\
                                            sequence[horizon-1:]).transpose()
    
    return expanded_sequence

def reduce_sequence(hankel_matrix):
    aux = []
    for feature_idx in range(hankel_matrix.shape[2]):
        aux.append( np.concatenate([hankel_matrix[0, :, feature_idx], hankel_matrix[1:, -1, feature_idx]], axis = 0) )
    return np.stack(aux, axis = 1)

src/data_utils/DataHandler_Keras_v2.py
eliatrevisan/IA-MPPI-LBM
import random

"""Code inspired in: https://github.com/fhennecker/deepdoom/blob/master/src/memory.py"""


class MemoryBuffer:
    def __init__(self, min_size=20, max_size=100):
        self.buffer = []
        self.min_size, self.max_size = min_size, max_size

    def full(self):
        return len(self.buffer) >= self.max_size

    def initialized(self):
        return len(self.buffer) >= self.min_size

    def add(self, step):
        if self.full():
            self.buffer.pop(0)
        self.buffer.append(step)

    def sample(self, batch_size):
        return [random.choice(self.buffer) for _ in range(batch_size)]
src/data_utils/MemoryBuffer.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import cv2 as cv
import src.data_utils.Support as sup
import matplotlib.pyplot as plt
import matplotlib.patches as patches
class OccupancyGrid():
	"""
	Occupancy grid class for capturing static object information.
	This occupancy grid is aligned with the Cartesian coordinate frame: 
		index 0: x-axis
		index 1: y-axis
	"""

	def __init__(self):
		self.gridmap = None
		self.resolution = None
		self.map_size = None
		#self.center = np.array([0.0, 0.0])
		self.center = np.array([None, None])
		
	def getIdx(self, pos_x, pos_y):
		"""
		Get indices of position. 
		pos_x and pos_y are the positions w.r.t. the center of the map.
		"""

		# Custom function to extract correct indices from Herengracht map
		idx_y = int((pos_x / self.resolution) + (self.center[0]*-1 / self.resolution)) 
		idx_x = self.map_size[1] - int((pos_y / self.resolution) + (self.center[1]*-1 / self.resolution))

		
		# Projecting index on map if out of bounds
		idx_x = max(0, min(idx_x, self.map_size[0] / self.resolution))
		idx_y = max(0, min(idx_y, self.map_size[1] / self.resolution))

		return idx_x, idx_y
		
	
	def getSubmapByIndices(self, center_idx_x, center_idx_y, span_x, span_y):
		"""
		Extract a submap of span (span_x, span_y) around 
		center index (center_idx_x, center_idx_y)
		"""
		debug_info = {}
		start_idx_x = max(0, int(center_idx_x - np.floor(span_x / 2)))
		start_idx_y = max(0, int(center_idx_y - np.floor(span_y / 2)))
		
		# Compute end indices (assure size of submap is correct, if out pf bounds)
		max_idx_x = self.gridmap.shape[0] - 1
		max_idx_y = self.gridmap.shape[1] - 1
		
		end_idx_x = start_idx_x + span_x
		if end_idx_x > max_idx_x:
			end_idx_x = max_idx_x
			start_idx_x = end_idx_x - span_x
		end_idx_y = start_idx_y + span_y
		if end_idx_y > max_idx_y:
			end_idx_y = max_idx_y
			start_idx_y = end_idx_y - span_y
		
		# Collect debug information
		debug_info["start_x"] = start_idx_x
		debug_info["start_y"] = start_idx_y
		debug_info["end_x"] = end_idx_x
		debug_info["end_y"] = end_idx_y

		#print("Shape in func: ", self.gridmap[start_idx_x:end_idx_x, start_idx_y:end_idx_y].shape)

		submap = self.gridmap[start_idx_x:end_idx_x, start_idx_y:end_idx_y]
		
		"""
		print(self.resolution)
		print(center_idx_x, center_idx_y)
		print(span_x, span_y)

		plt.figure()
		plt.imshow(self.gridmap, cmap="gray_r")
		plt.show()

		plt.figure()
		plt.imshow(submap, cmap="gray_r")
		plt.show()
		"""
		
		return submap, debug_info
		
	def getSubmapByCoords(self, center_pos_x, center_pos_y, size_x, size_y):
		"""
		Get submap around specified coordinates. 
		The sizes in x and y direction are within the same coordinate frame as the center coordinates.
		"""
		center_idx_x, center_idx_y = self.getIdx(center_pos_x, center_pos_y)
		span_x = int(size_x / self.resolution)
		span_y = int(size_y / self.resolution)

		#print("coordinates: ", center_pos_x, center_pos_y)
		
		return self.getSubmapByIndices(center_idx_x, center_idx_y, span_x, span_y)[0]

src/data_utils/OccupancyGrid.py
eliatrevisan/IA-MPPI-LBM
from copyreg import add_extension
import os
import numpy as np
import math
import pylab as pl
import sys
if sys.version_info[0] < 3:
	import Support as sup
else:
	import src.data_utils.Support as sup


def compute_trajectory_prediction_mse(args,ground_truth, predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	avg_mse = 0
	cnt = 0
	mse_list = []
	for pred, gt in zip(predictions, ground_truth):
		avg_mse = 0
		cnt = 0
		# compute average per trajectory
		for t in range(len(pred)):
			# real trajectory global frame
			real_vel_global_frame = gt.vel_vec[t+args.prev_horizon+1:t+args.prediction_horizon+args.prev_horizon+1,:2]
			real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=real_vel_global_frame,
			                                           dt=args.dt)
			vel_pred = np.zeros((args.prediction_horizon, args.output_dim))
			error = 0
			pred_t = pred[t][0]
			if args.n_mixtures<=1:
				# plot predicted trajectory global frame
				for sample_id in range(1):
					for i in range(args.prediction_horizon):
						idx = i * args.output_pred_state_dim
						idy = i * args.output_pred_state_dim + 1
						mu_x = pred_t[sample_id, idx]
						mu_y = pred_t[sample_id, idy]
						vel_pred[i, :] = [mu_x, mu_y]

				pred_vel_global_frame = vel_pred
				traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),pred_vel=np.squeeze(pred_vel_global_frame), dt=args.dt)
				for pred_step in range(args.prediction_horizon):
					error += np.linalg.norm(real_traj_global_frame[pred_step, :] - traj_pred[pred_step, :])/args.prediction_horizon
				avg_mse = (avg_mse*cnt+error)/(cnt+1)
				cnt += 1
			else:
				min_error = np.zeros((pred_t.shape[0]))
				for sample_id in range(pred_t.shape[0]):
					error = np.zeros((args.n_mixtures))
					for mix_idx in range(args.n_mixtures):
						# plot predicted trajectory global frame
						for i in range(args.prediction_horizon):
							idx = i * args.output_pred_state_dim * args.n_mixtures + mix_idx
							idy = i * args.output_pred_state_dim * args.n_mixtures + mix_idx + args.n_mixtures
							if args.normalize_data:
								mu_x = pred_t[sample_id,idx] / args.sx_vel + args.min_vel_x
								mu_y = pred_t[sample_id,idy] / args.sy_vel + args.min_vel_y
							else:
								mu_x = pred_t[sample_id,idx]
								mu_y = pred_t[sample_id,idy]
							vel_pred[i,:] = [mu_x, mu_y]

						traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),pred_vel=vel_pred, dt=args.dt)
						for pred_step in range(args.prediction_horizon):
							error[mix_idx] += np.linalg.norm(real_traj_global_frame[pred_step, :] - traj_pred[pred_step, :])/args.prediction_horizon
					min_error[sample_id] = min(error)
				avg_mse = (avg_mse * cnt + min(min_error)) / (cnt + 1)
				cnt += 1
		mse_list.append(avg_mse)
	return avg_mse, mse_list

def compute_trajectory_fde(args,ground_truth, predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	avg_fde = 0
	cnt = 0
	avg_fde_list = []
	real_traj_global_frame = np.zeros((args.prediction_horizon, args.output_dim))
	for pred, gt in zip(predictions, ground_truth):
		avg_fde = 0
		cnt = 0
		# compute average per trajectory
		for t in range(len(pred)):
			# real trajectory global frame
			real_vel_global_frame = gt.vel_vec[t+args.prev_horizon+1:t+args.prediction_horizon+args.prev_horizon+1,:2]
			real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=real_vel_global_frame,
			                                           dt=args.dt)
			vel_pred = np.zeros((args.prediction_horizon, args.output_dim))
			error = 0
			pred_t = pred[t][0]
			if args.n_mixtures<=1:
				# plot predicted trajectory global frame
				for sample_id in range(1):
					for i in range(args.prediction_horizon):
						idx = i * args.output_pred_state_dim
						idy = i * args.output_pred_state_dim + 1
						mu_x = pred_t[sample_id, idx]
						mu_y = pred_t[sample_id, idy]
						vel_pred[i, :] = [mu_x, mu_y]

				pred_vel_global_frame = vel_pred
				traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),pred_vel=np.squeeze(pred_vel_global_frame), dt=args.dt)
				error = np.linalg.norm(real_traj_global_frame[-1, :] - traj_pred[-1, :])
				avg_fde = (avg_fde*cnt+error)/(cnt+1)
				cnt += 1
			else:
				min_error = np.zeros((pred_t.shape[0]))
				for sample_id in range(pred_t.shape[0]):
					error = np.zeros((args.n_mixtures))
					for mix_idx in range(args.n_mixtures):
						# plot predicted trajectory global frame
						for i in range(args.prediction_horizon):
							idx = i * args.output_pred_state_dim * args.n_mixtures + mix_idx
							idy = i * args.output_pred_state_dim * args.n_mixtures + mix_idx + args.n_mixtures
							if args.normalize_data:
								mu_x = pred_t[sample_id,idx] / args.sx_vel + args.min_vel_x
								mu_y = pred_t[sample_id,idy] / args.sy_vel + args.min_vel_y
							else:
								mu_x = pred_t[sample_id,idx]
								mu_y = pred_t[sample_id,idy]
							vel_pred[i,:] = [mu_x, mu_y]

						traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),pred_vel=vel_pred, dt=args.dt)
						error[mix_idx] = np.linalg.norm(real_traj_global_frame[-1, :] - traj_pred[-1, :])
					min_error[sample_id] = min(error)
				avg_fde = (min(min_error)+ avg_fde*cnt)/(cnt+1)
				cnt +=1
			avg_fde_list.append(avg_fde)
	return avg_fde, avg_fde_list

def compute_rolling_trajectory_fde(args, horizon, ground_truth, predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	avg_fde = 0
	cnt = 0
	avg_fde_list = []
	real_traj_global_frame = np.zeros((args.prediction_horizon, args.output_dim))
	for pred, gt in zip(predictions, ground_truth):
		avg_fde = 0
		cnt = 0
		# compute average per trajectory
		for t in range(len(pred)):
			# real trajectory global frame
			real_vel_global_frame = gt.vel_vec[t+args.prev_horizon+1:t+args.prediction_horizon+args.prev_horizon+1,:2]
			#real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=real_vel_global_frame,
			#                                           dt=args.dt)
			real_traj_global_frame = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=real_vel_global_frame,
			                                           dt=args.dt)
			vel_pred = np.zeros((args.prediction_horizon, args.output_dim))
			error = 0
			pred_t = pred[t][0]
			if args.n_mixtures<=1:
				# plot predicted trajectory global frame
				for sample_id in range(1):
					for i in range(args.prediction_horizon):
						idx = i * args.output_pred_state_dim
						idy = i * args.output_pred_state_dim + 1
						mu_x = pred_t[sample_id, idx]
						mu_y = pred_t[sample_id, idy]
						vel_pred[i, :] = [mu_x, mu_y]

				pred_vel_global_frame = vel_pred
				traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]),pred_vel=np.squeeze(pred_vel_global_frame), dt=args.dt)
				error = np.linalg.norm(real_traj_global_frame[horizon, :] - traj_pred[horizon, :])
				avg_fde = (avg_fde*cnt+error)/(cnt+1)
				cnt += 1
			else:
				min_error = np.zeros((pred_t.shape[0]))
				for sample_id in range(pred_t.shape[0]):
					error = np.zeros((args.n_mixtures))
					for mix_idx in range(args.n_mixtures):
						# plot predicted trajectory global frame
						for i in range(args.prediction_horizon):
							idx = i * args.output_pred_state_dim * args.n_mixtures + mix_idx
							idy = i * args.output_pred_state_dim * args.n_mixtures + mix_idx + args.n_mixtures
							if args.normalize_data:
								mu_x = pred_t[sample_id,idx] / args.sx_vel + args.min_vel_x
								mu_y = pred_t[sample_id,idy] / args.sy_vel + args.min_vel_y
							else:
								mu_x = pred_t[sample_id,idx]
								mu_y = pred_t[sample_id,idy]
							vel_pred[i,:] = [mu_x, mu_y]

						traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),pred_vel=vel_pred, dt=args.dt)
						error[mix_idx] = np.linalg.norm(real_traj_global_frame[horizon, :] - traj_pred[horizon, :])
					min_error[sample_id] = min(error)
				avg_fde = (min(min_error)+ avg_fde*cnt)/(cnt+1)
				cnt +=1
			avg_fde_list.append(avg_fde)
			std = np.std(avg_fde_list)
	return avg_fde, avg_fde_list, std

def compute_2_wasserstein(args, predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	avg_wasser = 0
	cnt = 0
	wasser_list = []
	for prediction in predictions: # number of trajectories
		avg_wasser = 0
		cnt = 0
		# compute average per trajectory
		for pred_step in prediction: # trajectory length
			error = 0
			for i in range(args.prediction_horizon):  # prediction length
				if args.n_mixtures == 2:
					idx = i * args.output_pred_state_dim * args.n_mixtures
					idy = i * args.output_pred_state_dim * args.n_mixtures + args.n_mixtures
					mu_x = pred_step[0][ 0, idx]
					mu_y = pred_step[0][ 0, idy]
					mu_x1 = pred_step[0][ 0, idx+1]
					mu_y1 = pred_step[0][ 0, idy+1]
					sigma_x = pred_step[0][0, idy + args.n_mixtures]
					sigma_y = pred_step[0][0, idy + 2*args.n_mixtures]
					sigma_x1 = pred_step[0][0, idy + args.n_mixtures+1]
					sigma_y1 = pred_step[0][0, idy + 2*args.n_mixtures+1]
					wasser1 = np.linalg.norm(np.array([mu_x,mu_y])-np.array([mu_x1,mu_y1])) + sigma_x+sigma_x1+sigma_y+sigma_y1-2*np.sqrt(sigma_x*sigma_x1)-2*np.sqrt(sigma_y*sigma_y1)
					error += wasser1/args.n_mixtures
					avg_wasser = (avg_wasser * cnt + error) / (cnt + 1)
					cnt += 1
				if args.n_mixtures == 3:
					wasser1 = 0
					for mix_idx in range(args.n_mixtures):
						idx = i * args.output_pred_state_dim * args.n_mixtures + mix_idx
						idy = i * args.output_pred_state_dim * args.n_mixtures + args.n_mixtures + mix_idx
						mu_x = pred_step[0][ 0, idx]
						mu_y = pred_step[0][ 0, idy]
						sigma_x = pred_step[0][ 0, idy + args.n_mixtures]
						sigma_y = pred_step[0][ 0, idy + 2*args.n_mixtures]
						if mix_idx+1 >= args.n_mixtures:
							mu_x1 = pred_step[0][ 0, idx - mix_idx]
							mu_y1 = pred_step[0][ 0, idy - mix_idx]
							sigma_x1 = pred_step[0][ 0, idy + args.n_mixtures - mix_idx]
							sigma_y1 = pred_step[0][ 0, idy + 2 * args.n_mixtures - mix_idx]
						else:
							mu_x1 = pred_step[0][ 0, idx+1]
							mu_y1 = pred_step[0][ 0, idy+1]
							sigma_x1 = pred_step[0][ 0, idy + args.n_mixtures+1]
							sigma_y1 = pred_step[0][ 0, idy + 2*args.n_mixtures+1]
						wasser1 += np.linalg.norm(np.array([mu_x,mu_y])-np.array([mu_x1,mu_y1]))+sigma_x+sigma_x1+sigma_y+sigma_y1-2*np.sqrt(sigma_x*sigma_x1)-2*np.sqrt(sigma_y*sigma_y1)
					error += wasser1/args.n_mixtures
					avg_wasser = (avg_wasser * cnt + error) / (cnt + 1)
					cnt += 1
		wasser_list.append(avg_wasser)
	return avg_wasser, wasser_list

def compute_nll(args, ground_truth,predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	cnt = 0
	avg_list = []
	for pred, gt in zip(predictions, ground_truth): # number of trajectories
		avg_nll = 0
		cnt = 0
		# compute average per trajectory
		for t in range(pred.shape[0]): # trajectory length
			nll = 0
			for i in range(args.prediction_seq_length):  # prediction length
				term1 = 0
				for mix_id in range(args.n_mixtures):
					id = i * args.output_state_dim
					idx = i * args.output_pred_state_dim * args.n_mixtures
					idy = i * args.output_pred_state_dim * args.n_mixtures + args.n_mixtures
					mu_x = pred[t, idx+ mix_id]
					mu_y = pred[t, idy+ mix_id]
					sigma_x = max(pred[t, idy + args.n_mixtures+ mix_id],0)
					sigma_y = max(pred[t, idy + 2 * args.n_mixtures+ mix_id],0)
					pi = pred[t, idy + 3 * args.n_mixtures+ mix_id]
					normx = gt[t, id] - mu_x
					normy = gt[t, id + 1] - mu_y
					# Calculate sx*sy
					sxsy = sigma_x*sigma_y
					# Calculate the exponential factor
					z = np.square(normx/sigma_x) + np.square(normy/sigma_y)
					# Numerator
					result = np.exp(-z/2)
					# Normalization constant
					denom = 2 * np.pi * sxsy
					# Final PDF calculation
					result = result / denom
					term1 += pi*result
				nll += -np.log(np.maximum(term1, 1e-5))
			nll = nll / args.prediction_seq_length
			avg_nll = (avg_nll * cnt + nll) / (cnt + 1)
			cnt += 1
		avg_list.append(avg_nll)
	return avg_nll, avg_list

def compute_nll2(args, ground_truth,predictions):
	"""
		inputs:
			args: model parameters
			ground_truth: list of groudn truth velocities in absolute frame
			predictions: list of predicted velocities 		"""
	cnt = 0
	avg_list = []
	for pred, gt in zip(predictions, ground_truth): # number of trajectories
		avg_nll = 0
		cnt = 0
		# compute average per trajectory
		for t in range(pred.shape[0]): # trajectory length
			nll = 0
			pi = pred[t, args.prediction_seq_length*args.n_mixtures*args.output_pred_state_dim:]
			for i in range(args.prediction_seq_length):  # prediction length
				term1 = 0
				for mix_id in range(args.n_mixtures):
					id = i * args.output_state_dim
					idx = i * args.output_pred_state_dim * args.n_mixtures
					idy = i * args.output_pred_state_dim * args.n_mixtures + args.n_mixtures
					mu_x = pred[t, idx+ mix_id]
					mu_y = pred[t, idy+ mix_id]
					sigma_x = pred[t, idy + args.n_mixtures+ mix_id]
					sigma_y = pred[t, idy + 2 * args.n_mixtures+ mix_id]
					w = pred[t, idy + 3 * args.n_mixtures+ mix_id]
					normx = gt[t, id] - mu_x
					normy = gt[t, id + 1] - mu_y
					# Calculate sx*sy
					sxsy = sigma_x*sigma_y
					# Calculate the exponential factor
					z = (np.square(normx/sigma_x) + np.square(normy/sigma_y))/2.0
					# Numerator
					result = math.exp(-z)
					# Normalization constant
					denom = 2 * np.pi * sxsy
					# Final PDF calculation
					result2 = result / denom #pi[mix_id]*
					term1 += result2 / args.n_mixtures
				nll += -np.log(np.maximum(term1, 1e-5))
			nll = nll / args.prediction_seq_length
			avg_nll = (avg_nll * cnt + nll) / (cnt + 1)
			cnt += 1
		avg_list.append(avg_nll)
	return avg_nll, avg_list


def compute_ade(args, trajectories, predictions):
	"""
	Compute Average Displacement Error for the constant velocity predictions
	"""
	ade = []
	
	for idx in range(0,len(trajectories)):
		traj = trajectories[idx].pose_vec
		pred = predictions[idx]

		for step in range(0, len(pred)):
			tr = traj[step+args.prev_horizon+1:step+args.prev_horizon+args.prediction_horizon+1]
			pr = pred[step] 

			sum = 0
			for i in range(0, args.prediction_horizon):
				sum += np.linalg.norm(tr[i, :2] - pr[i, :])/args.prediction_horizon
			mse = sum #/ (i + 1)
			ade.append(mse)
	return np.mean(ade)

def compute_fde(args, trajectories, predictions):
	# Calculate Final Displacement Error for the CV predictions
	fde = []
	for idx in range(0, len(trajectories)):
		traj = trajectories[idx].pose_vec
		pred = predictions[idx]
		for step in range(0, len(pred)):
			tr = traj[step + args.prev_horizon + args.prediction_horizon] # Trajectory sample at horizon
			pr = pred[step][args.prediction_horizon-1] # Prediction at horizon
			error = np.linalg.norm(tr[:2]-pr[:])
			fde.append(error)
	return np.mean(fde)

def compute_rolling_fde(args, trajectories, predictions):
	res = []
	
	for idx in range(0,len(trajectories)):
		traj = trajectories[idx].pose_vec
		pred = predictions[idx]

		for step in range(0, len(pred)):

			fde = []
			for p in range(0, args.prediction_horizon):
				tr = traj[step + args.prev_horizon + p + 1] # Trajectory sample at horizon
				pr = pred[step][p] # Prediction at horizon
				error = np.linalg.norm(tr[:2]-pr[:])
				fde.append(error)
			res.append(fde)
	return np.mean(res, axis=0), np.std(res, axis=0)



src/data_utils/Performance.py
eliatrevisan/IA-MPPI-LBM
import os
import numpy as np
import cv2
import pylab as pl
from matplotlib.animation import FFMpegWriter
import sys
if sys.version_info[0] < 3:
	import Support as sup
else:
	import src.data_utils.Support as sup
import math
from matplotlib.patches import Ellipse
from scipy import interpolate
import matplotlib.pyplot as plt

def rgba2rgb(rgba):
    # rgba is a list of 4 color elements btwn [0.0, 1.0]
    # or a 2d np array (num_colors, 4)
    # returns a list of rgb values between [0.0, 1.0] accounting for alpha and background color [1, 1, 1] == WHITE
    if isinstance(rgba, list):
        alpha = rgba[3]
        r = max(min((1 - alpha) * 1.0 + alpha * rgba[0],1.0),0.0)
        g = max(min((1 - alpha) * 1.0 + alpha * rgba[1],1.0),0.0)
        b = max(min((1 - alpha) * 1.0 + alpha * rgba[2],1.0),0.0)
        return [r,g,b]
    elif rgba.ndim == 2:
        alphas = rgba[:,3]
        r = np.clip((1 - alphas) * 1.0 + alphas * rgba[:,0], 0, 1)
        g = np.clip((1 - alphas) * 1.0 + alphas * rgba[:,1], 0, 1)
        b = np.clip((1 - alphas) * 1.0 + alphas * rgba[:,2], 0, 1)
        return np.vstack([r,g,b]).T

class Recorder():

	def __init__(self,args,grid):

		# Parameters
		self.args = args
		self.gridmap = grid
		# Animation
		self.fig_animate = pl.figure('Animation')
		self.fig_width = 12  # width in inches
		self.fig_height = 12  # height in inches
		self.fig_size = [self.fig_width, self.fig_height]
		self.fontsize = 9
		self.colors = ['r', 'g', 'b','y','m','c','k']

		self.params = {'backend': 'ps',
		          'axes.labelsize': self.fontsize,
		          'font.size': self.fontsize,
		          'xtick.labelsize': self.fontsize,
		          'ytick.labelsize': self.fontsize,
		          'figure.figsize': self.fig_size}
		pl.rcParams.update(self.params)
		self.ax_pos = pl.subplot()
		#self.ax_pos_local = pl.subplot2grid((6, 12), (0, 6), colspan=6, rowspan=6)
		#self.ax_grid = pl.subplot2grid((9, 12), (7, 0), colspan=2, rowspan=2)
		#self.ax_ped_grid = pl.subplot2grid((9, 12), (7, 2), colspan=2, rowspan=2)
		#self.ax_vel = pl.subplot2grid((9, 12), (7, 4), colspan=2, rowspan=2)
		pl.rcParams.update(self.params)
		#pl.rcParams['animation.ffmpeg_path'] = "/home/bbrito/ffmpeg/ffmpeg"

		metadata = dict(title='Movie Test', artist='Matplotlib')
		self.writer = FFMpegWriter(fps=10, metadata=metadata,codec='mpeg4')

	def animate_local(self,input_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,trajectories,test_args):

		fig_animate = pl.figure('Local Scenario')
		fig_width = 12  # width in inches
		fig_height = 10  # height in inches
		fig_size = [fig_width, fig_height]
		fontsize = 26
		colors = ['r', 'g', 'y']
		params = {'backend': 'ps',
		               'axes.labelsize': fontsize,
		               'font.size': fontsize,
		               'xtick.labelsize': fontsize,
		               'ytick.labelsize': fontsize,
		               'figure.figsize': fig_size}

		pl.rcParams.update(params)
		ax_pos = pl.subplot()

		for animation_idx in range(len(input_list)):
			input = input_list[animation_idx]
			grid = grid_list[animation_idx]

			traj = trajectories[animation_idx]
			gt_vel = y_ground_truth_list[animation_idx]

			if not (y_pred_list_global is None):
				model_vel_pred = y_pred_list_global[animation_idx]
			else:
				model_vel_pred = None
			other_agents_pos = other_agents_list[animation_idx]

			for step in range(input.shape[0]): # trajectory length considering data from getTrajectoryAsBatch

				"""
				fig, ax = plt.subplots(1,2)
				sup.plot_grid_roboat(ax[0], self.gridmap.center, self.gridmap.gridmap, self.gridmap.resolution, self.gridmap.map_size)
				print("xy: ",input[step][0],input[step][1])
				ax[0].plot(input[step][0],input[step][1], markersize=50)
				ax[1].imshow(grid[step], cmap="gray_r")
				plt.show()

				"""
				self.plot_local_scenario(ax_pos,input, grid, model_vel_pred[step], gt_vel, other_agents_pos, step,
					                    rotate=False,n_samples=test_args.n_samples)

				ax_pos.axis("on")
				ax_pos.set_xlabel('x [m]',fontsize=26)
				ax_pos.set_ylabel('y [m]',fontsize=26)
				ax_pos.set_xlim([-10.0, 10.0])
				ax_pos.set_ylim([-10.0, 10.0])
				ax_pos.set_aspect('equal')
				pl.rcParams.update(params)
				ax_pos.set_aspect('equal')

				if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario+"/figs/"):
					os.makedirs(self.args.model_path + '/results/' + self.args.scenario+"/figs/")
				if test_args.save_figs:
					fig_animate.savefig(self.args.model_path + '/results/' + self.args.scenario+"/figs/local_"+str(animation_idx)+"_"+str(step)+".png")


	def animate(self,input_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,trajectories,test_args):

		#pl.show(block=False)
		if not os.path.exists(self.args.model_path + '/../videos/'):
			os.makedirs(self.args.model_path + '/../videos/')
		if test_args.unit_testing:
			video_file = self.args.model_path + '/../videos/' + str(self.args.exp_num) + "_unit_tests.mp4"
		elif test_args.freeze_other_agents:
			video_file = self.args.model_path + '/../videos/' + str(self.args.exp_num) + "_frozen.mp4"
		else:
			video_file = self.args.model_path + '/../videos/' + str(self.args.exp_num) + "_final.mp4"
		print("Recoding to: " + video_file)
		with self.writer.saving(self.fig_animate, video_file, 100):
			for animation_idx in range(len(input_list)):
				input = input_list[animation_idx]
				grid = grid_list[animation_idx]
				ped_grid = ped_grid_list[animation_idx]
				traj = trajectories[animation_idx]
				gt_vel = y_ground_truth_list[animation_idx]

				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				other_agents_pos = other_agents_list[animation_idx]

				for step in range(input.shape[0]): # trajectory length considering data from getTrajectoryAsBatch
					# Eucldean / position space
					self.ax_pos.clear()
					#self.ax_pos.plot(input[:, 0], input[:, 1], color='c', alpha=0.4, lw=1)
					self.ax_pos.plot(traj.pose_vec[-1, 0], traj.pose_vec[-1, 1], color='g', marker='o')
					#Real Trajectory
					self.ax_pos.plot(traj.pose_vec[:, 0], traj.pose_vec[:, 1], color='b', lw=2)

					self.ax_pos.plot(input[step, 0], input[step, 1], color='c', marker='o')
					bottom_left_submap = input[step, 0:2] - np.array([self.args.submap_size * self.gridmap.resolution / 2.0, self.args.submap_size * self.gridmap.resolution / 2.0])
					self.ax_pos.add_patch(pl.Rectangle(bottom_left_submap, self.args.submap_width * self.gridmap.resolution, self.args.submap_size * self.gridmap.resolution, fill=None))
					self.plot_local_scenario(input, grid, model_vel_pred[step], gt_vel, other_agents_pos, step,
					                    rotate=False,n_samples=test_args.n_samples)

					# Other agents
					for id in range(other_agents_pos[step].shape[0]): # number of agents
						self.ax_pos.plot(other_agents_pos[step][id, 0], other_agents_pos[step][id, 1], marker='o', color='b')
					vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
					sigmax = np.zeros((self.args.prediction_horizon, 1))
					sigmay = np.zeros((self.args.prediction_horizon, 1))
					#Predicted trajectory
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							# plot predicted trajectory global frame
							prediction_sample = model_vel_pred[step][0]
							for i in range(self.args.prediction_horizon):
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								mu_x = prediction_sample[0, idx]
								mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]

							pred_vel_global_frame = vel_pred
							traj_pred = sup.path_from_vel(initial_pos=np.array(
								[input[step, 0],
								 input[step, 1]]),
								pred_vel=pred_vel_global_frame, dt=self.args.dt)
							self.ax_pos.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[0], label='Predicted trajectory')
						else:
							for sample_id in range(test_args.n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								for mix_idx in range(self.args.n_mixtures):
									# plot predicted trajectory global frame
									for pred_step in range(self.args.prediction_horizon):
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										mu_x = prediction_sample[0, idx]
										mu_y = prediction_sample[0, idx + self.args.n_mixtures]
										if self.args.output_pred_state_dim>2:
											sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
											sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										vel_pred[pred_step, :] = [mu_x, mu_y]

									pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array(
										[input[step, 0],
										 input[step, 1]]),
										pred_vel=pred_vel_global_frame, dt=self.args.dt)

									self.ax_pos.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')

									if self.args.output_pred_state_dim > 2:
										# prior of 0 on the uncertainty of the pedestrian velocity
										sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
										sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

										e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
										             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
										e1.set_alpha(0.5)
										self.ax_pos.add_patch(e1)
										#self.ax_pos.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')
										for pred_step in range(1, self.args.prediction_horizon):
											sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
											sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt
											e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
											             height=np.sqrt(sigma_y) / 2,
											             angle=0 / np.pi * 180)
											self.ax_pos.add_patch(e1)

					# Real trajectory
					model_vel_real = np.zeros((self.args.prediction_horizon, self.args.output_dim))
					for i in range(self.args.prediction_horizon):
						idx = i * self.args.output_dim
						idy = i * self.args.output_dim + 1
						mu_x = gt_vel[step, idx]
						mu_y = gt_vel[step, idy]
						model_vel_real[i, :] = [mu_x, mu_y]
					if False:
						real_vel = sup.rotate_predicted_vel_to_global_frame(model_vel_real, batch_x[step, 2:])
					else:
						real_vel = model_vel_real
					traj_real = sup.path_from_vel(initial_pos=np.array([0, 0]),
					                              pred_vel=real_vel, dt=self.args.dt)
					self.ax_pos_local.plot(traj_real[:, 0], traj_real[:, 1], color='m', label='Agent real trajectory')

					sup.plot_grid(self.ax_pos, np.array([0.0, 0.0]), self.gridmap.gridmap, self.gridmap.resolution,
					              self.gridmap.map_size)

					ax_pos.axis('equal')
					self.ax_pos.axis("on")
					self.ax_pos.set_xlabel('x [m]')
					self.ax_pos.set_ylabel('y [m]')
					self.ax_pos.set_aspect('equal')

					#self.fig_animate.canvas.draw()
					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario+"/figs/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario+"/figs/")
					if test_args.save_figs:
						self.fig_animate.savefig(self.args.model_path + '/results/' + self.args.scenario+"/figs/result_"+str(animation_idx)+"_"+str(step)+".jpg")
					self.writer.grab_frame()
				

	def animate_group_of_agents(self,trajectories,y_pred_list_global,test_args=None):

		if not os.path.exists(self.args.model_path + '/videos/'):
			os.makedirs(self.args.model_path + '/videos/')

		fig_animate = pl.figure('Global Scenario')
		fig_width = 12  # width in inches
		fig_height = 10  # height in inches
		fig_size = [fig_width, fig_height]
		fontsize = 26

		params = {'backend': 'ps',
		          'axes.labelsize': fontsize,
		          'font.size': fontsize,
		          'xtick.labelsize': fontsize,
		          'ytick.labelsize': fontsize,
		          'figure.figsize': fig_size}

		plt_colors = []
		plt_colors.append([0.8500, 0.3250, 0.0980])  # orange
		plt_colors.append([0.0, 0.4470, 0.7410])  # blue
		plt_colors.append([0.4660, 0.6740, 0.1880])  # green
		plt_colors.append([0.4940, 0.1840, 0.5560])  # purple
		plt_colors.append([0.9290, 0.6940, 0.1250])  # yellow
		plt_colors.append([0.3010, 0.7450, 0.9330])  # cyan
		plt_colors.append([0.6350, 0.0780, 0.1840])  # chocolate
		plt_colors.append([0.6350, 0.0780, 0.1840])  # chocolate
		plt_colors.append([0.6350, 0.0780, 0.1840])  # chocolate
		scenario = self.args.scenario.split('/')[0]
		ax_pos = pl.subplot()

		ax_pos.set_xlabel('x [m]')
		ax_pos.set_ylabel('y [m]')
		ax_pos.axis('on')
		ax_pos.set_aspect('equal')

		n_agents =  len(trajectories[0])
		axbackground = fig_animate.canvas.copy_from_bbox(ax_pos.bbox)

		# Create plotting variables
		current_pos_list_ = []
		pred_line_list_ = []
		real_line_list_ = []
		ellipses_list_ = []
		for ag_id in range(n_agents):
			current_pos_, = ax_pos.plot([], color=plt_colors[ag_id], marker='o')
			current_pos_list_.append(current_pos_)
			pred_line, = ax_pos.plot([], color=plt_colors[ag_id], label='Predicted trajectory Agent ' + str(ag_id))
			pred_line_list_.append(pred_line)
			real_line, = ax_pos.plot([], color=plt_colors[ag_id], alpha=0.5, label='Real trajectory Agent ' + str(ag_id))
			real_line_list_.append(real_line)
			c = rgba2rgb(plt_colors[ag_id] + [float(1)])
			ellipses = [Ellipse(xy=(0, 0), width=0.6,
			                    height=0.6,
			                    fc=c, ec=plt_colors[ag_id],
			                    fill=True) for pred_step in range(self.args.prediction_horizon)]
			for ellipse in ellipses:
				ax_pos.add_artist(ellipse)
			ellipses_list_.append(ellipses)

		ax_pos.legend()

		with self.writer.saving(fig_animate, self.args.model_path + '/videos/' +str(self.args.exp_num)+scenario+".mp4", 100):
			for animation_idx in range(len(trajectories)):
				traj = trajectories[animation_idx]

				model_vel_pred = y_pred_list_global[animation_idx]

				vel_pred = np.zeros((self.args.prediction_horizon, self.args.output_pred_state_dim))
				sigmax = np.zeros((self.args.prediction_horizon, 1))
				sigmay = np.zeros((self.args.prediction_horizon, 1))

				ax_pos.set_xlim([0.0, 20.0])
				ax_pos.set_ylim([-5.0, 5.0])

				for step in range(traj[0].time_vec.shape[0]):
					for ag_id in range(n_agents):
					  # trajectory length considering data from getTrajectoryAsBatch
						fig_animate.canvas.restore_region(axbackground)

						try:
							current_pos_list_[ag_id].set_data(traj[ag_id].pose_vec[step,0]/self.args.sx_pos,traj[ag_id].pose_vec[step, 1]/self.args.sy_pos)
						except:
							current_pos_list_[ag_id].set_data(traj[ag_id].pose_vec[-1, 0] / self.args.sx_pos,
							                                  traj[ag_id].pose_vec[-1, 1] / self.args.sy_pos)
						#Predicted trajectory
						for sample_id in range(1):
							try:
								prediction_sample = model_vel_pred[ag_id][step][sample_id]
								init_pose = traj[ag_id].pose_vec[step,:2]
							except:
								prediction_sample =  np.zeros([self.args.output_dim*self.args.prediction_horizon])
								init_pose = traj[ag_id].pose_vec[-1, :2]
							for mix_idx in range(self.args.n_mixtures):
								# plot predicted trajectory global frame
								for pred_step in range(self.args.prediction_horizon):
									idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
									mu_x = prediction_sample[idx]/self.args.sx_vel
									mu_y = prediction_sample[idx + self.args.n_mixtures]/self.args.sy_vel
									if self.args.output_pred_state_dim > 2:
										sigmax[pred_step, :] = prediction_sample[idx + 2 * self.args.n_mixtures]
										sigmay[pred_step, :] = prediction_sample[idx + 3 * self.args.n_mixtures]
									vel_pred[pred_step, :] = [mu_x, mu_y]

								pred_vel_global_frame = vel_pred
								traj_pred = sup.path_from_vel(initial_pos=init_pose/self.args.sx_pos,
								                              pred_vel=pred_vel_global_frame, dt=self.args.dt)

								pred_line_list_[ag_id].set_data(traj_pred[:, 0], traj_pred[:, 1])

								if self.args.output_pred_state_dim > 2:
									# TODO: ADAPT TO NEW CODE
									# prior of 0 on the uncertainty of the pedestrian velocity
									sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
									sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

									e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
									             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
									e1.set_alpha(0.5)

									for pred_step in range(1, self.args.prediction_horizon):
										sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
										sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt

										e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
										             height=np.sqrt(sigma_y) / 2,
										             angle=0 / np.pi * 180)
										ax_pos.add_patch(e1)
								else:
									for pred_step in range(0, self.args.prediction_horizon):
										alpha = 1/ (1+pred_step)
										c = rgba2rgb(plt_colors[ag_id] + [float(alpha)])
										ellipses_list_[ag_id][pred_step].set_facecolor(c)
										ellipses_list_[ag_id][pred_step].set_center((traj_pred[pred_step, 0], traj_pred[pred_step, 1]))

						# Real trajectory
						real_line_list_[ag_id].set_data(traj[ag_id].pose_vec[:,0]/self.args.sx_pos, traj[ag_id].pose_vec[:,1]/self.args.sy_pos)

					fig_animate.canvas.blit(ax_pos.bbox)
					self.writer.grab_frame()

					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario + "/figs/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario + "/figs/")
					if test_args.save_figs:
						fig_animate.savefig(self.args.model_path + '/results/' + self.args.scenario + "/figs/result_" + str(
							animation_idx) + "_" + str(step) + ".jpg")

					fig_animate.canvas.flush_events()

	def animate_global(self,input_list,grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,y_ground_truth_list_exp,all_traj_likelihood, cv_pred_list, test_args=None ):

		if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario):
			os.makedirs(self.args.model_path + '/results/' + self.args.scenario)

		fig_animate = pl.figure('Global Scenario')
		fig_animate.tight_layout()
		fig_width = 12  # width in inches
		fig_height = 10  # height in inches
		fig_size = [fig_width, fig_height]
		fontsize = 26

		params = {'backend': 'ps',
		          'axes.labelsize': fontsize,
		          'font.size': fontsize,
		          'xtick.labelsize': fontsize,
		          'ytick.labelsize': fontsize,
		          'figure.figsize': fig_size,
		          'legend.loc': 'upper right',
		          'legend.fontsize': 14}

		#pl.rcParams.update(params)

		plt_colors = []
		plt_colors.append([0.8500, 0.3250, 0.0980])  # orange
		plt_colors.append([0.0, 0.4470, 0.7410])  # blue
		scenario = self.args.scenario.split('/')[0]
		ax_pos = pl.subplot()

		ax_pos.set_xlabel('x [m]')
		ax_pos.set_ylabel('y [m]')
		ax_pos.axis('on')
		#ax_pos.set_aspect('equal')

		# Load Map
		"""
		map_file = os.path.join(self.args.data_path + self.args.scenario, 'map.png')
		if os.path.exists(map_file):
			#img = np.uint8(cv2.imread(map_file) * -1 + 255)
			img = cv2.imread(map_file)
			#img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
			resolution = 0.1
			width = int(img.shape[1] * resolution)
			height = int(img.shape[0] * resolution)
			dim = (width, height)
			resized = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
		else:
			print('[INF] No map file')
		ax_pos.imshow(resized,extent=[0, 30, -3, 2])
		"""
		#sup.plot_grid(ax_pos, np.array([0.0, 0.0]), self.gridmap.gridmap, self.gridmap.resolution,  # np.array([0.0, 0.0])
		#              self.gridmap.map_size)

		sup.plot_grid_roboat(ax_pos, self.gridmap.center, self.gridmap.gridmap, self.gridmap.resolution,  # np.array([0.0, 0.0])
		              self.gridmap.map_size)

		axbackground = fig_animate.canvas.copy_from_bbox(ax_pos.bbox)
		current_pos_, = ax_pos.plot([], color='b', marker='o')
		other_agent_traj, = ax_pos.plot([], color='r', marker='.', alpha=0.5, label='Real Trajectory Agent 2')
		other_agent_traj_, = ax_pos.plot([], color='r', marker='.', alpha=0.5, label='Real Trajectory Agent 3')
		other_agent_traj_2, = ax_pos.plot([], color='r', marker='.', alpha=0.5, label='Real Trajectory Agent 4')
		cv_predictions, = ax_pos.plot([], color='g', marker='.', alpha=0.5, label='Constant Velocity')
		other_agents_line_ = []
		pred_line_list = []
		for i in range(self.args.n_mixtures):
			pred_line, = ax_pos.plot([], color=self.colors[2], label='Predicted trajectory Agent ' + str(i))
			pred_line_list.append(pred_line)

		real_line, = ax_pos.plot([], color='c',marker='.',alpha=0.5, label='Real trajectory Agent 1')
		c = rgba2rgb(plt_colors[1] + [float(1)])

		ellipses = []
		for i in range(self.args.n_mixtures):
			ellipse = [Ellipse(xy=(0,0), width=0.6,
				             height=0.6,
				             fc=c, ec=plt_colors[1],
				             fill=True) for pred_step in range(self.args.prediction_horizon)]
			ellipses.append(ellipse)

		for ellipse in ellipses:
			for el in ellipse:
				ax_pos.add_artist(el)

		#ax_pos.legend()


		with self.writer.saving(fig_animate, self.args.model_path + '/results/' + self.args.scenario+"/video.mp4", 100):
			for animation_idx in range(0,len(input_list)):
				input = input_list[animation_idx] #* (1 / self.gridmap.resolution)
				grid = grid_list[animation_idx]
				if test_args.constant_velocity:
					cv_pred = cv_pred_list[animation_idx]

				#ax_pos.set_xlim([int(np.min(input/self.args.sx_pos))-2.0, int(np.max(input/self.args.sx_pos))+2.0])
				#ax_pos.set_ylim([int(np.min(input/self.args.sy_pos))-2.0, int(np.max(input/self.args.sy_pos))+2.0])
				#ax_pos.set_xlim([-self.gridmap.map_size[0]/2, self.gridmap.map_size[0]/2])
				#ax_pos.set_ylim([-self.gridmap.map_size[1]/2, self.gridmap.map_size[1]/2])
				#ax_pos.set_xlim([self.gridmap.map_size[0], -self.gridmap.map_size[0]])
				#ax_pos.set_ylim([self.gridmap.map_size[1], -self.gridmap.map_size[1]])

				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				other_agents_pos = other_agents_list[animation_idx]
				traj_likelihood = all_traj_likelihood[animation_idx]

				number_of_agents = len(other_agents_pos[0])
				for ag_id in range(number_of_agents):
					other_agent, = ax_pos.plot([], color='r', marker='o')
					other_agents_line_.append(other_agent)
				vel_pred = np.zeros((self.args.prediction_horizon, 2))
				sigmax = np.zeros((self.args.prediction_horizon))
				sigmay = np.zeros((self.args.prediction_horizon))

				### 
				# Plotting agent real trajectories

				other_agent_traj.set_data(None,None)
				other_agent_traj_.set_data(None,None)
				other_agent_traj_2.set_data(None,None)
			
				x_real = list()
				y_real = list()
				for test in range(input.shape[0]):
					x_real.append(other_agents_pos[test][:,0])
					y_real.append(other_agents_pos[test][:,1])

				if len(x_real[0]) == 1:
					other_agent_traj.set_data(x_real[:], y_real[:])
				elif len(x_real[0]) == 2:
					# Two (other) agents
					n_agents = len(x_real[0])
 
					x_real_one = list()
					x_real_two = list()
					y_real_one = list()
					y_real_two = list()

					for test in range(input.shape[0]):
						x_real_one.append(x_real[test][0])
						x_real_two.append(x_real[test][1])
						y_real_one.append(y_real[test][0])	
						y_real_two.append(y_real[test][1])

					other_agent_traj.set_data(x_real_one[:], y_real_one[:])
					other_agent_traj_.set_data(x_real_two[:], y_real_two[:])
				elif len(x_real[0]) == 3:
					# Three (other) agents
					n_agents = len(x_real[0])
 
					x_real_one = list()
					x_real_two = list()
					x_real_three = list()
					y_real_one = list()
					y_real_two = list()
					y_real_three = list()

					for test in range(input.shape[0]):
						x_real_one.append(x_real[test][0])
						x_real_two.append(x_real[test][1])
						x_real_three.append(x_real[test][2])
						y_real_one.append(y_real[test][0])	
						y_real_two.append(y_real[test][1])
						y_real_three.append(y_real[test][2])

					other_agent_traj.set_data(x_real_one[:], y_real_one[:])
					other_agent_traj_.set_data(x_real_two[:], y_real_two[:])
					other_agent_traj_2.set_data(x_real_three[:], y_real_three[:])
	

				for step in range(input.shape[0]): # trajectory length considering data from getTrajectoryAsBatch
					fig_animate.canvas.restore_region(axbackground)
					# Eucldean / position space
					current_pos_.set_data(input[step, 0]/self.args.sx_pos, input[step, 1]/self.args.sy_pos)

					if test_args.constant_velocity:
						cv_predictions.set_data(cv_pred[step][:,0], cv_pred[step][:,1])
					
					bottom_left_submap = input[step, 0:2] - np.array([self.args.submap_size / 2.0, self.args.submap_size / 2.0])

					# Other agents
					for id in range(other_agents_pos[step].shape[0]): # number of agents
						other_agents_line_[id].set_data(other_agents_pos[step][id, 0], other_agents_pos[step][id, 1])
					#Predicted trajectory
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							# plot predicted trajectory global frame
							for i in range(self.args.prediction_horizon):
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								mu_x = model_vel_pred[step, idx]/self.args.sx_vel
								mu_y = model_vel_pred[step, idy]/self.args.sy_vel
								vel_pred[0, 2 * i:2 * i + 2] = [mu_x, mu_y]

							pred_vel_global_frame = vel_pred
							traj_pred = sup.path_from_vel(initial_pos=np.array(
								[input[step, 0],
								 input[step, 1]]),
								pred_vel=np.squeeze(pred_vel_global_frame), dt=self.args.dt)
							pred_line_list[mix_idx].set_data(traj_pred[:, 0], traj_pred[:, 1])
						else:
							for sample_id in range(1):
								prediction_sample = model_vel_pred[step][sample_id]
								most_likely_traj = np.argmax(traj_likelihood[step][0,0])
								for mix_idx in range(self.args.n_mixtures):
									# plot predicted trajectory global frame
									for pred_step in range(self.args.prediction_horizon):
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										mu_x = prediction_sample[0, idx]/self.args.sx_vel
										mu_y = prediction_sample[0, idx + self.args.n_mixtures]/self.args.sy_vel
										if self.args.output_pred_state_dim > 2:
											sigmax[pred_step] = np.maximum(np.minimum(prediction_sample[0, idx + 2 * self.args.n_mixtures],0.5),0.05)
											sigmay[pred_step] = np.maximum(np.minimum(prediction_sample[0, idx + 3 * self.args.n_mixtures],0.5),0.05)
										vel_pred[pred_step, :] = [mu_x, mu_y]

									pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=input[step,:2]/self.args.sx_pos,
									                              pred_vel=pred_vel_global_frame, dt=self.args.dt)
									# hack for prediciting positions
									#traj_pred = vel_pred+input[step,:2]
									if mix_idx != most_likely_traj:
										pred_line_list[mix_idx].set_data(traj_pred[:, 0], traj_pred[:, 1])
										pred_line_list[mix_idx].set_color(self.colors[2])
									else:
										pred_line_list[mix_idx].set_data(traj_pred[:, 0], traj_pred[:, 1])
										pred_line_list[mix_idx].set_color(self.colors[1])
									pred_line_list[mix_idx].set_label('Predicted trajectory ' + str(traj_likelihood[step][0,0,mix_idx]))
									
									if self.args.output_pred_state_dim > 2:
										# prior of 0 on the uncertainty of the pedestrian velocity
										sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
										sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

										alpha = 1
										c = rgba2rgb(plt_colors[0] + [float(alpha)])
										ellipses[mix_idx][0].set_facecolor(c)
										ellipses[mix_idx][0].set_center((traj_pred[0, 0], traj_pred[0, 1]))
										ellipses[mix_idx][0].width = np.sqrt(sigma_x)
										ellipses[mix_idx][0].height = np.sqrt(sigma_y)
										for pred_step in range(1, self.args.prediction_horizon):
											sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
											sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt
											#sigma_x += sigmax[pred_step]
											#sigma_y += sigmay[pred_step]

											alpha = 1 / (1 + pred_step)
											c = rgba2rgb(plt_colors[0] + [float(alpha)])
											ellipses[mix_idx][pred_step].set_facecolor(c)
											ellipses[mix_idx][pred_step].set_center((traj_pred[pred_step, 0], traj_pred[pred_step, 1]))
											ellipses[mix_idx][pred_step].width = np.sqrt(sigma_x)
											ellipses[mix_idx][pred_step].height = np.sqrt(sigma_y)
									else:
										for pred_step in range(0, self.args.prediction_horizon):
											alpha = 1/ (1+pred_step)
											c = rgba2rgb(plt_colors[1] + [float(alpha)])
											ellipses[mix_idx][pred_step].set_facecolor(c)
											ellipses[mix_idx][pred_step].set_center((traj_pred[pred_step, 0], traj_pred[pred_step, 1]))
									
									"""
									# From animate local script
									if self.args.output_pred_state_dim > 2:
										# prior of 0 on the uncertainty of the pedestrian velocity
										sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
										sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

										e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
																height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
										e1.set_alpha(0.5)

										for pred_step in range(1, self.args.prediction_horizon):
											sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
											sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt

											e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
																	height=np.sqrt(sigma_y) / 2,
																	angle=0 / np.pi * 180)
											ax_pos.add_patch(e1)
										"""

					# Real trajectory
					vel_real = np.zeros((self.args.prediction_horizon, self.args.output_dim))

					# plot trajectory global frame
					real_line.set_data(input[:, 0]/self.args.sx_pos, input[:, 1]/self.args.sy_pos)

					#sup.plot_grid(ax_pos, np.array([0.0, 0.0]), self.gridmap.gridmap, self.gridmap.resolution,
					#              self.gridmap.map_size)
					fig_animate.set_tight_layout(True)
					fig_animate.canvas.blit(ax_pos.bbox)

					self.writer.grab_frame()

					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario + "/figs/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario + "/figs/")
					if test_args.save_figs:
						fig_animate.savefig(self.args.model_path + '/results/' + self.args.scenario + "/figs/result_" + str(
							animation_idx) + "_" + str(step) + ".png")

					fig_animate.canvas.flush_events()


	def plot_local_scenario(self,ax_pos_local, batch_x, batch_grid, model_vel_pred, model_vel_real, other_agents_pos,step,
	            rotate=False,n_samples=1):

		"""
			inputs:
				global_grid: map grid of the scenario
				batch_grid: batch_grid numpy array (gridmap) in the global frame [Batch_size][truncated_back_propagation_time][width][height]
				batch_agent_real_traj: agent ground truth trajectory in global frame
				model_vel_pred: agent predicted trajectory in global frame
			"""
		ax_pos_local.clear()


		#plot initial and goal pose
		x_init_global_frame = batch_x[step, 0]
		y_init_global_frame = batch_x[step, 1]

		# plot initial and goal pose
		ax_pos_local.plot(0, 0, color='c', marker='o', label='Agent initial pos')
		

		# Local Frame Goal Position
		x_goal = batch_x[-1, 0] - x_init_global_frame
		y_goal = batch_x[-1, 1] - y_init_global_frame
		#self.ax_pos_local.plot(x_goal, y_goal, color='g', marker='o', label='Agent goal pos')

		# plot predicted trajectory local frame
		if not (model_vel_pred is None):
			vel_pred = np.zeros((self.args.prediction_horizon, self.args.output_dim))
			sigmax = np.zeros((self.args.prediction_horizon, 1))
			sigmay = np.zeros((self.args.prediction_horizon, 1))
			if self.args.n_mixtures == 0:
				# plot predicted trajectory global frame
				for i in range(self.args.prediction_horizon):
					idx = i * self.args.output_pred_state_dim
					idy = i * self.args.output_pred_state_dim + 1
					mu_x = model_vel_pred[0][0, idx]
					mu_y = model_vel_pred[0][0, idy]
					vel_pred[i, :] = [mu_x, mu_y]
				if rotate:
					pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, batch_x[step, 2:])
				else:
					pred_vel_global_frame = vel_pred
				traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),
					pred_vel=pred_vel_global_frame, dt=self.args.dt)
				ax_pos_local.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[0], label='Predicted trajectory')
			else:
				for sample_id in range(n_samples):
					prediction_sample = model_vel_pred[sample_id]
					for mix_idx in range(self.args.n_mixtures):
						# plot predicted trajectory global frame
						for pred_step in range(self.args.prediction_horizon):
							idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
							mu_x = prediction_sample[0, idx]
							mu_y = prediction_sample[0, idx + self.args.n_mixtures]
							if self.args.output_pred_state_dim > 2:
								sigmax[pred_step, :] = prediction_sample[0, idx + 2 * self.args.n_mixtures]
								sigmay[pred_step, :] = prediction_sample[0, idx + 3 * self.args.n_mixtures]
							vel_pred[pred_step, :] = [mu_x, mu_y]

						if rotate:
							pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, batch_x[step, 2:])
						else:
							pred_vel_global_frame = vel_pred
						traj_pred = sup.path_from_vel(initial_pos=np.array([0,0]),
							pred_vel=pred_vel_global_frame, dt=self.args.dt)

						ax_pos_local.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx],
							                       label='Predicted trajectory')

						if self.args.output_pred_state_dim > 2:
							# prior of 0 on the uncertainty of the pedestrian velocity
							sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
							sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

							e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
							             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
							e1.set_alpha(0.5)

							for pred_step in range(1, self.args.prediction_horizon):
								sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
								sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt

								e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
								             height=np.sqrt(sigma_y) / 2,
								             angle=0 / np.pi * 180)
								ax_pos_local.add_patch(e1)

		# Plot real trajectory
		gt_vel = np.zeros((self.args.prediction_horizon, self.args.output_dim))
		for i in range(self.args.prediction_horizon):
			idx = i * self.args.output_dim
			idy = i * self.args.output_dim + 1
			mu_x = model_vel_real[step, idx]
			mu_y = model_vel_real[step, idy]
			gt_vel[i, :] = [mu_x, mu_y]
		if False:
			real_vel = sup.rotate_predicted_vel_to_global_frame(gt_vel, batch_x[step, 2:])
		else:
			real_vel = gt_vel
		traj_real = sup.path_from_vel(initial_pos=np.array([0, 0]),
		                              pred_vel=real_vel, dt=self.args.dt)
		ax_pos_local.plot(traj_real[:, 0], traj_real[:, 1], color='m', label='Agent real trajectory')

		# Other agents
		for id in range(other_agents_pos[step].shape[0]):  # number of agents
			ax_pos_local.plot(other_agents_pos[step][id, 0]-x_init_global_frame
			                 , other_agents_pos[step][id, 1]-y_init_global_frame, marker='o', color='b')
		
		# Local grid
		
		if rotate:
			heading = math.atan2(batch_x[step, 3], batch_x[step, 2])
			grid = sup.rotate_grid_around_center(batch_grid[step, :, :], -heading * 180 / math.pi)  # rotation in degrees
			#sup.plot_grid(ax_pos_local, np.array([0.0, 0.0]), grid, self.gridmap.resolution,
			#              np.array([self.args.submap_width, self.args.submap_height]))
		else:
			"""
			if real_vel[0,0] > 0.1:
				sup.plot_grid(ax_pos_local, np.array([3.0, 0.0]), batch_grid[step, :, :], self.gridmap.resolution,
				              np.array([self.args.submap_width, self.args.submap_height]))
			elif real_vel[0,0] < -0.1:
				sup.plot_grid(ax_pos_local, np.array([-3.0, 0.0]), batch_grid[step, :, :], self.gridmap.resolution,
				              np.array([self.args.submap_width, self.args.submap_height]))
			else:
			"""
			#batch_grid[step, :, :] = list(zip(*batch_grid[step, :, :][::-1]))

			#fig, ax = plt.subplots()
			#ax.imshow(batch_grid[step, :, :])
			#plt.show()
			#ax_pos_local.imshow(batch_grid[step,:,:], cmap="gray_r")
			
			#sup.plot_submap(ax_pos_local, self.gridmap.center, batch_grid[step, :, :], self.gridmap.resolution)

			sup.plot_grid(ax_pos_local, self.gridmap.center, batch_grid[step, :, :], self.gridmap.resolution,
		            np.array([self.args.submap_size, self.args.submap_size]))



	def save_global_scenario(self,input_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,rotate=False,n_samples=1,format=".pdf"):
		if not os.path.exists(self.args.model_path + '/results/'+ self.args.scenario+"/figs"):
			os.makedirs(self.args.model_path + '/results/'+ self.args.scenario+"/figs")
		for animation_idx in range(len(input_list)):
			input = input_list[animation_idx]
			grid = grid_list[animation_idx]
			ped_grid = ped_grid_list[animation_idx]

			gt_vel = y_ground_truth_list[animation_idx]
			if not (y_pred_list_global is None):
				model_vel_pred = y_pred_list_global[animation_idx]
			other_agents_pos = other_agents_list[animation_idx]

			for step in range(input.shape[0]):  # trajectory length considering data from getTrajectoryAsBatch
				# Eucldean / position space
				self.ax_global_scenario.clear()
				#self.ax_global_scenario.plot(input[:, 0], input[:, 1], color='b', alpha=0.4, lw=1)
				#self.ax_global_scenario.plot(input[:step, 0], input[:step, 1], color='b', lw=2)
				self.ax_global_scenario.plot(input[step, 0], input[step, 1], color='c', marker='o')
				bottom_left_submap = input[step, 0:2] - np.array([self.args.submap_width / 2.0, self.args.submap_height / 2.0])
				#self.ax_global_scenario.add_patch(
				#	pl.Rectangle(bottom_left_submap, self.args.submap_width, self.args.submap_height, fill=None))

				# Plot real trajectory
				real_vel = np.zeros((1, self.args.prediction_horizon * self.args.output_dim))
				for i in range(self.args.prediction_horizon):
					idx = i * self.args.output_dim
					idy = i * self.args.output_dim + 1
					mu_x = gt_vel[step, idx]
					mu_y = gt_vel[step, idy]
					real_vel[0, 2 * i:2 * i + 2] = [mu_x, mu_y]

				traj_real = sup.path_from_vel(initial_pos=np.array([input[step, 0], input[step, 1]]),
				                              pred_vel=np.squeeze(real_vel), dt=self.args.dt)
				self.ax_global_scenario.plot(traj_real[:, 0], traj_real[:, 1], color='m', label='Agent real trajectory')

				# Other agents
				for id in range(other_agents_pos[step].shape[0]):  # number of agents
					self.ax_global_scenario.plot(other_agents_pos[step][id, 0], other_agents_pos[step][id, 1], marker='o', color='b',markersize=10)
				vel_pred = np.zeros((1, self.args.prediction_horizon * self.args.output_dim))
				# Predicted trajectory
				if not (y_pred_list_global is None):
					if self.args.n_mixtures == 0:
						# plot predicted trajectory global frame
						for i in range(self.args.prediction_horizon):
							idx = i * self.args.output_pred_state_dim
							idy = i * self.args.output_pred_state_dim + 1
							mu_x = model_vel_pred[step, idx]
							mu_y = model_vel_pred[step, idy]
							vel_pred[0, 2 * i:2 * i + 2] = [mu_x, mu_y]
						if rotate:
							pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
						else:
							pred_vel_global_frame = vel_pred
						traj_pred = sup.path_from_vel(initial_pos=np.array(
							[input[step, 0] * self.args.normalization_constant_x,
							 input[step, 1] * self.args.normalization_constant_y]),
							pred_vel=np.squeeze(pred_vel_global_frame), dt=self.args.dt)
						self.ax_global_scenario.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[0], label='Predicted trajectory')
					else:
						for sample_id in range(n_samples):
							for mix_idx in range(self.args.n_mixtures):
								# plot predicted trajectory global frame
								time = np.zeros([self.args.prediction_horizon])
								for pred_step in range(self.args.prediction_horizon):
									time[pred_step] = pred_step * self.args.dt
									idx = i * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id
									idy = i * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx + self.args.n_mixtures + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id
									mu_x = model_vel_pred[step, idx]
									mu_y = model_vel_pred[step, idy]
									vel_pred[0, 2 * i:2 * i + 2] = [mu_x, mu_y]

								if rotate:
									pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
								else:
									pred_vel_global_frame = vel_pred
								traj_pred = sup.path_from_vel(initial_pos=np.array(
									[input[step, 0] * self.args.normalization_constant_x,
									 input[step, 1] * self.args.normalization_constant_y]),
									pred_vel=np.squeeze(pred_vel_global_frame), dt=self.args.dt)
								# sub-sample to make smoother trajectories
								_, traj_pred_smooth, _ = self.smoothenTrajectory(time, traj_pred, vel_pred)


								self.ax_global_scenario.plot(traj_pred_smooth[:, 0], traj_pred_smooth[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')

								# prior of 0 on the uncertainty of the pedestrian velocity
								sigma_x = np.square(
									np.maximum(model_vel_pred[0, self.args.n_mixtures * 2 + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id],0.75)) * self.args.dt * self.args.dt + 0.5
								sigma_y = np.square(
									np.maximum(model_vel_pred[0, self.args.n_mixtures * 3 + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id],0.75)) * self.args.dt * self.args.dt +0.5

								e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
								             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
								e1.set_alpha(0.5)
								self.ax_global_scenario.add_patch(e1)

								for i in range(1, self.args.prediction_horizon):
									idx = i * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures * 2 + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id
									idy = i * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures * 3 + mix_idx + self.args.output_pred_state_dim * self.args.n_mixtures*self.args.prediction_horizon*sample_id
									sigma_x += np.square(
										np.maximum(model_vel_pred[i - 1, idx],0.75)) * self.args.dt * self.args.dt
									sigma_y += np.square(
										np.maximum(model_vel_pred[i - 1, idy],0.75)) * self.args.dt * self.args.dt
									e1 = Ellipse(xy=(traj_pred[i, 0], traj_pred[i, 1]), width=np.sqrt(sigma_x) / 2,
									             height=np.sqrt(sigma_y) / 2,
									             angle=0 / np.pi * 180)
									e1.set_alpha(0.5)
									self.ax_global_scenario.add_patch(e1)

				sup.plot_grid(self.ax_global_scenario, np.array([0.0, 0.0]), self.gridmap.gridmap, self.gridmap.resolution,
				              self.gridmap.map_size)
				#self.ax_global_scenario.set_xlim([-self.gridmap.center[0], self.gridmap.center[0]])
				#self.ax_global_scenario.set_ylim([-self.gridmap.center[1], self.gridmap.center[1]])
				self.ax_global_scenario.set_xlim([input[step, 0]-7, input[step, 0]+7])
				self.ax_global_scenario.set_ylim([input[step, 1]-7, input[step, 1]+7])

				#self.ax_global_scenario.set_xlabel('x [m]')
				#self.ax_global_scenario.set_ylabel('y [m]')
				self.ax_global_scenario.axis('off')
				self.ax_global_scenario.set_aspect('equal')
				self.fig_global_scenario.savefig(self.args.model_path + '/results/'+ self.args.scenario+"/figs/"+"scene"+str(animation_idx)+"_step_"+str(step)+format, bbox_inches='tight')

	def smoothenTrajectory(self, time_vec,pose_vec,vel_vec,dt=0.05):
		"""
		Cubic interpolation with provided values in order to smoothen the trajectory and obtain
		a sequence with the specified dt.
		"""
		x_interpolator = interpolate.interp1d(time_vec, pose_vec[:, 0], kind='cubic', axis=0,fill_value='extrapolate')
		y_interpolator = interpolate.interp1d(time_vec, pose_vec[:, 1], kind='cubic', axis=0,fill_value='extrapolate')
		vx_interpolator = interpolate.interp1d(time_vec, vel_vec[:, 0], kind='cubic', axis=0,fill_value='extrapolate')
		vy_interpolator = interpolate.interp1d(time_vec, vel_vec[:, 1], kind='cubic', axis=0,fill_value='extrapolate')
		n_elem = self.args.prediction_horizon*int(self.args.dt/dt)
		new_time_vec = np.linspace(0, 0 + (n_elem - 1) * dt, n_elem)

		new_time_vec = np.zeros([n_elem])
		new_pose_vec = np.zeros([n_elem, 2])
		new_vel_vec = np.zeros([n_elem, 2])

		for ii in range(n_elem):
			t = 0 + ii * dt
			new_time_vec[ii] = t
			new_pose_vec[ii, 0] = x_interpolator(t)
			new_pose_vec[ii, 1] = y_interpolator(t)
			new_vel_vec[ii, 0] = vx_interpolator(t)
			new_vel_vec[ii, 1] = vy_interpolator(t)

		return new_time_vec, new_pose_vec, new_vel_vec

	def plot_on_video(self,input_list,grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,traj_list,all_traj_likelihood,test_args,social_trajectories=None):
		homography_file = os.path.join( self.args.data_path+ self.args.scenario, 'H.txt')
		if os.path.exists(homography_file):
			Hinv = np.linalg.inv(np.loadtxt(homography_file))
		else:
			print('[INF] No homography file')
		resolution = np.abs(1/Hinv[1,0])
		print("Resolution: " + str(resolution))
		scenario = self.args.scenario.split('/')[-1]
		video_file = self.args.data_path + self.args.scenario+'/' + scenario + '.avi'
		if os.path.exists(video_file):
			print('[INF] Using video file ' + video_file)
			cap = cv2.VideoCapture(video_file)
			if "zara_02" in self.args.scenario:
				time_offset = -22
			else:
				time_offset = 0

			if cap:
				dt = 0.4  # seconds (equivalent to 2.5 fps)
				#if 'seq_eth' in self.args.scenario:
				#	frames_between_annotation = 6.0
				#else:
				frames_between_annotation = 10.0

				frame_width = int(cap.get(3))
				frame_height = int(cap.get(4))

				# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
				if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario):
					os.makedirs(self.args.model_path + '/results/' + self.args.scenario)
				if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario+"/figs"):
					os.makedirs(self.args.model_path + '/results/' + self.args.scenario+"/figs")
				if social_trajectories:
					out = cv2.VideoWriter(self.args.model_path + '/results/' + self.args.scenario+'/outpy'+str(test_args.n_samples)+'_social.avi', cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (frame_width, frame_height))
				else:
					out = cv2.VideoWriter(self.args.model_path + '/results/' + self.args.scenario+'/outpy'+str(test_args.n_samples)+'.avi', cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (frame_width, frame_height))

			vel_real = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			sigmax = np.zeros((self.args.prediction_horizon, 1))
			sigmay = np.zeros((self.args.prediction_horizon, 1))
			pis = np.zeros((self.args.prediction_horizon, 3))
			for animation_idx in range(len(input_list)):
				traj = traj_list[animation_idx]

				if social_trajectories:
					social_traj = social_trajectories[animation_idx]
				gt_vel = y_ground_truth_list[animation_idx]
				input = input_list[animation_idx]
				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				traj_likelihood = all_traj_likelihood[animation_idx]
				# plot real trajectory global frame
				for step in range(input.shape[0]):

					time_stamp = traj.time_vec[step] / dt * frames_between_annotation + time_offset
					cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, time_stamp +self.args.prev_horizon* frames_between_annotation))
					ret, im = cap.read()
					overlay = im.copy()

					overlay_ellipses = im.copy()
					if not ret:
						print("Could not open the video")

					# Initial positions
					if self.args.normalize_data:
						x0 = input[step, 0] / self.args.sx_pos + self.args.min_pos_x
						y0 = input[step, 1] / self.args.sy_pos + self.args.min_pos_y
					else:
						x0 = input[step, 0]
						y0 = input[step, 1]

					# Real Trajectory
					for i in range(self.args.prediction_horizon):
						idx = i * self.args.output_dim
						idy = i * self.args.output_dim + 1
						mu_x = gt_vel[step, idx]
						mu_y = gt_vel[step, idy]
						vel_real[i, :] = [mu_x, mu_y]
					if self.args.rotated_grid:
						real_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_real, input[step, 2:])
					else:
						real_vel_global_frame = vel_real
					traj_real = sup.path_from_vel(initial_pos=np.array([x0,
					                                                    y0]),
					                              pred_vel=real_vel_global_frame, dt=self.args.dt)

					obsv_XY = sup.to_image_frame(Hinv, traj_real)
					sup.line_cv(overlay, obsv_XY, (255, 0, 0), 3) # bgr convention

					# Plot social trajectory
					if social_trajectories:
						for sample_id in range(3):
							social_traj_pred = sup.path_from_vel(initial_pos=np.array([input[step, 0],
						                                                    input[step, 1]]),
						                              pred_vel=social_traj[step][sample_id], dt=1)
							obsv_XY = sup.to_image_frame(Hinv, social_traj_pred)
							sup.line_cv(overlay, obsv_XY, (255, 0, 0), 3)  # bgr convention

					# Plot real predicted traj from positions
					traj_real = traj.pose_vec[step:step+self.args.prev_horizon+self.args.prediction_horizon,:2]
					obsv_XY = sup.to_image_frame(Hinv, traj_real)
					sup.line_cv(overlay, obsv_XY, (0, 0, 0), 3)  # bgr convention

					# Predicted trajectory
					colors = [(0,0,255),(0,255,0),(0,255,255)]
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							prediction_sample = model_vel_pred[step][0]
							# plot predicted trajectory global frame
							time = np.zeros([self.args.prediction_horizon])
							for i in range(self.args.prediction_horizon):
								time[i] = i*self.args.dt
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								if self.args.normalize_data:
									mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
									mu_y = prediction_sample[0, idy] / self.args.sy_vel + self.args.min_vel_y
								else:
									mu_x = prediction_sample[0, idx]
									mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]
							if self.args.rotated_grid:
								pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
							else:
								pred_vel_global_frame = vel_pred

							if self.args.predict_positions:
								traj_pred = vel_pred
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array(
									[x0,
									 y0]),
									pred_vel=pred_vel_global_frame, dt=self.args.dt)
							# sub-sample to make smoother trajectories
							_, traj_pred_smooth, _ = self.smoothenTrajectory(time,traj_pred,vel_pred)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred_smooth)
							sup.line_cv(overlay, obsv_XY, (0, 0, 255),2)
						else:
							for sample_id in range(test_args.n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								for mix_idx in range(self.args.n_mixtures): #self.args.n_mixtures
									# plot predicted trajectory global frame
									time = np.zeros([self.args.prediction_horizon])
									for pred_step in range(self.args.prediction_horizon):
										time[pred_step] = pred_step * self.args.dt
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										if self.args.normalize_data:
											mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
											mu_y = prediction_sample[0, idx + self.args.n_mixtures] / self.args.sy_vel + self.args.min_vel_y
										else:
											mu_x = prediction_sample[0, idx]
											mu_y = prediction_sample[0, idx + self.args.n_mixtures]
										if self.args.output_pred_state_dim > 2:
											sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
											sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										if self.args.output_pred_state_dim ==5:
											pis[pred_step, mix_idx] = prediction_sample[0][idx + 4 * self.args.n_mixtures]
										if math.isnan(mu_x) | math.isnan(mu_y):
											continue
										vel_pred[pred_step, :] = [mu_x, mu_y]

									if self.args.rotated_grid:
										pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
									else:
										pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array(
										[x0,
										 y0]),
										pred_vel=pred_vel_global_frame, dt=self.args.dt)
									# sub-sample to make smoother trajectories
									self.smoothenTrajectory(time, traj_pred, vel_pred)
									obsv_XY = sup.to_image_frame(Hinv, traj_pred)

									sup.line_cv(overlay, obsv_XY, colors[mix_idx], 2)

									if self.args.output_pred_state_dim > 2:
										# prior of 0 on the uncertainty of the pedestrian velocity
										sigma_x = np.square(min(sigmax[0],0.5)) * self.args.dt * self.args.dt +0.2
										sigma_y = np.square(min(sigmay[0],0.5)) * self.args.dt * self.args.dt +0.2
										if math.isnan(sigma_x) | math.isnan(sigma_y):
											continue
										axis = (int(np.sqrt(sigma_x) /resolution),int(np.sqrt(sigma_y) /resolution))
										center = (obsv_XY[0, 1], obsv_XY[0, 0])
										cv2.ellipse(overlay_ellipses, center,       axis,    0, 0, 360, (255,153,51),-1)

										for pred_step in range(1, self.args.prediction_horizon):
											sigma_x += np.square(min(sigmax[pred_step],0.5)) * self.args.dt * self.args.dt
											sigma_y += np.square(min(sigmay[pred_step],0.5)) * self.args.dt * self.args.dt
											if math.isnan(sigma_x) | math.isnan(sigma_y):
												continue
											axis = (int(np.sqrt(sigma_x) / resolution),
											        int(np.sqrt(sigma_y) / resolution))
											if (sigma_x<0) or (sigma_y<0):
												print("Negative ellipse")
											center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
											cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (255,153,51),-1)

					# Plot other agents
					plot_other_agents = True
					if plot_other_agents:
						other_agents = other_agents_list[animation_idx][step]
						axis = (int(np.sqrt(0.3) / resolution),
						        int(np.sqrt(0.3) / resolution))

						for agent_id in range(other_agents.shape[0]):
							obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(other_agents[agent_id,:2],axis=0))
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (153, 153, 51), -1)

					image_new = cv2.addWeighted(overlay, 0.6, im, 0.4, 0)
					image_new = cv2.addWeighted(image_new, 0.6, overlay_ellipses, 0.4, 0)

					# Adding legend
					font = cv2.FONT_HERSHEY_SIMPLEX
					cv2.line(image_new, (frame_width - 270, 30), (frame_width - 250, 30), (0, 0, 0), 4)
					cv2.putText(image_new, "Real Trajectory", (frame_width - 240, 30), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)
					cv2.line(image_new, (frame_width - 270, 50), (frame_width - 250, 50), (0, 0, 255), 4)
					if self.args.output_pred_state_dim == 5:
						cv2.putText(image_new, "Predicted Trajectory Prob " + str(np.sum(pis[:,0])/self.args.prediction_horizon*100), (frame_width - 240, 50), font, 0.5, (0, 0, 255), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 70), (frame_width - 250, 70), (0, 255, 0), 4)
						cv2.putText(image_new, "Predicted Trajectory Prob " + str(np.sum(pis[:,1])/self.args.prediction_horizon*100), (frame_width - 240, 70), font, 0.5, (0, 255, 0), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 90), (frame_width - 250, 90), (0, 255, 255), 4)
						cv2.putText(image_new, "Predicted Trajectory Prob " + str(np.sum(pis[:,2])/self.args.prediction_horizon*100), (frame_width - 240, 90), font, 0.5, (0, 255, 255), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 110), (frame_width - 250, 110), (255, 153, 51), 4)
					else:
						cv2.putText(image_new, "Predicted Trajectory: " + str(np.around(traj_likelihood[step][0,0,0],decimals=2)), (frame_width - 240, 50), font, 0.5, (0, 0, 255), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 70), (frame_width - 250, 70), (0, 255, 0), 4)
						cv2.putText(image_new, "Predicted Trajectory: " + str(np.around(traj_likelihood[step][0,0,1],decimals=2)), (frame_width - 240, 70), font, 0.5, (0, 255, 0), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 90), (frame_width - 250, 90), (0, 255, 255), 4)
						cv2.putText(image_new, "Predicted Trajectory: " + str(np.around(traj_likelihood[step][0,0,2],decimals=2)), (frame_width - 240, 90), font, 0.5, (0, 255, 255), 2,
						            cv2.LINE_AA)
						cv2.line(image_new, (frame_width - 270, 110), (frame_width - 250, 110), (255, 153, 51), 4)
					cv2.putText(image_new, "1-sigma Uncertainty", (frame_width - 240, 110), font, 0.5, (255,153,51), 2,
					            cv2.LINE_AA)

					out.write(image_new)

					#cv2.imwrite(self.args.model_path + '/results/' + self.args.scenario+"/figs/result_"+str(animation_idx)+"_"+str(step)+".jpg", image_new);

			# When everything done, release the video capture and video write objects
			cap.release()
			out.release()

			# Closes all the frames
			cv2.destroyAllWindows()

		else:
			"Video not found"

	def plot_on_image(self,input_list, grid_list, y_pred_list_global, y_ground_truth_list, other_agents_list, traj_list, test_args, social_trajectories=None):
		"""
		Not working yet for Herengracht map. Seems to be something wrong with the H.txt file and size of plotting
		"""

		scenario = self.args.scenario.split('/')[-1]

		if not os.path.exists(self.args.model_path + '/../videos/'):
			os.makedirs(self.args.model_path + '/../videos/')
		if test_args.freeze_other_agents:
			video_file = self.args.model_path + '/../videos/'+scenario + str(self.args.exp_num) + "_frozen.avi"
			map_file = self.args.data_path + self.args.scenario + '/map.png'
			homography_file = os.path.join(self.args.data_path + self.args.scenario, 'H.txt')
		elif test_args.unit_testing:
			video_file = self.args.model_path + '/../videos/'+scenario + str(self.args.exp_num) + "_unit_tests.avi"
			map_file = self.args.data_path + self.args.scenario + '/unit_map.png'
			homography_file = os.path.join(self.args.data_path + self.args.scenario, 'unit_H.txt')
		else:
			video_file = self.args.model_path + '/../videos/'+scenario + str(self.args.exp_num) + "_final.avi"
			map_file = self.args.data_path + self.args.scenario + '/map.png'
			homography_file = os.path.join(self.args.data_path + self.args.scenario, 'H.txt')

		if os.path.exists(homography_file):
			Hinv = np.linalg.inv(np.loadtxt(homography_file))
		else:
			print('[INF] No homography file')
		resolution = self.args.submap_resolution
		print("Resolution: " + str(resolution))

		im = np.uint8(cv2.imread(map_file) * -1 + 255)
		#im = cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE)
		frame_height, frame_width, layers = im.shape
		scale_factor = 2
		resized_img = cv2.resize(im, (frame_width * scale_factor, frame_height * scale_factor),
		                         interpolation=cv2.INTER_AREA)
		print(resized_img.shape)

		h, w, c = resized_img.shape
		blank_image = 255 * np.ones(shape=(h, w, c), dtype=np.uint8)


		if os.path.exists(map_file):
			print('[INF] Using video file ' + video_file)
			cap = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 20,
			                      (frame_width * scale_factor, frame_height * scale_factor))

			vel_real = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			sigmax = np.zeros((self.args.prediction_horizon, 1))
			sigmay = np.zeros((self.args.prediction_horizon, 1))
			pis = np.zeros((self.args.prediction_horizon, 3))
			for animation_idx in range(len(input_list)):
				traj = traj_list[animation_idx]

				if social_trajectories:
					social_traj = social_trajectories[animation_idx]
				gt_vel = y_ground_truth_list[animation_idx]
				input = input_list[animation_idx]
				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				# plot real trajectory global frame
				for step in range(input.shape[0]):

					overlay = resized_img.copy()
					#overlay = blank_image.copy()

					# Initial positions
					if self.args.normalize_data:
						x0 = input[step, 0] / self.args.sx_pos + self.args.min_pos_x
						y0 = input[step, 1] / self.args.sy_pos + self.args.min_pos_y
					else:
						x0 = input[step, 0]
						y0 = input[step, 1]

					# Real Trajectory
					for i in range(self.args.prediction_horizon):
						idx = i * self.args.output_dim
						idy = i * self.args.output_dim + 1
						mu_x = gt_vel[step, idx] + input[step, 2]
						mu_y = gt_vel[step, idy] + input[step, 3]
						vel_real[i, :] = [mu_x, mu_y]

					real_vel_global_frame = vel_real
					traj_real = sup.path_from_vel(initial_pos=np.array([x0,
					                                                    y0]),
					                              pred_vel=real_vel_global_frame, dt=self.args.dt)

					obsv_XY = sup.to_image_frame(Hinv, traj_real)*scale_factor
					sup.line_cv(overlay, obsv_XY, (255, 0, 0), 3) # bgr convention

					# Plot social trajectory
					if social_trajectories:
						for sample_id in range(3):
							social_traj_pred = sup.path_from_vel(initial_pos=np.array([input[step, 0],
						                                                    input[step, 1]]),
						                              pred_vel=social_traj[step][sample_id], dt=1)
							obsv_XY = sup.to_image_frame(Hinv, social_traj_pred)*scale_factor
							sup.line_cv(overlay, obsv_XY, (255, 0, 0), 3)  # bgr convention

					# Plot real predicted traj from positions
					#traj_real = traj.pose_vec[step:step+self.args.prev_horizon+self.args.prediction_horizon,:2]
					#obsv_XY = sup.to_image_frame(Hinv, traj_real)*scale_factor
					#sup.line_cv(overlay, obsv_XY, (0, 0, 0), 3)  # bgr convention

					# Predicted trajectory
					colors = [(0,0,255),(0,255,0),(0,255,255)]
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							prediction_sample = model_vel_pred[step][0]
							# plot predicted trajectory global frame
							time = np.zeros([self.args.prediction_horizon])
							for i in range(self.args.prediction_horizon):
								time[i] = i*self.args.dt
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								if self.args.normalize_data:
									mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
									mu_y = prediction_sample[0, idy] / self.args.sy_vel + self.args.min_vel_y
								else:
									mu_x = prediction_sample[0, idx]
									mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]

							pred_vel_global_frame = vel_pred

							if self.args.predict_positions:
								traj_pred = vel_pred
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array(
									[x0,
									 y0]),
									pred_vel=pred_vel_global_frame, dt=self.args.dt)
							# sub-sample to make smoother trajectories
							#_, traj_pred_smooth, _ = self.smoothenTrajectory(time,traj_pred,vel_pred)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)*scale_factor
							sup.line_cv(overlay, obsv_XY, (0, 0, 255),2)
							#print(obsv_XY)
						else:
							for sample_id in range(test_args.n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								for mix_idx in range(self.args.n_mixtures): #self.args.n_mixtures
									# plot predicted trajectory global frame
									time = np.zeros([self.args.prediction_horizon])
									for pred_step in range(self.args.prediction_horizon):
										time[pred_step] = pred_step * self.args.dt
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										if self.args.normalize_data:
											mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
											mu_y = prediction_sample[0, idx + self.args.n_mixtures] / self.args.sy_vel + self.args.min_vel_y
										else:
											mu_x = prediction_sample[0, idx] + input[step, 2]
											mu_y = prediction_sample[0, idx + self.args.n_mixtures] + input[step, 3]
										if self.args.output_pred_state_dim > 2:
											sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
											sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										if self.args.output_pred_state_dim ==5:
											pis[pred_step, mix_idx] = prediction_sample[0][idx + 4 * self.args.n_mixtures]
										if math.isnan(mu_x) | math.isnan(mu_y):
											continue
										vel_pred[pred_step, :] = [mu_x, mu_y]

									pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array(
										[x0,
										 y0]),
										pred_vel=pred_vel_global_frame, dt=self.args.dt)
									# sub-sample to make smoother trajectories
									#self.smoothenTrajectory(time, traj_pred, vel_pred)
									obsv_XY = sup.to_image_frame(Hinv, traj_pred)*scale_factor

									sup.line_cv(overlay, obsv_XY, colors[mix_idx], 2)
									print(obsv_XY)

									try:
										if self.args.output_pred_state_dim > 2:
											# prior of 0 on the uncertainty of the pedestrian velocity
											sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt +0.2
											sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt +0.2
											if math.isnan(sigma_x) | math.isnan(sigma_y):
												continue
											axis = (int(np.sqrt(sigma_x) /resolution*scale_factor),int(np.sqrt(sigma_y) /resolution*scale_factor))
											y = obsv_XY[0, 1]*scale_factor
											x = obsv_XY[0, 0]*scale_factor
											center = (y , x)
											cv2.ellipse(overlay, center,       axis,    0, 0, 360, (255,153,51),-1)

											for pred_step in range(1, self.args.prediction_horizon):
												sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
												sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt
												if math.isnan(sigma_x) | math.isnan(sigma_y):
													continue
												axis = (int(np.sqrt(sigma_x) / resolution*scale_factor),
												        int(np.sqrt(sigma_y) / resolution*scale_factor))
												if (sigma_x<0) or (sigma_y<0):
													print("Negative ellipse")
												y = obsv_XY[pred_step, 1] * scale_factor
												x = obsv_XY[pred_step, 0] * scale_factor
												center = (y, x)
												cv2.ellipse(overlay, center, axis, 0, 0, 360, (255,153,51),-1)
									except:
										print("Failed to add ellipse")

					# Plot other agents
					plot_other_agents = True
					if plot_other_agents:
						other_agents = other_agents_list[animation_idx][step]
						axis = (int(np.sqrt(0.3) / resolution),
						        int(np.sqrt(0.3) / resolution))

						for agent_id in range(other_agents.shape[0]):
							obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(other_agents[agent_id,:2],axis=0))
							center = (int(obsv_XY[0, 1]*scale_factor), int(obsv_XY[0, 0]*scale_factor))
							cv2.ellipse(overlay, center, axis, 0, 0, 360, (153, 153, 51), -1)

					# Adding legend
					font = cv2.FONT_HERSHEY_SIMPLEX
					cv2.line(overlay, (int(frame_height*scale_factor*0.9-50), 30), (int(frame_height*scale_factor*0.9-40), 30), (255, 0, 0), 4)
					cv2.putText(overlay, "Real Trajectory", (int(frame_height*scale_factor*0.9-30), 30), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
					cv2.line(overlay, (int(frame_height*scale_factor*0.9-50), 50), (int(frame_height*scale_factor*0.9-40), 50), (0, 0, 255), 4)
					cv2.putText(overlay, "Predicted Trajectory", (int(frame_height*scale_factor * 0.9 - 40), 50), font, 0.5, (0, 0, 255), 2,
					            cv2.LINE_AA)

					cap.write(overlay)

					#cv2.imwrite(self.args.model_path + '/results/' + self.args.scenario+"/figs/result_"+str(animation_idx)+"_"+str(step)+".jpg", overlay);

			# When everything done, release the video capture and video write objects
			cap.release()

			# Closes all the frames
			cv2.destroyAllWindows()

		else:
			print("map_file not found")

	def plot_GA3C(self,input_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,traj_list,rotate=False,n_samples=1,social_trajectories=None):
		homography_file = os.path.join( self.args.data_path+ self.args.scenario, 'H.txt')
		if os.path.exists(homography_file):
			Hinv = np.linalg.inv(np.loadtxt(homography_file))
		else:
			print('[INF] No homography file')
		resolution = np.abs(1/Hinv[0,0])/2
		print("Resolution: " + str(resolution))
		scenario = self.args.scenario.split('/')[-1]
		map_file = self.args.data_path + self.args.scenario+ '/map.png'
		im = cv2.imread(map_file) +255
		height, width, layers = im.shape
		scale_factor = 2
		resized_img = cv2.resize(im, (height*scale_factor, width*scale_factor), interpolation=cv2.INTER_AREA)

		video_file = self.args.model_path + '/results/' + self.args.scenario+'/outpy'+'.avi'
		if os.path.exists(map_file):
			print('[INF] Using map file ' + map_file)
			cap = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width*scale_factor, height*scale_factor))


			vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
			sigmax = np.zeros((self.args.prediction_horizon, 1))
			sigmay = np.zeros((self.args.prediction_horizon, 1))
			pis = np.zeros((self.args.prediction_horizon, 3))

			for animation_idx in range(len(input_list)):
				traj = traj_list[animation_idx]

				gt_vel = y_ground_truth_list[animation_idx]
				input = input_list[animation_idx]
				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None

				# Real Trajectory
				traj_real = np.zeros((len(traj), self.args.output_dim))
				for i in range(len(traj)):
					x = traj[i]["pedestrian_state"]["position"][0]
					y = traj[i]["pedestrian_state"]["position"][1]
					traj_real[i, :] = [x, y]

				# plot real trajectory global frame
				for step in range(len(model_vel_pred)):

					overlay = resized_img.copy()

					# Initial positions
					x0 = traj[0]["pedestrian_state"]["position"][0]
					y0 = traj[0]["pedestrian_state"]["position"][1]
					obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(np.array([x0,y0]),axis=0))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					axis = (int(np.sqrt(0.3) / resolution),
					        int(np.sqrt(0.3) / resolution))
					cv2.ellipse(overlay, center, axis, 0, 0, 360, (255, 0, 0), -1)

					# Final positions
					x0 = traj[-1]["pedestrian_state"]["position"][0]
					y0 = traj[-1]["pedestrian_state"]["position"][1]
					obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(np.array([x0, y0]), axis=0))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					axis = (int(np.sqrt(0.3) / resolution),
					        int(np.sqrt(0.3) / resolution))
					cv2.ellipse(overlay, center, axis, 0, 0, 360, (0, 255, 0), -1)

					obsv_XY = sup.to_image_frame(Hinv, traj_real)
					sup.line_cv(overlay, obsv_XY, (0, 0, 0), 3) # bgr convention

					# Predicted trajectory
					colors = [(0,0,255),(0,255,0),(0,255,255)]
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							prediction_sample = model_vel_pred[step][0]
							# plot predicted trajectory global frame
							time = np.zeros([self.args.prediction_horizon])
							for i in range(self.args.prediction_horizon):
								time[i] = i*self.args.dt
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								if self.args.normalize_data:
									mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
									mu_y = prediction_sample[0, idy] / self.args.sy_vel + self.args.min_vel_y
								else:
									mu_x = prediction_sample[0, idx]
									mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]

							if True: #self.args.predict_positions
								traj_pred = vel_pred
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array(
									[x0,
									 y0]),
									pred_vel=pred_vel_global_frame, dt=self.args.dt)
							# sub-sample to make smoother trajectories
							#_, traj_pred_smooth, _ = self.smoothenTrajectory(time,traj_pred,vel_pred)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)
							sup.line_cv(overlay, obsv_XY, (0, 0, 255),2)
						else:
							for sample_id in range(n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								for mix_idx in range(self.args.n_mixtures): #self.args.n_mixtures
									# plot predicted trajectory global frame
									time = np.zeros([self.args.prediction_horizon])
									for pred_step in range(self.args.prediction_horizon):
										time[pred_step] = pred_step * self.args.dt
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										if self.args.normalize_data:
											mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
											mu_y = prediction_sample[0, idx + self.args.n_mixtures] / self.args.sy_vel + self.args.min_vel_y
										else:
											mu_x = prediction_sample[0, idx]
											mu_y = prediction_sample[0, idx + self.args.n_mixtures]
										sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
										sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										if self.args.output_pred_state_dim ==5:
											pis[pred_step, mix_idx] = prediction_sample[0][idx + 4 * self.args.n_mixtures]
										if math.isnan(mu_x) | math.isnan(mu_y):
											continue
										vel_pred[pred_step, :] = [mu_x, mu_y]

									if rotate:
										pred_vel_global_frame = sup.rotate_predicted_vel_to_global_frame(vel_pred, input[step, 2:])
									else:
										pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array(
										[x0,
										 y0]),
										pred_vel=pred_vel_global_frame, dt=self.args.dt)
									# sub-sample to make smoother trajectories
									#self.smoothenTrajectory(time, traj_pred, vel_pred)
									obsv_XY = sup.to_image_frame(Hinv, traj_pred)

									sup.line_cv(overlay, obsv_XY, colors[mix_idx], 2)

									# prior of 0 on the uncertainty of the pedestrian velocity
									sigma_x = np.square(min(sigmax[0],0.5)) * self.args.dt * self.args.dt +0.2
									sigma_y = np.square(min(sigmay[0],0.5)) * self.args.dt * self.args.dt +0.2
									if math.isnan(sigma_x) | math.isnan(sigma_y):
										continue
									axis = (int(np.sqrt(sigma_x) /resolution),int(np.sqrt(sigma_y) /resolution))
									center = (obsv_XY[0, 1], obsv_XY[0, 0])
									cv2.ellipse(overlay, center,       axis,    0, 0, 360, (255,153,51),-1)

									for pred_step in range(1, self.args.prediction_horizon):
										sigma_x += np.square(min(sigmax[pred_step],0.5)) * self.args.dt * self.args.dt
										sigma_y += np.square(min(sigmay[pred_step],0.5)) * self.args.dt * self.args.dt
										if math.isnan(sigma_x) | math.isnan(sigma_y):
											continue
										axis = (int(np.sqrt(sigma_x) / resolution),
										        int(np.sqrt(sigma_y) / resolution))
										if (sigma_x<0) or (sigma_y<0):
											print("Negative ellipse")
										center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
										cv2.ellipse(overlay, center, axis, 0, 0, 360, (255,153,51),-1)

					# Plot other agents
					plot_other_agents = True
					if plot_other_agents:
						other_agents = np.expand_dims(other_agents_list[animation_idx][0,step],axis=0)
						axis = (int(np.sqrt(0.3) / resolution),
						        int(np.sqrt(0.3) / resolution))

						for agent_id in range(other_agents.shape[0]):
							obsv_XY = sup.to_image_frame(Hinv, np.expand_dims(other_agents[agent_id],axis=0))
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							cv2.ellipse(overlay, center, axis, 0, 0, 360, (153, 153, 51), -1)

					# Adding legend
					font = cv2.FONT_HERSHEY_SIMPLEX
					cv2.line(overlay, (width - 70, 30), (width - 50, 30), (0, 0, 0), 4)
					cv2.putText(overlay, "Real Trajectory", (width - 40, 30), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)
					cv2.line(overlay, (width - 70, 50), (width - 50, 50), (0, 0, 255), 4)
					if self.args.output_pred_state_dim == 5:
						cv2.putText(overlay, "Predicted Trajectory Prob " + str(np.sum(pis[:,0])/self.args.prediction_horizon*100), (width - 40, 50), font, 0.5, (0, 0, 255), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 70), (width - 50, 70), (0, 255, 0), 4)
						cv2.putText(overlay, "Predicted Trajectory Prob " + str(np.sum(pis[:,1])/self.args.prediction_horizon*100), (width - 40, 70), font, 0.5, (0, 255, 0), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 90), (width - 50, 90), (0, 255, 255), 4)
						cv2.putText(overlay, "Predicted Trajectory Prob " + str(np.sum(pis[:,2])/self.args.prediction_horizon*100), (width - 40, 90), font, 0.5, (0, 255, 255), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 110), (width - 50, 110), (255, 153, 51), 4)
					else:
						cv2.putText(overlay, "Predicted Trajectory Mode 0", (width - 40, 50), font, 0.5, (0, 0, 255), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 70), (width - 50, 70), (0, 255, 0), 4)
						cv2.putText(overlay, "Goal Position", (width - 40, 70), font, 0.5, (0, 255, 0), 2,
						            cv2.LINE_AA)
						cv2.line(overlay, (width - 70, 90), (width - 50, 90), (255, 0, 0), 4)
						cv2.putText(overlay, "Initial", (width - 40, 90), font, 0.5, (255, 0, 0), 2,
						            cv2.LINE_AA)

					cap.write(overlay)

					cv2.imwrite(self.args.model_path + '/results/' + self.args.scenario+"/figs/result_"+str(animation_idx)+"_"+str(step)+".jpg", overlay);

			# When everything done, release the video capture and video write objects
			cap.release()

			# Closes all the frames
			cv2.destroyAllWindows()

		else:
			print("Video not found")

	def animate_GA3C(self,input_list,goal_list,grid_list,ped_grid_list,y_pred_list_global,y_ground_truth_list,other_agents_list,traj_list,test_args):
		#self.args.dt = 0.1
		pl.show(block=False)
		if not os.path.exists(self.args.model_path + '/results/'+ self.args.scenario):
			os.makedirs(self.args.model_path + '/results/'+ self.args.scenario)
		if test_args.freeze_other_agents:
			video_file = self.args.model_path + '/results/' + self.args.scenario+"/frozen.mp4"
		else:
			video_file = self.args.model_path + '/results/' + self.args.scenario + "/final.mp4"
		with self.writer.saving(self.fig_ga3c, video_file, 100):
			for animation_idx in range(len(input_list)):
				input = input_list[animation_idx]
				goal = goal_list[animation_idx]
				traj = traj_list[animation_idx]
				if not (y_pred_list_global is None):
					model_vel_pred = y_pred_list_global[animation_idx]
				else:
					model_vel_pred = None
				other_agents_pos = other_agents_list[animation_idx]

				# Real Trajectory
				traj_real = np.zeros((len(traj), self.args.output_dim))
				for i in range(len(traj)):
					x = traj[i]["pedestrian_state"]["position"][0]
					y = traj[i]["pedestrian_state"]["position"][1]
					traj_real[i, :] = [x, y]

				for step in range(input.shape[0]): # trajectory length considering data from getTrajectoryAsBatch
					# Eucldean / position space
					self.ga3c_plot.clear()
					#self.ga3c_plot.plot(input[:, 0], input[:, 1], color='c', alpha=0.4, lw=1)
					#self.ga3c_plot.plot(input[:step, 0], input[:step, 1], color='b', lw=2)
					# Initial positions
					x0 = traj[0]["pedestrian_state"]["position"][0]
					y0 = traj[0]["pedestrian_state"]["position"][1]
					self.ga3c_plot.plot(x0, y0, color='r', marker='o')

					xf = traj[-1]["pedestrian_state"]["position"][0]
					yf = traj[-1]["pedestrian_state"]["position"][1]
					self.ga3c_plot.plot(xf, yf, color='g', marker='x')

					xf = goal[0,0]
					yf = goal[0,1]
					self.ga3c_plot.plot(xf, yf, color='g', marker='o')

					# Other agents
					for id in range(other_agents_pos.shape[0]): # number of agents
						self.ga3c_plot.plot(other_agents_pos[id][step, 0], other_agents_pos[id][step, 1], marker='o', color='b')
					vel_pred = np.zeros((self.args.prediction_horizon,  self.args.output_dim))
					sigmax = np.zeros((self.args.prediction_horizon, 1))
					sigmay = np.zeros((self.args.prediction_horizon, 1))
					#Predicted trajectory
					if not (y_pred_list_global is None):
						if self.args.n_mixtures == 0:
							prediction_sample = model_vel_pred[step][0]
							# plot predicted trajectory global frame
							time = np.zeros([self.args.prediction_horizon])
							for i in range(self.args.prediction_horizon):
								time[i] = i*self.args.dt
								idx = i * self.args.output_pred_state_dim
								idy = i * self.args.output_pred_state_dim + 1
								if self.args.normalize_data:
									mu_x = prediction_sample[0, idx] / self.args.sx_vel + self.args.min_vel_x
									mu_y = prediction_sample[0, idy] / self.args.sy_vel + self.args.min_vel_y
								else:
									mu_x = prediction_sample[0, idx]
									mu_y = prediction_sample[0, idy]
								vel_pred[i, :] = [mu_x, mu_y]
							x0 = traj[step]["pedestrian_state"]["position"][0]
							y0 = traj[step]["pedestrian_state"]["position"][1]
							if self.args.predict_positions:
								traj_pred = vel_pred
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array([x0,y0]),
									pred_vel=vel_pred, dt=self.args.dt)
							self.ga3c_plot.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[0], label='Predicted trajectory')
						else:
							for sample_id in range(test_args.n_samples):
								prediction_sample = model_vel_pred[step][sample_id]
								x0 = traj[step]["pedestrian_state"]["position"][0]
								y0 = traj[step]["pedestrian_state"]["position"][1]
								for mix_idx in range(self.args.n_mixtures):
									# plot predicted trajectory global frame
									for pred_step in range(self.args.prediction_horizon):
										idx = pred_step * self.args.output_pred_state_dim * self.args.n_mixtures + mix_idx
										mu_x = prediction_sample[0, idx]
										mu_y = prediction_sample[0, idx + self.args.n_mixtures]
										sigmax[pred_step, :] = prediction_sample[0][idx + 2 * self.args.n_mixtures]
										sigmay[pred_step, :] = prediction_sample[0][idx + 3 * self.args.n_mixtures]
										vel_pred[pred_step, :] = [mu_x, mu_y]

									pred_vel_global_frame = vel_pred
									traj_pred = sup.path_from_vel(initial_pos=np.array([x0,y0]),pred_vel=pred_vel_global_frame, dt=self.args.dt)

									self.ga3c_plot.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')

									# prior of 0 on the uncertainty of the pedestrian velocity
									sigma_x = np.square(sigmax[0]) * self.args.dt * self.args.dt
									sigma_y = np.square(sigmay[0]) * self.args.dt * self.args.dt

									e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
									             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
									e1.set_alpha(0.5)
									self.ga3c_plot.add_patch(e1)
									#self.ga3c_plot.plot(traj_pred[:, 0], traj_pred[:, 1], color=self.colors[mix_idx], label='Predicted trajectory')
									for pred_step in range(1, self.args.prediction_horizon):
										sigma_x += np.square(sigmax[pred_step]) * self.args.dt * self.args.dt
										sigma_y += np.square(sigmay[pred_step]) * self.args.dt * self.args.dt
										e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
										             height=np.sqrt(sigma_y) / 2,
										             angle=0 / np.pi * 180)
										self.ga3c_plot.add_patch(e1)

						# Plot other agents
						plot_other_agents = True
						if plot_other_agents:
							other_agents = np.expand_dims(other_agents_list[animation_idx][0, step], axis=0)

							for agent_id in range(other_agents.shape[0]):
								self.ga3c_plot.plot(other_agents[agent_id,0], other_agents[agent_id,1], color='b', marker='o')

					# Real trajectory
					self.ga3c_plot.plot(traj_real[:, 0], traj_real[:, 1], color="k",marker='x', label='Real trajectory')

					model_vel_real = np.zeros((self.args.prediction_horizon, self.args.output_dim))
					for i in range(self.args.prediction_horizon):
						mu_x = traj[step+i]["pedestrian_state"]["velocity"][0]
						mu_y = traj[step+i]["pedestrian_state"]["velocity"][1]
						model_vel_real[i, :] = [mu_x, mu_y]
					if False:
						real_vel = sup.rotate_predicted_vel_to_global_frame(model_vel_real, batch_x[step, 2:])
					else:
						real_vel = model_vel_real
					traj_real_from_v = sup.path_from_vel(initial_pos=np.array([x0, y0]),
					                              pred_vel=real_vel, dt=self.args.dt)
					self.ga3c_plot.plot(traj_real_from_v[:, 0], traj_real_from_v[:, 1], color='m', label='Agent real trajectory')


					self.ga3c_plot.set_xlim([-15, 15])
					self.ga3c_plot.set_ylim([-15, 15])
					self.ga3c_plot.axis("on")
					self.ga3c_plot.set_xlabel('x [m]')
					self.ga3c_plot.set_ylabel('y [m]')
					self.ga3c_plot.set_aspect('equal')

					self.fig_ga3c.canvas.draw()
					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario+"/figs/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario+"/figs/")
					if not os.path.exists(self.args.model_path + '/results/' + self.args.scenario + "/figs_frozen/"):
						os.makedirs(self.args.model_path + '/results/' + self.args.scenario + "/figs_frozen/")
					if test_args.save_figs:
						if test_args.freeze_other_agents:
							self.fig_ga3c.savefig(self.args.model_path + '/results/' + self.args.scenario+"/figs_frozen/result_"+str(animation_idx)+"_"+str(step)+".jpg")
						else:
							self.fig_ga3c.savefig(self.args.model_path + '/results/' + self.args.scenario + "/figs/result_" + str(
								animation_idx) + "_" + str(step) + ".jpg")
					self.writer.grab_frame()

src/data_utils/Recorder.py
eliatrevisan/IA-MPPI-LBM
import os
import numpy as np
from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter
from matplotlib.patches import Ellipse
# from mpl_toolkits.mplot3d import Axes3D
# import mpl_toolkits.mplot3d.axes3d as p3

from src.data_utils.DataHandler_Keras_v2 import prediction_to_trajectory

class Recorder(object): 
    def __init__(self, all_trajectories, all_predictions, args, save = False, display = True, figsize = (8,6), nframes = 10000): 
        self.dt = args.dt
        self.all_predictions = all_predictions
        self.all_trajectories = all_trajectories
        self.args = args
        # self.fps = int(1/self.dt)
        self.fps = int(10/self.dt)
        
        self.fig = plt.figure(figsize=figsize)
        self.ax = plt.axes(xlim=(-15, 15), ylim=(-15, 15))
        
        self.stream = self.data_stream()

        total_frames = 0
        for prediction in all_predictions:
            total_frames += prediction.shape[0]

        self.nframes = min(nframes, total_frames) 

        self.ani = animation.FuncAnimation(self.fig, self.animate, frames = self.nframes, init_func=self.setup_plot, interval=1/self.fps, blit=False, repeat=False)

        if save:
            self.save_animation(self.ani)

        if display:
            plt.show()


    def setup_plot(self):
        # Axis setup
        self.ax.set_xlabel('X (m)')
        self.ax.set_ylabel('Y (m)')

        pedestrian, = self.ax.plot([], [], marker="o", color='b')
        pedestrian_goal, = self.ax.plot([], [],  marker="D", color='b')
        robot, = self.ax.plot([], [], marker="o", color='r')
        robot_goal, = self.ax.plot([], [], marker="D", color='r')
        complete_path, = self.ax.plot([], [], marker="o", fillstyle='none', markersize=5, lw=1.5)
        predicted_path, = self.ax.plot([], [], marker="x", markersize=5, lw=1.5, color='g')
        self.objects = {
            "pedestrian": pedestrian,
            "pedestrian_goal": pedestrian_goal,
            "robot": robot,
            "robot_goal": robot_goal,
            "complete_path": complete_path,
            "predicted_path": predicted_path
        }
        return self.objects,

    def animate(self, iteration):
        data = next(self.stream)
        
        prediction = prediction_to_trajectory(data["pedestrian_position"], data["prediction"], self.args)
        
        self.objects["pedestrian"].set_data(data["pedestrian_position"][0], data["pedestrian_position"][1])
        self.objects["robot"].set_data(data["robot_position"][0], data["robot_position"][1])
        self.objects["pedestrian_goal"].set_data(data["pedestrian_goal"][0], data["pedestrian_goal"][1])
        self.objects["robot_goal"].set_data(data["robot_goal"][0], data["robot_goal"][1])
        self.objects["complete_path"].set_data(data["complete_path"][:, 0], data["complete_path"][:, 1])
        self.objects["predicted_path"].set_data( prediction[:, 0], prediction[:, 1] )
        
        return self.objects,

    def data_stream(self):
        """ Generator for the data to be plotted """
        while True:
            for prediction_array, trajectory in zip(self.all_predictions, self.all_trajectories):
                complete_path = np.zeros((len(trajectory), 2))
                for idx, datapoint in enumerate(trajectory):
                    complete_path[idx, :] = datapoint['pedestrian_state']['position']
                    
                for prediction, datapoint in zip(prediction_array, trajectory):
                    yield {
                        "prediction": prediction,
                        "pedestrian_position": datapoint['pedestrian_state']['position'],
                        "robot_position": datapoint['robot_state'],
                        "pedestrian_goal": datapoint["pedestrian_goal_position"],
                        # "robot_goal": trajectory[-1]["robot_state"][0:2],
                        "robot_goal": datapoint["goal_position"],
                        "complete_path": complete_path
                    }

    def save_animation(self, ani):
        self.writer =  animation.writers['ffmpeg']
        self.writer = self.writer(fps=self.fps, metadata=dict(artist='Me'), bitrate=1800, extra_args=['-vcodec', 'libx264'])
        video_dir = os.path.join(self.args.model_path, self.args.scenario)
        Path(video_dir).mkdir(parents=True, exist_ok=True)
        video_name = os.path.join(video_dir, 'video.mp4')
        self.ani.save(video_name, writer=self.writer)

src/data_utils/Recorder_Keras.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import pylab as pl
import math
import sys
import os

ros_path = '/opt/ros/kinetic/lib/python2.7/dist-packages'
if ros_path in sys.path:
    sys.path.remove(ros_path)
import cv2
import cv2 as cv
sys.path.append('../../src')

def compute_radial_distance_vector(number_elements, relative_positions, max_range=10.0, min_angle=0, max_angle=2*np.pi, normalize=False):
  """
  Compute the distance to surrounding objects in a radially discretized way.
  For each element there will be a floating point distance to the closest object in this sector.
  This allows to preserve the continuous aspect of the distance vs. a standard grid.

  !!! Attention: 0 angle is at the negative x-axis.

  number_elements: radial discretization
  relative_positions: relative positions of the surrounding objects in the local frame (numpy array)
  max_range: maximum range of the distance measurement

  returns:
  radial_dist_vector: contains the distance to the closest object in each sector
  """
  radial_dist_vector = max_range * np.ones([number_elements])
  radial_resolution = (max_angle-min_angle) / float(number_elements)
  if np.any(np.isnan(relative_positions)):
    print("inf")
  for ii in range(relative_positions.shape[0]):
    phi = math.atan2(relative_positions[ii, 1], relative_positions[ii, 0]) + np.pi
    rad_idx = int((phi - min_angle) / radial_resolution)
    if rad_idx >= 0 and rad_idx < number_elements:
      radial_dist_vector[rad_idx] = min(radial_dist_vector[rad_idx], np.linalg.norm(relative_positions[ii,:]))
  if normalize:
    radial_dist_vector /= max_range

  return radial_dist_vector

def sigmoid(x):
  return 1.0/(1.0+np.exp(-x))

def str2bool(v):
  if isinstance(v, bool):
    return v
  if v.lower() in ('yes', 'true', 't', 'y', '1'):
    return True
  elif v.lower() in ('no', 'false', 'f', 'n', '0'):
    return False
  else:
    raise argparse.ArgumentTypeError('Boolean value expected.')

def positions_in_local_frame(ego_pos, heading, other_pos):
  proj_matrix = np.array([[np.cos(heading), np.sin(heading)], [-np.sin(heading), np.cos(heading)]])
  local_frame_pos = np.zeros([0, 2])

  for ii in range(other_pos.shape[0]):
    rel_pos = other_pos[ii, :] - ego_pos
    local_frame_pos = np.append(local_frame_pos, np.expand_dims(np.matmul(proj_matrix, rel_pos), axis=0), axis=0)

  return local_frame_pos

def rotate_grid_around_center(grid, angle):
    """
    inputs:
      grid: numpy array (gridmap) that needs to be rotated
      angle: rotation angle in degrees
    """
    # Rotate grid into direction of initial heading
    grid = grid.copy()
    rows, cols = grid.shape
    M = cv2.getRotationMatrix2D(center=(rows / 2, cols / 2), angle=angle, scale=1)
    grid = cv2.warpAffine(grid, M, (rows, cols))

    return grid

def rotate_batch(batch_y, batch_x):
  """
    inputs:
      batch_y: trainning velocities batch containing vx, vy on the global frame
      batch_x: trainning or prediction state batch containing x,y, vx, vy on the global frame
    """
  # rotate initial vx and vy in global frame to quary-agent local frame
  by = batch_y.copy()  # rotated output batch values
  heading = math.atan2(batch_x[1], batch_x[0])
  rot_mat = np.array([[np.cos(-heading), -np.sin(-heading)], [np.sin(-heading), np.cos(-heading)]])
  for pred_step in range(by.shape[0]):
    by[pred_step,:] = np.dot(rot_mat, by[pred_step,:])
  return by

def rotate_batch_to_local_frame(batch_y, batch_x):
  """
    inputs:
      batch_y: trainning velocities batch containing vx, vy on the global frame
      batch_x: trainning or prediction state batch containing x,y, vx, vy on the global frame
    """
  # rotate initial vx and vy in global frame to quary-agent local frame
  bx = batch_x.copy()  # rotated input batch values
  by = batch_y.copy()  # rotated output batch values
  for batch_idx in range(batch_x.shape[0]):
    for tbp_step in range(batch_x.shape[1]):
      heading = math.atan2(bx[batch_idx, tbp_step, 3], bx[batch_idx, tbp_step, 2])
      rot_mat = np.array([[np.cos(-heading), -np.sin(-heading)], [np.sin(-heading), np.cos(-heading)]])
      bx[batch_idx, tbp_step, 2:] = np.dot(rot_mat, bx[batch_idx, tbp_step, 2:])
      for pred_step in range(int(by.shape[2] / 2)):
        by[batch_idx, tbp_step, 2 * pred_step:2 * pred_step + 2] = np.dot(rot_mat, by[batch_idx, tbp_step,
                                                                                   2 * pred_step:2 * pred_step + 2])
  return bx , by

def plot_submap(ax, origin, grid, resolution, submap_size):
    ax.imshow(grid, cmap='gray_r')

                    
    pl.show()
    

def plot_grid(ax, center, grid, grid_resolution, submap_size):
    """
    Plot a binary occupancy grid in the axis ax.
    Specify center, resolution and size in order to align it with the coordinate frame.
    """
    plot_values = np.zeros([0, 2])
    for idx_x in range(grid.shape[0]):
        for idx_y in range(grid.shape[1]):
            grid_coordinate_local = np.array([idx_x, idx_y]) * grid_resolution - np.array(submap_size) / 2.0
            grid_coordinate_global = grid_coordinate_local + center
            if grid[idx_x, idx_y] == 1:
                plot_values = np.append(plot_values, np.array([[grid_coordinate_global[0], grid_coordinate_global[1]]]),
                                        axis=0)
    return ax.plot(plot_values[:, 0], plot_values[:, 1], marker='s', alpha=0.8, color='k', lw=0)

def plot_grid_roboat(ax, origin, grid, resolution, size):
    ax.imshow(grid,
               extent = (origin[0], origin[0] + size[0] * resolution,
                         origin[1], origin[1] + size[1] * resolution),
               cmap='gray_r')
    ax.set_facecolor('#262626')
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    ax.yaxis.set_label_coords(-0.08, .5)
    ax.xaxis.set_label_coords(0.5, -0.09)

    if resolution == 0.4378: # Amstel
      ax.set_xlim((-280,-230))
      ax.set_ylim((-25,25))
    
    if resolution == 0.22: # Bloemgracht
      ax.set_xlim((-20,50))
      ax.set_ylim((-40,30))

def create_map_from_png(file_name,resolution,map_size,map_center,data_path='../data',):
  # Create grid for SF model
  print("create_map_from_png")
  grid_map = {}
  grid_map['Resolution'] = resolution
  grid_map['Size'] = map_size  # map size in [m]
  grid_map['Map'] = np.zeros((int(grid_map['Size'][0] / grid_map['Resolution']),
                              int(grid_map['Size'][1] / grid_map['Resolution'])))

  #map_center = grid_map['Size'] / 2  # hack for my dataset
  H = np.genfromtxt(os.path.join(data_path, 'H.txt'),
                    delimiter='  ',
                    unpack=True).transpose()

  # Extract static obstacles
  obst_threshold = 200
  static_obst_img = cv.imread(os.path.join(data_path, file_name), 0)
  #static_obst_img = cv.imread(os.path.join(load_data_path, 'map.png'), 0)
  obstacles = np.zeros([0, 3])
  #grid = (static_obst_img/255*-1)+1
  #static_obst_img = np.transpose(static_obst_img)
  for xx in range(static_obst_img.shape[0]):
    for yy in range(static_obst_img.shape[1]):
      if static_obst_img[xx, yy] > obst_threshold:
        obstacles = np.append(obstacles,
                              np.dot(H, np.array([[xx], [yy], [1]])).transpose(),
                              axis=0)
  Hinv = np.linalg.inv(H)

  # Compute obstacles in 2D
  obstacles_2d = np.zeros([obstacles.shape[0], 2])
  obstacles_2d[:, 0] = obstacles[:, 0] / obstacles[:, 2]
  obstacles_2d[:, 1] = obstacles[:, 1] / obstacles[:, 2]

  # Get obstacle idx on map
  obst_idx = []
  for obst_ii in range(obstacles_2d.shape[0]):
    idx = np.dot(Hinv, np.array([[obstacles_2d[obst_ii, 0]], [obstacles_2d[obst_ii, 1]], [1]])).transpose()
    obst_idx.append((int(idx[0,0]),int(idx[0,1])))
    grid_map['Map'][obst_idx[-1]] = 1


  grid_map['Closest X'] = np.zeros_like(grid_map['Map']) + 10000
  grid_map['Closest Y'] = np.zeros_like(grid_map['Map']) + 10000
  """ """
  # Calculate closest obstacle
  obst_idx = np.array(obst_idx)
  for xx in range(int(grid_map['Size'][0] / grid_map['Resolution'])):
    for yy in range(int(grid_map['Size'][1] / grid_map['Resolution'])):
      delta_idx = obst_idx - np.array([xx, yy])
      distances = np.sqrt(np.sum(np.square(delta_idx), axis=1))
      closest_obj = np.argmin(distances)
      grid_map['Closest X'][xx, yy] = obst_idx[closest_obj][0]
      grid_map['Closest Y'][xx, yy] = obst_idx[closest_obj][1]

  grid_map['Closest X'] = grid_map['Closest X'] * grid_map['Resolution'] - map_center[0]
  grid_map['Closest Y'] = grid_map['Closest Y'] * grid_map['Resolution'] - map_center[1]

  np.save(os.path.join(data_path, 'map'), grid_map)

def idx_from_pos(x, y, center, res=0.1):
  idx_x = round((x + float(center[0])) / res)
  idx_y = round((y + float(center[1])) / res)

  # Projecting index on map if out of bounds
  idx_x = max(0, min(idx_x, -1 + center[0] * 2. / res))
  idx_y = max(0, min(idx_y, -1 + center[1] * 2. / res))
  return int(idx_x), int(idx_y)

def path_from_vel(initial_pos, pred_vel,v0=np.array([0,0]), dt=0.1, n_vx=1, n_vy=1):
  """
  Extract a path from vector of predicted velocities by applying Euler forward integration.
  """
  vel_x = pred_vel[:,0]*n_vx
  vel_y = pred_vel[:,1]*n_vy
  n_steps = len(vel_x)
  pos_vec = np.zeros([n_steps, 2])
  pos_vec[0, :] = initial_pos + v0*dt
  for ii in range(1, n_steps):
    pos_vec[ii, :] = pos_vec[ii-1,:] + dt * np.array([vel_x[ii], vel_y[ii]])

  return pos_vec

def to_image_frame(Hinv, loc):
  """
	Given H^-1 and world coordinates, returns (u, v) in image coordinates.
	"""
  locHomogenous = np.hstack((loc, np.ones((loc.shape[0], 1))))
  if locHomogenous.ndim > 1:
    loc_tr = np.transpose(locHomogenous)
    loc_tr = np.matmul(Hinv, loc_tr)  # to camera frame
    locXYZ = np.transpose(loc_tr / loc_tr[2])  # to pixels (from millimeters)
    return locXYZ[:, :2].astype(int)
  else:
    locHomogenous = np.dot(Hinv, locHomogenous)  # to camera frame
    locXYZ = locHomogenous / locHomogenous[2]  # to pixels (from millimeters)
    return locXYZ[:2].astype(int)

def to_pos_frame(Hinv, loc):
  """
	Given H^-1 and world coordinates, returns (u, v) in image coordinates.
	"""
  locHomogenous = np.hstack((loc, np.ones((loc.shape[0], 1))))
  if locHomogenous.ndim > 1:
    loc_tr = np.transpose(locHomogenous)
    loc_tr = np.matmul(Hinv, loc_tr)  # to camera frame
    locXYZ = np.transpose(loc_tr / loc_tr[2])  # to pixels (from millimeters)
    return locXYZ[:, :2]
  else:
    locHomogenous = np.dot(Hinv, locHomogenous)  # to camera frame
    locXYZ = locHomogenous / locHomogenous[2]  # to pixels (from millimeters)
    return locXYZ[:2]

def line_cv(im, ll, value, width):
  for tt in range(ll.shape[0] - 1):
    cv2.line(im, (int(ll[tt][1]), int(ll[tt][0])), (int(ll[tt + 1][1]), int(ll[tt + 1][0])), value, width)

def plotGrid(grid, ax, color='k', alpha=1.0):
  plot_values_ped = np.zeros([0, 2])
  plot_values_grid = np.zeros([0, 2])
  for idx_x in range(grid.shape[0]):
    for idx_y in range(grid.shape[1]):
      if grid[idx_x, idx_y] > 0.8 and grid[idx_x, idx_y] < 1.5:
        plot_values_grid = np.append(plot_values_grid, np.array([[idx_x, idx_y]]), axis=0)
      elif grid[idx_x, idx_y] > 1.5:
        plot_values_ped = np.append(plot_values_ped, np.array([[idx_x, idx_y]]), axis=0)
  ax.plot(plot_values_grid[:, 0], plot_values_grid[:, 1], marker='o', color=color, lw=0)
  ax.plot(plot_values_ped[:, 0], plot_values_ped[:, 1], marker='o', color='r', lw=0)
  ax.set_xlim([0, grid.shape[0]])
  ax.set_ylim([0, grid.shape[1]])


src/data_utils/Support.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
from scipy import interpolate

class Trajectory():
  
  def __init__(self, time_vec = np.zeros([0]), pose_vec = np.zeros([0,3]), vel_vec = np.zeros([0,3]), goal = np.zeros([1,2])):
    self.time_vec = time_vec    # timesteps in [ns]
    self.pose_vec = pose_vec    # [x, y, heading]
    self.vel_vec = vel_vec      # [vx, vy, omega]
    self.goal = goal            # [x,y]
    self.pose_interp = None
    self.vel_interp = None
    self.other_agents_positions = []  # store indices of other agents trajectories for each time step of this trajectory (with which other positions does it need to be compared at a certain time)
    self.other_agents_velocities = []
    
  def __len__(self):
    if len(self.time_vec) == len(self.pose_vec) == len(self.vel_vec):
      return len(self.time_vec)
    else: 
      raise Exception('Member vectors are of different length. Cannot get unique length.')
    
  def addData(self, timestamp, pose, vel, goal):
    """
    Inserts the provided data at the right index in the trajectory.
    """
    assert np.all(goal == self.goal), "Goals have to be matching."
    # Add data at the right position. In the log the data might not be ordered properly.
    if self.time_vec.shape[0] > 0 and timestamp < self.time_vec[-1]:
      try: 
        idx_insert = np.where(timestamp > self.time_vec)[-1][-1] + 1
      except:
        idx_insert = 0  # In case the timestamp is smaller than the first element of time_vec
      self.time_vec = np.insert(self.time_vec, idx_insert, timestamp)
      self.pose_vec = np.insert(self.pose_vec, idx_insert, np.atleast_2d(pose), axis=0)
      self.vel_vec = np.insert(self.vel_vec, idx_insert, np.atleast_2d(vel), axis=0)
    else:
      self.time_vec = np.append(self.time_vec, timestamp)
      self.pose_vec = np.append(self.pose_vec, np.atleast_2d(pose), axis=0)
      self.vel_vec = np.append(self.vel_vec, np.atleast_2d(vel), axis=0)
    
    if self.time_vec.shape[0] > 1:
      self.updateInterpolators()
      
  def subsample(self, subsampling_factor):
    """
    Subsample the trajectory in order to reduce the overall number of samples and 
    increase the time difference between two samples.
    """
    if int(self.time_vec[0]*10) % 2 == 1:
      self.time_vec = self.time_vec[1:]
      self.pose_vec = self.pose_vec[1:, :]
      self.vel_vec = self.vel_vec[1:, :]
    self.time_vec = self.time_vec[0::subsampling_factor]
    self.pose_vec = self.pose_vec[0::subsampling_factor, :]
    self.vel_vec = self.vel_vec[0::subsampling_factor, :]

    
    #self.updateInterpolators()
        
  def updateInterpolators(self):
    self.pose_interp = interpolate.interp1d(self.time_vec, self.pose_vec, kind='linear', axis=0)  # TODO: correct heading interpolation
    self.vel_interp = interpolate.interp1d(self.time_vec, self.vel_vec, kind='linear', axis=0)
      
  def contains(self, query_time):
    return query_time >= self.time_vec[0] and query_time <= self.time_vec[-1]
   
  def getPoseAtTime(self, query_time):
    pose = self.pose_interp(query_time)
    return pose
   
  def getVelAtTime(self, query_time):
    return self.vel_interp(query_time)
  
  def getDataAtTime(self, query_time):
    pose_interpol = self.getPoseAtTime(query_time)
    vel_interpol = self.getVelAtTime(query_time)
    return pose_interpol, vel_interpol
  
  def getMinTime(self):
    """
    Return the starting time of the trajectory.
    """
    return self.time_vec[0]
  
  def getMaxTime(self):
    """
    Return the end time of the trajectory.
    """
    return self.time_vec[-1]  
  
  def getDuration(self):
    return self.getMaxTime() - self.getMinTime()
  
  def smoothenTrajectory(self, dt=0.3):
    """
    Cubic interpolation with provided values in order to smoothen the trajectory and obtain
    a sequence with the specified dt.
    """
    x_interpolator = interpolate.interp1d(self.time_vec, self.pose_vec[:,0], kind='cubic', axis=0)
    y_interpolator = interpolate.interp1d(self.time_vec, self.pose_vec[:,1], kind='cubic', axis=0)
    vx_interpolator = interpolate.interp1d(self.time_vec, self.vel_vec[:,0], kind='cubic', axis=0)
    vy_interpolator = interpolate.interp1d(self.time_vec, self.vel_vec[:,1], kind='cubic', axis=0)
    n_elem = int((self.getMaxTime() - self.getMinTime()) / dt)
    new_time_vec = np.linspace(self.getMinTime(), self.getMinTime() + (n_elem - 1) * dt, n_elem)
    
    new_time_vec = np.zeros([n_elem])
    new_pose_vec = np.zeros([n_elem, 3])
    new_vel_vec = np.zeros([n_elem, 3])
    
    for ii in range(n_elem):
      t = self.getMinTime() + ii*dt
      new_time_vec[ii] = t
      new_pose_vec[ii, 0] = x_interpolator(t)
      new_pose_vec[ii, 1] = y_interpolator(t)
      new_vel_vec[ii, 0] = vx_interpolator(t)
      new_vel_vec[ii, 1] = vy_interpolator(t)
    
    self.goal = new_pose_vec[-1,:2]
    self.time_vec = new_time_vec
    self.pose_vec = new_pose_vec
    self.vel_vec = new_vel_vec
    

  def __getstate__(self):
    d = dict(self.__dict__)
    del d['pose_interp']
    del d['vel_interp']
    return d
    
src/data_utils/Trajectory.py
eliatrevisan/IA-MPPI-LBM
import numpy as np

class Config:
src/data_utils/config.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import os
import random
import sys
import math
import cv2
import pickle as pkl
import time
if sys.version_info[0] < 3:
  sys.path.append('../src/data_utils')
  import Support as sup
else:
	import src.data_utils.Support as sup
import matplotlib.pyplot as pl
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter
from matplotlib.patches import Ellipse

def plot_scenario(enc_seq, dec_seq,traj, predictions, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	metadata = dict(title='Movie Test', artist='Matplotlib')
	writer = FFMpegWriter(fps=10, metadata=metadata)

	fig = pl.figure("Trajectory Predictions")
	ax_in = pl.subplot()

	ax_in.clear()

	#x_lim = [-global_grid.map_size[0]/2 ,global_grid.map_size[0]/2]
	#y_lim = [-global_grid.map_size[1]/2,global_grid.map_size[1]/2]

	x_lim = [-5 ,5]
	y_lim = [-5,5]

	ax_in.clear()
	pedestrian_trajectory = np.zeros((len(traj),2))
	for t in range(len(traj)):
		pedestrian_trajectory[t,:] = traj[t]["pedestrian_state"]["position"]
	with writer.saving(fig, args.data_path + "/movie.mp4",100):
		for t in range(len(enc_seq)):
				ax_in.clear()
				# plot scenario grid
				#sup.plot_grid(ax_in, np.array([0, 0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
				ax_in.plot(pedestrian_trajectory[t:, 0], pedestrian_trajectory[t:, 1], color='red', label='Agent Real Traj')
				ax_in.plot(enc_seq[t][:,0], enc_seq[t][:,1], color='g', marker='o', label='Robot Prev Real Traj')
				ax_in.plot(enc_seq[t][:,4], enc_seq[t][:,5], color='red', marker='o', label='Agent Prev Traj')
				ax_in.plot(dec_seq[t][:, 0], dec_seq[t][:, 1], color='blue',marker='o', label='Agent Real')
				ax_in.plot(predictions[t][:, 0], predictions[t][:, 1], color='blue', marker='x', label='Agent Pred Traj')


				ax_in.set_xlim(x_lim)
				ax_in.set_ylim(y_lim)

				ax_in.legend()
				fig.canvas.draw()
				pl.show(block=False)
				#time.sleep(t)
				writer.grab_frame()

def plot_scenario_vel_tbp(enc_seq, dec_seq,traj, predictions, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	metadata = dict(title='Movie Test', artist='Matplotlib')
	writer = FFMpegWriter(fps=10, metadata=metadata)

	fig = pl.figure("Trajectory Predictions")
	ax_in = pl.subplot()

	ax_in.clear()

	#x_lim = [-global_grid.map_size[0]/2 ,global_grid.map_size[0]/2]
	#y_lim = [-global_grid.map_size[1]/2,global_grid.map_size[1]/2]

	x_lim = [-5 ,5]
	y_lim = [-5,5]

	ax_in.clear()
	pedestrian_trajectory = np.zeros((len(traj),2))
	for t in range(len(traj)):
		pedestrian_trajectory[t,:] = traj[t]["pedestrian_state"]["position"]
	with writer.saving(fig, args.data_path + "/movie.mp4",100):
		for t in range(len(enc_seq)):
				ax_in.clear()
				# plot scenario grid
				#sup.plot_grid(ax_in, np.array([0, 0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
				ax_in.plot(pedestrian_trajectory[t:, 0], pedestrian_trajectory[t:, 1], color='red', label='Agent Real Traj')
				ax_in.plot(enc_seq[t][0,0::args.input_dim], enc_seq[t][0,1::args.input_dim], color='g', marker='o', label='Robot Prev Real Traj')
				ax_in.plot(enc_seq[t][0,4::args.input_dim], enc_seq[t][0,5::args.input_dim], color='red', marker='o', label='Agent Prev Traj')
				ax_in.plot(pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 0], pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 1], color='blue',marker='o', label='Agent Real')
				x0 = pedestrian_trajectory[t+args.prev_horizon,0]
				y0 = pedestrian_trajectory[t+args.prev_horizon,1]
				pred_vel = np.zeros((args.prediction_horizon,args.output_dim))
				for i in range(args.prediction_horizon):
					idx = i * args.output_dim
					idy = idx + 1
					vx = predictions[t][0,idx]
					vy = predictions[t][0,idy]
					pred_vel[i, :] = [vx, vy]

				traj_pred = sup.path_from_vel(initial_pos=np.array([x0 , y0]),pred_vel=pred_vel, dt=args.dt)
				ax_in.plot(traj_pred[:, 0], traj_pred[:, 1], color='blue', marker='x', label='Agent Pred Traj')
				ax_in.set_xlim(x_lim)
				ax_in.set_ylim(y_lim)

				ax_in.legend()
				fig.canvas.draw()
				pl.show(block=False)
				#time.sleep(t)
				writer.grab_frame()

#def plot_scenario_vel(enc_seq, dec_seq,traj, predictions, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	"""	metadata = dict(title='Movie Test', artist='Matplotlib')
	writer = FFMpegWriter(fps=1, metadata=metadata)

	fig = pl.figure("Trajectory Predictions")
	ax_in = pl.subplot()

	ax_in.clear()

	#x_lim = [-global_grid.map_size[0]/2 ,global_grid.map_size[0]/2]
	#y_lim = [-global_grid.map_size[1]/2,global_grid.map_size[1]/2]

	x_lim = [-5 ,5]
	y_lim = [-5,5]

	ax_in.clear()
	pedestrian_trajectory = np.zeros((len(traj),2))
	for t in range(len(traj)):
		pedestrian_trajectory[t,:] = traj[t]["pedestrian_state"]["position"]
	with writer.saving(fig, args.data_path + "/movie.mp4",100):
		for t in range(len(enc_seq)):
				ax_in.clear()
				# plot scenario grid
				#sup.plot_grid(ax_in, np.array([0, 0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
				ax_in.plot(pedestrian_trajectory[t:, 0], pedestrian_trajectory[t:, 1], color='red', label='Agent Real Traj')
				ax_in.plot(enc_seq[t][:,0], enc_seq[t][:,1], color='g', marker='o', label='Robot Prev Real Traj')
				ax_in.plot(enc_seq[t][:,4], enc_seq[t][:,5], color='red', marker='o', label='Agent Prev Traj')
				ax_in.plot(pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 0], pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 1], color='blue',marker='o', label='Agent Real')
				x0 = pedestrian_trajectory[t+args.prev_horizon,0]
				y0 = pedestrian_trajectory[t+args.prev_horizon,1]
				pred_vel = predictions[t]
				traj_pred = sup.path_from_vel(initial_pos=np.array([x0 , y0]),pred_vel=pred_vel, dt=0.1)
				ax_in.plot(traj_pred[:, 0], traj_pred[:, 1], color='k', marker='x', label='Agent Pred Traj')
				ax_in.set_xlim(x_lim)
				ax_in.set_ylim(y_lim)

				ax_in.legend()
				fig.canvas.draw()
				pl.show(block=False)
				#time.sleep(t)
				writer.grab_frame()
"""

def plot_scenario_vel_OpenCV_simdata(trajectories, batch_target, all_predictions, args, exp_num=0,display = True):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
	"""
	homography_file = os.path.join(args.data_path + args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')
	scenario = args.scenario
	video_file = args.model_path +"/"  + args.model_name+ "_" + str(exp_num) + '.avi'
	img_file = args.data_path + args.scenario +'/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = ~cv2.imread(img_file)
		img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
		height, width, channels = img.shape

		scale_percent = 400  # percent of original size
		width = int(img.shape[1] * scale_percent / 100)
		height = int(img.shape[0] * scale_percent / 100)
		dim = (width, height)
		# resize image
		resized_img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		out = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width, height))

		# Adding legend
		font = cv2.FONT_HERSHEY_SIMPLEX
		cv2.line(resized_img, (width - 150, 30), (width - 130, 30), (0, 255, 0), 4)
		cv2.putText(resized_img, "Query Agent", (width - 120, 30), font, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 50), (width - 130, 50), (255, 0, 0), 4)
		cv2.putText(resized_img, "Other Agents", (width - 120, 50), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 70), (width - 130, 70), (255, 255, 0), 4)
		cv2.putText(resized_img, "Query Agent traj", (width - 120, 70), font, 0.5, (255, 255, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 90), (width - 130, 90), (128, 128, 128), 4)
		cv2.putText(resized_img, "Pred traj 1", (width - 120, 90), font, 0.5, (128, 128, 128), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 110), (width - 130, 110), (0, 128, 255), 4)
		cv2.putText(resized_img, "Pred traj 2", (width - 120, 110), font, 0.5, (0, 128, 255), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 130), (width - 130, 130), (255, 0, 127), 4)
		cv2.putText(resized_img, "Pred traj 3", (width - 120, 130), font, 0.5, (255, 0, 127), 2, cv2.LINE_AA)

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]
			ped_vel = batch_target[exp_idx][0]

			for t in range(len(predictions)):
				pedestrian_trajectory = traj.pose_vec[:,:2]
				overlay_ellipses = resized_img.copy()

				# Plot Query Agent Positions
				query_agent_pos = traj.pose_vec[t,:2]
				query_agent = np.expand_dims(np.array((query_agent_pos[0], query_agent_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, query_agent)*int(scale_percent / 100)
				sigma_x = 0.2
				sigma_y = 0.2
				axis = (
					int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100), int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
				center = (obsv_XY[0, 1], obsv_XY[0, 0])
				cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (0, 255, 0), -1)

				# Plot other pedestrians positions
				agent_pos = traj.other_agents_positions[t]
				for agent_id in range(agent_pos.shape[0]):
					ped_pos = np.expand_dims(np.array((agent_pos[agent_id, 0], agent_pos[agent_id, 1])), axis=0).astype(float)
					obsv_XY = sup.to_image_frame(Hinv, ped_pos)*int(scale_percent / 100)
					axis = (
						int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100),
						int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (255, 0, 0), 1)

				# Query-agent real trajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:, :])*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 0, 0), 1)  # bgr convention

				# Query-agent predicted real trajectory from positions
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:t+args.prediction_horizon, :]) * int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 1)  # bgr convention

				#Query-agent Predicted Real TRajectory
				x0 = traj.pose_vec[t+1+args.prev_horizon,0]
				y0 = traj.pose_vec[t+1+args.prev_horizon,1]
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.truncated_backprop_length):
					if args.normalize_data:
						pedestrian_velocity[t_idx, 0] = ped_vel[t, t_idx * 2] / args.sx_vel + args.min_vel_x
						pedestrian_velocity[t_idx, 1] = ped_vel[t, t_idx * 2+1] / args.sy_vel + args.min_vel_y
					else:
						pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_traj = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pedestrian_velocity, dt=args.dt)
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_traj)*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 1)  # bgr convention

				# Query-agent real trajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:t+args.prediction_horizon, :])*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (0, 0, 0), 1)  # bgr convention

				# Plot Predictions
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				for sample_id in range(len(predictions[t])):
					prediction_sample = predictions[t][sample_id]
					for mix_id in range(args.n_mixtures):
						for pred_step in range(args.prediction_horizon):
							idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
							if args.normalize_data:
								pred_vel[pred_step, 0] = prediction_sample[0][idx] / args.sx_vel + args.min_vel_x
								pred_vel[pred_step, 1] = prediction_sample[0][idx+args.n_mixtures] / args.sy_vel + args.min_vel_y
							else:
								pred_vel[pred_step, :] = np.array((prediction_sample[0][idx],prediction_sample[0][idx+args.n_mixtures]
							                                   ))
							if args.output_pred_state_dim > args.output_dim:
								sigmax[pred_step, :] = prediction_sample[0][idx + 2*args.n_mixtures]
								sigmay[pred_step, :] = prediction_sample[0][idx + 3*args.n_mixtures]
						try:
							if args.predict_positions:
								traj_pred = pred_vel
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						except:
							traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
						sup.line_cv(overlay_ellipses, obsv_XY, colors[mix_id], 1)  # bgr convention

						if args.output_pred_state_dim > args.output_dim:
							sigma_x = np.square(sigmax[0]) * (args.dt ** 2)
							sigma_y = np.square(sigmay[0]) * (args.dt ** 2)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
							axis = (
								int(np.sqrt(sigma_x) / args.submap_resolution * scale_percent / 100), int(np.sqrt(sigma_y) / args.submap_resolution * scale_percent / 100))
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							try:
								cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], 1)
							except:
								print("failed to plot ellipse")

							for pred_step in range(1,args.prediction_horizon):
								sigma_x += np.square(sigmax[pred_step]) * (args.dt ** 2)
								sigma_y += np.square(sigmay[pred_step]) * (args.dt ** 2)
								axis = (
									int(np.sqrt(sigma_x) / args.submap_resolution * scale_percent / 100), int(np.sqrt(sigma_y) / args.submap_resolution * scale_percent / 100))
								center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
								try:
									cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], 1)
								except:
									print("failed to plot ellipse")

				#image_new = cv2.addWeighted(resized_img, 0.2, overlay_ellipses, 0.8, 0)
				if display:
					cv2.imshow("image", overlay_ellipses)
					cv2.waitKey(100)
				out.write(overlay_ellipses)

		# When everything done, release the video capture and video write objects
		out.release()

	print("Done recording...")

def plot_local_scenario_vel_OpenCV(trajectories, batch_target, all_predictions, args, exp_num=0,display = True):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	homography_file = os.path.join(args.data_path + args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')
	scenario = args.scenario
	video_file = args.model_path +"/"  + args.model_name+ "_" + str(exp_num) + '_local.avi'
	img_file = args.data_path + args.scenario +'/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = ~cv2.imread(img_file)

		height, width, channels = img.shape

		scale_percent = 200  # percent of original size
		width = int(img.shape[1] * scale_percent / 100)
		height = int(img.shape[0] * scale_percent / 100)
		dim = (width, height)
		# resize image
		resized_img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		out = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (600, 600))

		# Adding legend
		font = cv2.FONT_HERSHEY_SIMPLEX
		cv2.line(resized_img, (width - 150, 30), (width - 130, 30), (0, 255, 0), 4)
		cv2.putText(resized_img, "Query Agent", (width - 120, 30), font, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 50), (width - 130, 50), (255, 0, 0), 4)
		cv2.putText(resized_img, "Other Agents", (width - 120, 50), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 70), (width - 130, 70), (255, 255, 0), 4)
		cv2.putText(resized_img, "Query Agent traj", (width - 120, 70), font, 0.5, (255, 255, 0), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 90), (width - 130, 90), (128, 128, 128), 4)
		cv2.putText(resized_img, "Pred traj 1", (width - 120, 90), font, 0.5, (128, 128, 128), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 110), (width - 130, 110), (0, 128, 255), 4)
		cv2.putText(resized_img, "Pred traj 2", (width - 120, 110), font, 0.5, (0, 128, 255), 2, cv2.LINE_AA)
		cv2.line(resized_img, (width - 150, 130), (width - 130, 130), (255, 0, 127), 4)
		cv2.putText(resized_img, "Pred traj 3", (width - 120, 130), font, 0.5, (255, 0, 127), 2, cv2.LINE_AA)

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]
			ped_vel = batch_target[exp_idx][0]

			for t in range(len(predictions)):
				pedestrian_trajectory = traj.pose_vec[:,:2]
				overlay_ellipses = resized_img.copy()

				# Plot Query Agent Positions
				query_agent_pos = traj.pose_vec[t,:2]
				query_agent = np.expand_dims(np.array((query_agent_pos[0], query_agent_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, query_agent)*int(scale_percent / 100)
				sigma_x = 0.4
				sigma_y = 0.4
				axis = (
					int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100), int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
				agent_center = (obsv_XY[0, 1], obsv_XY[0, 0])
				cv2.ellipse(overlay_ellipses, agent_center, axis, 0, 0, 360, (0, 255, 0), -1)

				# Plot ped real position
				agent_pos = traj.other_agents_positions[t]
				for agent_id in range(agent_pos.shape[0]):
					ped_pos = np.expand_dims(np.array((agent_pos[agent_id, 0], agent_pos[agent_id, 1])), axis=0).astype(float)
					obsv_XY = sup.to_image_frame(Hinv, ped_pos)*int(scale_percent / 100)
					axis = (
						int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100),
						int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (255, 0, 0), -1)

				# Query-agent real trajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:, :])*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 0, 0), 1)  # bgr convention

				#Query-agent Predicted Real TRajectory
				x0 = traj.pose_vec[t+1+args.prev_horizon,0]
				y0 = traj.pose_vec[t+1+args.prev_horizon,1]
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.truncated_backprop_length):
					if args.normalize_data:
						pedestrian_velocity[t_idx, 0] = ped_vel[t, t_idx * 2] / args.sx_vel + args.min_vel_x
						pedestrian_velocity[t_idx, 1] = ped_vel[t, t_idx * 2+1] / args.sy_vel + args.min_vel_y
					else:
						pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_trajectory = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pedestrian_velocity, dt=args.dt)
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory)*int(scale_percent / 100)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 1)  # bgr convention


				# Plot Predictions
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				for sample_id in range(len(predictions[t])):
					prediction_sample = predictions[t][sample_id]
					for mix_id in range(args.n_mixtures):
						for pred_step in range(args.prediction_horizon):
							idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
							if args.normalize_data:
								pred_vel[pred_step, 0] = prediction_sample[0][idx] / args.sx_vel + args.min_vel_x
								pred_vel[pred_step, 1] = prediction_sample[0][idx+args.n_mixtures] / args.sy_vel + args.min_vel_y
							else:
								pred_vel[pred_step, :] = np.array((prediction_sample[0][idx],prediction_sample[0][idx+args.n_mixtures]))
							if args.output_pred_state_dim > args.output_dim:
								sigmax[pred_step, :] = prediction_sample[0][idx + 2*args.n_mixtures]
								sigmay[pred_step, :] = prediction_sample[0][idx + 3*args.n_mixtures]
						try:
							if args.predict_positions:
								traj_pred = pred_vel
							else:
								traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						except:
							traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
						sup.line_cv(overlay_ellipses, obsv_XY, colors[mix_id], 1)  # bgr convention

						if args.output_pred_state_dim > args.output_dim:
							sigma_x = np.square(sigmax[0]) * (args.dt ** 2)
							sigma_y = np.square(sigmay[0]) * (args.dt ** 2)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
							axis = (
								int(np.sqrt(sigma_x) / args.submap_resolution * scale_percent / 100), int(np.sqrt(sigma_y) / args.submap_resolution * scale_percent / 100))
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], 1)

							for pred_step in range(1,args.prediction_horizon):
								sigma_x += np.square(sigmax[pred_step]) * (args.dt ** 2)
								sigma_y += np.square(sigmay[pred_step]) * (args.dt ** 2)
								axis = (
									int(np.sqrt(sigma_x) / args.submap_resolution * scale_percent / 100), int(np.sqrt(sigma_y) / args.submap_resolution * scale_percent / 100))
								center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
								try:
									cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], 1)
								except:
									print("failed to plot ellipse")

				crop_img = cv2.getRectSubPix(overlay_ellipses,(600,600),agent_center)
				if display:
					cv2.imshow("image", crop_img)
					cv2.waitKey(100)
				out.write(crop_img)

		# When everything done, release the video capture and video write objects
		out.release()

	print("Done recording...")

def plot_scenario_vel_OpenCV(trajectories, batch_target, all_predictions, args, exp_num=0,display = False):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	homography_file = os.path.join(args.data_path+ args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')
	scenario = args.scenario
	video_file = args.model_path +"/"  + args.model_name+ "_" + str(exp_num) + '.avi'
	map_resolution = 1 / Hinv[0, 0]
	print("Recording to: " + video_file)

	img_file = args.data_path+ args.scenario + '/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = cv2.imread(img_file)

		#height, width, channels = img.shape
		width = 600
		height = 600
		dim = (width, height)

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		out = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width, height))

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]
			ped_vel = batch_target[exp_idx][0]

			pedestrian_trajectory = np.zeros((len(traj), 2))
			for t in range(len(traj)):
				pedestrian_trajectory[t, :] = traj[t]["pedestrian_state"]["position"]

			for t in range(len(predictions)):
				overlay_ellipses = img.copy()

				# Plot Robot Positions
				robot_pos = traj[t]["robot_state"]
				robot = np.expand_dims(np.array((robot_pos[0], robot_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, robot)
				sigma_x = 0.4
				sigma_y = 0.4
				axis_agent = (
					int(np.sqrt(sigma_x) / map_resolution), int(np.sqrt(sigma_y) / map_resolution))
				center = (obsv_XY[0, 0], obsv_XY[0, 1])
				cv2.ellipse(overlay_ellipses, center, axis_agent, 0, 0, 360, (0, 0, 0), -1)

				# Plot ped real position
				agent_pos = traj[t]["pedestrian_state"]["position"]
				ped_pos = np.expand_dims(np.array((agent_pos[0], agent_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, ped_pos)

				center = (obsv_XY[0, 0], obsv_XY[0, 1])
				cv2.ellipse(overlay_ellipses, center, axis_agent, 0, 0, 360, (255, 0, 0), -1)

				# Pedestrian real trajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t:, :])
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 0, 0), 3)  # bgr convention

				#Pedestrian Predicted Real TRajectory
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory[t + args.prev_horizon:t + args.prev_horizon + args.prediction_horizon, :])
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 3)  # bgr convention

				#Pedestrian Predicted Real TRajectory
				x0 = pedestrian_trajectory[t, 0]
				y0 = pedestrian_trajectory[t, 1]
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.truncated_backprop_length):
					pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_traj = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pedestrian_velocity, dt=args.dt)
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_traj)
				sup.line_cv(overlay_ellipses, obsv_XY, (0, 0, 255), 3)  # bgr convention

				# Plot Predictions
				x0 = pedestrian_trajectory[t, 0]
				y0 = pedestrian_trajectory[t, 1]
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				for sample_id in range(len(predictions[t])):
					prediction_sample = predictions[t][sample_id]
					for mix_id in range(args.n_mixtures):
						for pred_step in range(args.prediction_horizon):
							idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
							pred_vel[pred_step, :] = np.array((prediction_sample[0][idx],prediction_sample[0][idx+args.n_mixtures]
							                                   ))
							if args.output_pred_state_dim > args.output_dim:
								sigmax[pred_step, :] = prediction_sample[0][idx + 2*args.n_mixtures]
								sigmay[pred_step, :] = prediction_sample[0][idx + 3*args.n_mixtures]
						traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
						obsv_XY = sup.to_image_frame(Hinv, traj_pred)
						sup.line_cv(overlay_ellipses, obsv_XY, colors[mix_id], 3)  # bgr convention

						if args.output_pred_state_dim > args.output_dim:
							sigma_x = np.square(np.exp(sigmax[0])) * (args.dt ** 2)
							sigma_y = np.square(np.exp(sigmay[0])) * (args.dt ** 2)

							obsv_XY = sup.to_image_frame(Hinv, traj_pred)
							axis = (
								int(np.sqrt(sigma_x) / map_resolution), int(np.sqrt(sigma_y) / map_resolution))
							#axis = (int(axis[0] + axis_agent[0]),int(axis[1] + axis_agent[1]))
							center = (obsv_XY[0, 0], obsv_XY[0, 1])
							cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], -1)

							for pred_step in range(1,args.prediction_horizon):
								sigma_x += np.square(np.exp(sigmax[pred_step])) * (args.dt ** 2)
								sigma_y += np.square(np.exp(sigmay[pred_step])) * (args.dt ** 2)
								axis = (
									int(np.sqrt(sigma_x) / map_resolution), int(np.sqrt(sigma_y) / map_resolution))
								#axis = (int(axis[0] + axis_agent[0]), int(axis[1] + axis_agent[1]))
								center = (obsv_XY[pred_step, 0], obsv_XY[pred_step, 1])
								cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], -1)

				image_new = cv2.addWeighted(img, 0.2, overlay_ellipses, 0.8, 0)

				# resize image
				resized_img = cv2.resize(image_new, dim, interpolation=cv2.INTER_AREA)
				if display:
					cv2.imshow("image", resized_img)
					cv2.waitKey(100)
				out.write(resized_img)

	# When everything done, release the video capture and video write objects
	out.release()

	print("Done recording...")

def plot_scenario_vel(trajectories, all_predictions, args, grid = None):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	metadata = dict(title='Movie Test', artist='Matplotlib')
	writer = FFMpegWriter(fps=2, metadata=metadata)

	fig = pl.figure("Trajectory Predictions")
	ax_in = pl.subplot()

	ax_in.clear()

	x_lim = [-5 ,5]
	y_lim = [-5,5]

	ax_in.clear()

	with writer.saving(fig, args.data_path + "/movie.mp4",100):

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]

			pedestrian_trajectory = np.zeros((len(traj), 2))
			for t in range(len(traj)):
				pedestrian_trajectory[t, :] = traj[t]["pedestrian_state"]["position"]

			for t in range(len(predictions)):
				ax_in.clear()
				# plot scenario grid
				#sup.plot_grid(ax_in, np.array([0, 0]), global_grid.gridmap, global_grid.resolution, global_grid.map_size)
				ax_in.plot(pedestrian_trajectory[t:, 0], pedestrian_trajectory[t:, 1], color='red', label='Agent Real Traj')
				robot_pos = traj[t]["robot_state"]
				ax_in.plot(robot_pos[0], robot_pos[1],
				           color='g', marker='o', label='Robot')
				agent_pos = traj[t]["pedestrian_state"]["position"]
				ax_in.plot(agent_pos[0], agent_pos[1],
				           color='red', marker='o', label='Agent')

				if grid == None:
					gridmap = np.zeros((int((x_lim[1] - x_lim[0]) / map_resolution),
					                   int((y_lim[1] - y_lim[0]) / map_resolution)))

				sup.add_other_agents_to_grid(gridmap,robot_pos,args)

				sup.plot_grid(ax_in, np.array([0.0, 0.0]), gridmap, map_resolution,
				              np.array((10,10)))

				ax_in.plot(pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 0], pedestrian_trajectory[t+args.prev_horizon:t+args.prev_horizon+args.prediction_horizon, 1], color='blue',marker='o', label='Agent Real')

				x0 = pedestrian_trajectory[t+args.prev_horizon,0]
				y0 = pedestrian_trajectory[t+args.prev_horizon,1]
				pred_vel = np.zeros((args.prediction_horizon,args.output_dim))
				for i in range(args.prediction_horizon):
					pred_vel[i,:] = predictions[t][i*args.output_pred_state_dim:i*args.output_pred_state_dim+2]
				traj_pred = sup.path_from_vel(initial_pos=np.array([x0 , y0]),pred_vel=pred_vel, dt=args.dt)
				ax_in.plot(traj_pred[:, 0], traj_pred[:, 1], color='k', marker='x', label='Agent Pred Traj')

				sigma_x = np.square(predictions[t][2]) * (args.dt **2)
				sigma_y = np.square(predictions[t][3]) * (args.dt **2)

				e1 = Ellipse(xy=(traj_pred[0, 0], traj_pred[0, 1]), width=np.sqrt(sigma_x) / 2,
				             height=np.sqrt(sigma_y) / 2, angle=0 / np.pi * 180)
				e1.set_alpha(0.5)
				ax_in.add_patch(e1)

				for pred_step in range(args.prediction_horizon):
					for mix_idx in range(args.n_mixtures):
						sigma_x += np.square(predictions[t][args.n_mixtures * pred_step * args.output_pred_state_dim +
						                                          mix_idx * args.output_pred_state_dim +2]) * (args.dt **2)
						sigma_y += np.square(predictions[t][args.n_mixtures * pred_step * args.output_pred_state_dim +
						                                          mix_idx * args.output_pred_state_dim +3])* (args.dt **2)

						e1 = Ellipse(xy=(traj_pred[pred_step, 0], traj_pred[pred_step, 1]), width=np.sqrt(sigma_x) / 2,
						             height=np.sqrt(sigma_y) / 2,
						             angle=0 / np.pi * 180)
						ax_in.add_patch(e1)

				ax_in.set_xlim(x_lim)
				ax_in.set_ylim(y_lim)

				ax_in.legend()
				fig.canvas.draw()
				pl.show(block=False)
				time.sleep(0.1)
				writer.grab_frame()

def plot_batch(batch_x, batch_grid, batch_ped_grid, batch_y, other_agents_pos,args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""

	fig = pl.figure("Batch Data")
	ax_in = pl.subplot()

	ax_in.clear()

    #x_lim = [-args.map_size[0]*args.resolution/2 ,args.map_size[0]*args.resolution/2]
	#y_lim = [-args.map_size[1]*args.resolution/2,args.map_size[1]*args.resolution/2]
	x_lim = [-args.grid_width*args.submap_resolution/2 ,args.grid_width*args.submap_resolution/2]
	y_lim = [-args.grid_height*args.submap_resolution/2,args.grid_height*args.submap_resolution/2]
	x_lim = [-5,5]
	y_lim = [-5, 5]
	ax_in.clear()

	for batch_id in range(batch_x.shape[0]):
		for t_step in range(batch_x.shape[1]):
			ax_in.clear()
			# plot scenario grid
			sup.plot_grid(ax_in, np.array([0, 0]), batch_grid[batch_id,t_step,:,:],
			              args.submap_resolution, np.array((args.grid_width,args.grid_height))*args.submap_resolution)

			robot_pos = other_agents_pos[batch_id,t_step,:]
			agent_pos = batch_x[batch_id, t_step, :2]
			ax_in.plot(robot_pos[0]-agent_pos[0], robot_pos[1]-agent_pos[1],
				           color='g', marker='o', label='Robot')

			ax_in.plot(0, 0,
				           color='red', marker='o', label='Agent')

			ax_in.set_xlim(x_lim)
			ax_in.set_ylim(y_lim)

			ax_in.legend()
			fig.canvas.draw()
			pl.show(block=False)

def plot_batch_OpenCV(step,batch_x, batch_grid, batch_ped_grid, batch_y, other_agents_pos,_model_prediction, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	homography_file = os.path.join(args.data_path+ args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')

	video_file = args.model_path +"/"  + args.model_name+ "_" + str(step) + '_training.avi'
	img_file = args.data_path+ args.scenario + '/map.png'
	font = cv2.FONT_HERSHEY_SIMPLEX
	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = cv2.imread(img_file)

		height, width, channels = img.shape

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		recorder = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width, height))

		# Adding legend
		cv2.putText(img, str(step), (30, 30), font, 1, (0, 0, 0), 2, cv2.LINE_AA)
		cv2.line(img, (850, 30), (900, 30), (0, 0, 0), 4)
		cv2.putText(img, "Robot Pos", (910, 30), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)
		cv2.line(img, (850, 50), (900, 50), (255, 0, 0), 4)
		cv2.putText(img, "Pedestrian", (910, 50), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
		cv2.line(img, (850, 70), (900, 70), (255, 255, 0), 4)
		cv2.putText(img, "Pedestrian traj", (910, 70), font, 0.5, (255, 255, 0), 2, cv2.LINE_AA)
		cv2.line(img, (850, 90), (900, 90), (128,128,128), 4)
		cv2.putText(img, "Pred traj 1", (910, 90), font, 0.5, (128,128,128), 2, cv2.LINE_AA)
		cv2.line(img, (850, 110), (900, 110), (0,128,255), 4)
		cv2.putText(img, "Pred traj 2", (910, 110), font, 0.5, (0,128,255), 2, cv2.LINE_AA)
		cv2.line(img, (850, 130), (900, 130), (255,0,127), 4)
		cv2.putText(img, "Pred traj 3", (910, 130), font, 0.5, (255,0,127), 2, cv2.LINE_AA)

		for batch_idx in range(batch_x.shape[0]):
			if "keras" in args.model_name:
				predictions = _model_prediction[batch_idx]
			else:
				predictions = np.transpose(_model_prediction,(1,0,2))[batch_idx]
			ped_vel = batch_y[batch_idx]

			for t in range(predictions.shape[0]):
				overlay_ellipses = img.copy()

				# Plot Robot Positions
				robot_pos = batch_x[batch_idx,t,:2]
				robot = np.expand_dims(np.array((robot_pos[0], robot_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, robot)
				sigma_x = 0.4
				sigma_y = 0.4
				axis = (
					int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
				center = (obsv_XY[0, 1], obsv_XY[0, 0])
				cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (0, 0, 0), -1)

				# Plot ped real position
				agent_pos = other_agents_pos[batch_idx][t]
				for agent_id in range(agent_pos.shape[0]):
					ped_pos = np.expand_dims(np.array((agent_pos[agent_id,0], agent_pos[agent_id,1])), axis=0).astype(float)
					obsv_XY = sup.to_image_frame(Hinv, ped_pos)
					axis = (
						int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
					center = (obsv_XY[0, 1], obsv_XY[0, 0])
					cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, (255, 0, 0), -1)

				#Pedestrian Predicted Real TRajectory
				x0 = batch_x[batch_idx, t, 0]
				y0 = batch_x[batch_idx, t, 1]
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.truncated_backprop_length):
					pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_trajectory = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pedestrian_velocity, dt=args.dt)
				obsv_XY = sup.to_image_frame(Hinv, pedestrian_trajectory)
				sup.line_cv(overlay_ellipses, obsv_XY, (255, 255, 0), 3)  # bgr convention

				# Plot Predictions
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				# orange grey and purple
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				for mix_id in range(args.n_mixtures):
					for pred_step in range(args.prediction_horizon):
						idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
						pred_vel[pred_step, :] = np.array((predictions[t][idx],predictions[t][idx+args.n_mixtures]
						                                   ))
						if args.output_pred_state_dim > args.output_dim:
							sigmax[pred_step, :] = predictions[t][idx + 2*args.n_mixtures]
							sigmay[pred_step, :] = predictions[t][idx + 3*args.n_mixtures]
					traj_pred = sup.path_from_vel(initial_pos=np.array([x0, y0]), pred_vel=pred_vel, dt=args.dt)
					obsv_XY = sup.to_image_frame(Hinv, traj_pred)
					sup.line_cv(overlay_ellipses, obsv_XY, colors[mix_id], 3)  # bgr convention

					if args.output_pred_state_dim > args.output_dim:
						#Propagate uncertainty and plot
						sigma_x = np.square(sigmax[0]) * (args.dt ** 2)
						sigma_y = np.square(sigmay[0]) * (args.dt ** 2)

						obsv_XY = sup.to_image_frame(Hinv, traj_pred)
						axis = (
							int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
						center = (obsv_XY[0, 1], obsv_XY[0, 0])
						cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], -1)

						for pred_step in range(1,args.prediction_horizon):
							sigma_x += np.square(sigmax[pred_step]) * (args.dt ** 2)
							sigma_y += np.square(sigmay[pred_step]) * (args.dt ** 2)
							axis = (
								int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
							center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
							cv2.ellipse(overlay_ellipses, center, axis, 0, 0, 360, colors[mix_id], -1)

				image_new = cv2.addWeighted(img, 0.2, overlay_ellipses, 0.8, 0)
				width = 1200
				height = 1200
				dim = (width, height)
				# resize image
				resized_img = cv2.resize(image_new, dim, interpolation=cv2.INTER_AREA)
				if False:
					cv2.imshow("image",resized_img)
					cv2.waitKey(100)
				recorder.write(resized_img)
		recorder.release()

def plot_local_scenario_vel_OpenCV2(trajectories, batch_target, all_predictions, args, exp_num=0, display = True):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
		"""
	homography_file = os.path.join(args.data_path + args.scenario, 'H.txt')
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')
	scenario = args.scenario
	video_file = args.model_path +"/"  + args.model_name+ "_" + str(exp_num) + '_local.avi'
	img_file = args.data_path + args.scenario +'/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = ~cv2.imread(img_file)

		scale_percent = 100  # percent of original size
		width = 300
		height = 300
		dim = (width, height)
		# resize image
		resized_img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)

		# Define the codec and create VideoWriter object.The output is stored in 'outpy.avi' file.
		out = cv2.VideoWriter(video_file,
		                      cv2.VideoWriter_fourcc('M', 'P', '4', '2'), 10, (width, height))

		# Adding legend
		font = cv2.FONT_HERSHEY_SIMPLEX

		for exp_idx in range(len(all_predictions)):
			predictions = all_predictions[exp_idx]
			traj = trajectories[exp_idx]
			ped_vel = batch_target[exp_idx][0]

			for t in range(len(predictions)):
				pedestrian_trajectory = traj.pose_vec[:,:2]

				# Plot Query Agent Positions
				query_agent_pos = traj.pose_vec[t,:2]
				query_agent = np.expand_dims(np.array((query_agent_pos[0], query_agent_pos[1])), axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, query_agent)*int(scale_percent / 100)
				sigma_x = 0.4
				sigma_y = 0.4
				axis = (
					int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
				query_agent_center = np.array([obsv_XY[0, 0], obsv_XY[0, 1]])
				axis_pos = query_agent_center - np.array([args.submap_width/args.submap_resolution/2,args.submap_height/args.submap_resolution/2])
				overlay_ellipses = resized_img.copy()
				# Crop image for obtain surrounding static environment
				crop_img = cv2.getRectSubPix(overlay_ellipses,(int(args.submap_height / args.submap_resolution),int(args.submap_width / args.submap_resolution)),(query_agent_center[0],query_agent_center[1]))

				# resize image
				resized_crop_img = cv2.resize(crop_img, dim, interpolation=cv2.INTER_AREA)

				cv2.line(resized_crop_img, (width - 150, 30), (width - 130, 30), (0, 255, 0), 4)
				cv2.putText(resized_crop_img, "Query Agent", (width - 120, 30), font, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 50), (width - 130, 50), (255, 0, 0), 4)
				cv2.putText(resized_crop_img, "Other Agents", (width - 120, 50), font, 0.5, (255, 0, 0), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 70), (width - 130, 70), (255, 255, 0), 4)
				cv2.putText(resized_crop_img, "Query Agent traj", (width - 120, 70), font, 0.5, (255, 255, 0), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 90), (width - 130, 90), (128, 128, 128), 4)
				cv2.putText(resized_crop_img, "Pred traj 1", (width - 120, 90), font, 0.5, (128, 128, 128), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 110), (width - 130, 110), (0, 128, 255), 4)
				cv2.putText(resized_crop_img, "Pred traj 2", (width - 120, 110), font, 0.5, (0, 128, 255), 2, cv2.LINE_AA)
				cv2.line(resized_crop_img, (width - 150, 130), (width - 130, 130), (255, 0, 127), 4)
				cv2.putText(resized_crop_img, "Pred traj 3", (width - 120, 130), font, 0.5, (255, 0, 127), 2, cv2.LINE_AA)

				# Plot Query Agent Positions
				center = (dim[0]//2, dim[1]//2)
				cv2.ellipse(resized_crop_img, center, axis, 0, 0, 360, (0, 255, 0), -1)

				# Plot other pedestrians
				agent_pos = traj.other_agents_positions[t]
				for agent_id in range(agent_pos.shape[0]):
					ped_pos = np.expand_dims(np.array((agent_pos[agent_id, 0], agent_pos[agent_id, 1])), axis=0).astype(float)
					obsv_XY = sup.to_image_frame(Hinv, ped_pos)*int(scale_percent / 100)
					axis = (
						int(np.sqrt(sigma_x) / args.submap_resolution / scale_percent * 100),
						int(np.sqrt(sigma_y) / args.submap_resolution / scale_percent * 100))
					center = (np.array([obsv_XY[0, 0], obsv_XY[0, 1]]) - axis_pos)
					if (center[0]>0) & (center[0]<width) & (center[1]>0) & (center[1]<height):
						agent_center = (int(center[0]),int(center[1]))
						cv2.ellipse(resized_crop_img, agent_center, axis, 0, 0, 360, (0, 0, 0), -1)

				# Query-agent real trajectory
				obsv_XY = (sup.to_image_frame(Hinv, pedestrian_trajectory[t:, :])- axis_pos)
				obsv_XY_filtered = []
				for i in range(obsv_XY.shape[0]):
					if (obsv_XY[i,0]>0) & (obsv_XY[i,0]<dim[0]) & (obsv_XY[i,1]>0) & (obsv_XY[i,1]<dim[1]):
						obsv_XY_filtered.append(obsv_XY[i])
				sup.line_cv(resized_crop_img, np.asarray(obsv_XY_filtered), (255, 0, 0), 1)  # bgr convention

				#Query-agent Predicted Real TRajectory
				pedestrian_velocity = np.zeros((args.prediction_horizon, args.output_dim))
				for t_idx in range(args.prediction_horizon):
					pedestrian_velocity[t_idx, :] = ped_vel[t, t_idx * 2:t_idx * 2 + 2]
				pedestrian_trajectory = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=pedestrian_velocity, dt=args.dt)/args.submap_resolution
				obsv_XY = (sup.to_image_frame(Hinv, pedestrian_trajectory)- axis_pos)*int(scale_percent / 100)
				obsv_XY_filtered = []
				for i in range(obsv_XY.shape[0]):
					if (obsv_XY[i, 0] > 0) & (obsv_XY[i, 0] < dim[0]) & (obsv_XY[i, 1] > 0) & (obsv_XY[i, 1] < dim[1]):
						obsv_XY_filtered.append(obsv_XY[i])
				sup.line_cv(resized_crop_img, np.asarray(obsv_XY_filtered), (255, 255, 0), 1)  # bgr convention

				# Plot Predictions
				pred_vel = np.zeros((args.prediction_horizon, args.output_dim))
				sigmax = np.zeros((args.prediction_horizon, 1))
				sigmay = np.zeros((args.prediction_horizon, 1))
				colors=[(128,128,128),(0,128,255),(255,0,127)]
				""""""
				for mix_id in range(args.n_mixtures):
					for pred_step in range(args.prediction_horizon):
						idx = pred_step * args.output_pred_state_dim * args.n_mixtures + mix_id
						pred_vel[pred_step, :] = np.array((predictions[t][idx],predictions[t][idx+args.n_mixtures]
						                                   ))
						if args.output_pred_state_dim > args.output_dim:
							sigmax[pred_step, :] = predictions[t][idx + 2*args.n_mixtures]
							sigmay[pred_step, :] = predictions[t][idx + 3*args.n_mixtures]
					traj_pred = sup.path_from_vel(initial_pos=np.array([0, 0]), pred_vel=pred_vel, dt=args.dt)/args.submap_resolution
					obsv_XY = (sup.to_image_frame(Hinv, traj_pred)- axis_pos)*int(scale_percent / 100)
					obsv_XY_filtered = []
					for i in range(obsv_XY.shape[0]):
						if (obsv_XY[i, 0] > 0) & (obsv_XY[i, 0] < dim[0]) & (obsv_XY[i, 1] > 0) & (obsv_XY[i, 1] < dim[1]):
							obsv_XY_filtered.append(obsv_XY[i])
					sup.line_cv(resized_crop_img, np.asarray(obsv_XY_filtered), colors[mix_id], 1)  # bgr convention

					if args.output_pred_state_dim > args.output_dim:
						sigma_x = np.square(sigmax[0]) * (args.dt ** 2)/ args.submap_resolution
						sigma_y = np.square(sigmay[0]) * (args.dt ** 2)/ args.submap_resolution

						obsv_XY = sup.to_image_frame(Hinv, traj_pred)*int(scale_percent / 100)
						axis = (
							int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
						if (obsv_XY[0, 0] > 0) & (obsv_XY[0, 0] < dim[0]) & (obsv_XY[0, 1] > 0) & (obsv_XY[0, 1] < dim[1]):
							center = (obsv_XY[0, 1], obsv_XY[0, 0])
							cv2.ellipse(resized_crop_img, center, axis, 0, 0, 360, colors[mix_id], -1)

						for pred_step in range(1,args.prediction_horizon):
							sigma_x += np.square(sigmax[pred_step]) * (args.dt ** 2)/ args.submap_resolution
							sigma_y += np.square(sigmay[pred_step]) * (args.dt ** 2)/ args.submap_resolution
							axis = (
								int(np.sqrt(sigma_x) / args.submap_resolution), int(np.sqrt(sigma_y) / args.submap_resolution))
							if (obsv_XY[pred_step, 0] > 0) & (obsv_XY[pred_step, 0] < dim[0]) & (obsv_XY[pred_step, 1] > 0) & (obsv_XY[pred_step, 1] < dim[1]):
								center = (obsv_XY[pred_step, 1], obsv_XY[pred_step, 0])
								cv2.ellipse(resized_crop_img, center, axis, 0, 0, 360, colors[mix_id], -1)

				if display:
					cv2.imshow("image", resized_crop_img)
					cv2.waitKey(10)
				out.write(resized_crop_img)

		# When everything done, release the video capture and video write objects
		out.release()

	print("Done recording...")

def plot_dataset(trajs, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
	"""
	homography_file = os.path.join(args.data_path + args.scenario, 'H.txt')
	print("Looking for homography at: " + homography_file)
	if os.path.exists(homography_file):
		Hinv = np.linalg.inv(np.loadtxt(homography_file))
	else:
		print('[INF] No homography file')

	img_file = args.data_path + args.scenario +'/map.png'

	if os.path.exists(img_file):
		print('[INF] Using image file ' + img_file)
		img = cv2.imread(img_file)
		map_resolution = 1/Hinv[0,1]
		height, width, channels = img.shape

		# Adding legend

		font = cv2.FONT_HERSHEY_SIMPLEX

		for exp_idx in range(len(trajs)):
			overlay_img = img.copy()
			save_img_to_file = args.data_path + "/" + args.scenario + "/" + str(exp_idx) + '.png'

			cv2.line(overlay_img, (width - 150, 30), (width - 130, 30), (0, 0, 0), 4)
			cv2.putText(overlay_img, "Agent", (width - 120, 30), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)
			cv2.line(overlay_img, (width - 150, 50), (width - 130, 50), (255, 0, 0), 4)
			cv2.putText(overlay_img, "Robot", (width - 120, 50), font, 0.5, (0, 0, 0), 2, cv2.LINE_AA)

			sigma_x = 1.0
			sigma_y = 1.0
			axis = (
				int(np.sqrt(sigma_x) / map_resolution), int(np.sqrt(sigma_y) / map_resolution))

			for t in range(len(trajs[exp_idx])):
				dict = trajs[exp_idx][t]
				robot_pos = np.expand_dims(dict["robot_state"][:2],axis=0).astype(float)
				ped_pos = np.expand_dims(dict["pedestrian_state"]["position"], axis=0).astype(float)
				obsv_XY = sup.to_image_frame(Hinv, robot_pos)
				center = (obsv_XY[0, 0], obsv_XY[0,1])
				cv2.ellipse(overlay_img, center, axis, 0, 0, 360, (0, 0, 0), -1)

				obsv_XY = sup.to_image_frame(Hinv, ped_pos)
				center = (obsv_XY[0, 0], obsv_XY[0, 1])
				cv2.ellipse(overlay_img, center, axis, 0, 0, 360, (255, 0, 0), -1)

			if True:
				cv2.imshow("image",overlay_img)
				cv2.waitKey(100)
			cv2.imwrite(save_img_to_file,overlay_img)

def matplot_dataset(trajs, args):
	"""
		inputs:
			enc_seq: Encoder Inouts Sequences
			dec_seq: Decoder input sequences
			traj: Trajectory Groud Truth of Robot and Pedestrians
			predictions: Pedestrians Predicted Trajectory
			global_grid: Static Environment
	"""
	fig = pl.figure('Training Performance ' + args.model_name)
	ax = pl.subplot()
	for exp_idx in range(len(trajs)):
		ax.clear()
		save_img_to_file = args.data_path + "/" + args.scenario + "/" + str(exp_idx) + '.png'

		sigma_x = 1.0
		sigma_y = 1.0
		ax.set_aspect('equal')
		ax.set_xlabel('x [m]')
		ax.set_ylabel('y [m]')
		robot_traj = np.zeros((len(trajs[exp_idx]),2))
		ped_traj = np.zeros((len(trajs[exp_idx]), 2))
		for t in range(len(trajs[exp_idx])):
			dict = trajs[exp_idx][t]
			robot_pos = dict["robot_state"][:2]
			ped_pos = dict["pedestrian_state"]["position"]
			robot_traj[t, :] = robot_pos
			ped_traj[t,:] = ped_pos

			e1 = Ellipse(xy=(robot_pos[0], robot_pos[1]), width=sigma_x / 2, height=sigma_y / 2,
				             angle=0 / np.pi * 180)
			e1.set_alpha(0.5)

			ax.add_patch(e1)

			e1 = Ellipse(xy=(ped_pos[0], ped_pos[1]), width=sigma_x / 2, height=sigma_y / 2,
				             angle=0 / np.pi * 180,facecolor='orange')
			e1.set_alpha(0.5)

			ax.add_patch(e1)
			pl.show(block=False)
		min_x = min(min(robot_traj[:,0]),min(ped_traj[:,0]))-1.0
		max_x = max(max(robot_traj[:, 0]), max(ped_traj[:, 0])) + 1.0
		min_y = min(min(robot_traj[:, 1]), min(ped_traj[:, 1])) - 1.0
		max_y = max(max(robot_traj[:, 1]), max(ped_traj[:, 1])) + 1.0
		ax.set_xlim([min_x, max_x])
		ax.set_ylim([min_y, max_y])
		fig.savefig(save_img_to_file)

src/data_utils/plot_utils.py
eliatrevisan/IA-MPPI-LBM
import os
import csv
import sys
import pickle as pkl
import argparse

sys.path.append('../src/data_utils')
if sys.version_info[0] < 3:
	import Support as sup
else:
	import src.data_utils.Support as sup

sys.path.append('../src/data_utils')
class bcolors:
	HEADER = '\033[95m'
	OKBLUE = '\033[94m'
	OKGREEN = '\033[92m'
	WARNING = '\033[93m'
	FAIL = '\033[91m'
	ENDC = '\033[0m'
	BOLD = '\033[1m'
	UNDERLINE = '\033[4m'


def write_summary(training_loss, args):
	if not os.path.isfile("trained_models_summary.csv"):
		with open("trained_models_summary.csv", 'w') as csvfile:
			# Write header
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			writer.writerow(["Model name", "Dataset", "Ped Info Size",
			                 "Batch size", "Training loss", "MSE", "FDE",
			                 "N epochs", "Learning rate", "Grad clip", "N mixtures","Test Dataset"])
			writer.writerow(
				[args.model_name, args.scenario, args.pedestrian_vector_dim, args.batch_size,
				 training_loss, 0, 0, args.n_epochs,
				 args.learning_rate_init, args.grads_clip, args.n_mixtures])
	else:
		with open("trained_models_summary.csv", 'r') as readFile:
			reader = csv.reader(readFile)
			lines = list(reader)
		with open("trained_models_summary.csv", 'w') as csvfile:
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			if len(lines) > args.exp_num + 1:
				for i in range(len(lines)):
					if i == args.exp_num + 1:
						writer.writerow(
							[args.model_name + "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size,
							 training_loss, 0, 0, args.n_epochs,
							 args.learning_rate_init, args.grads_clip, args.n_mixtures])
					else:
						writer.writerow(lines[i])
			else:
				writer.writerows(lines)
				writer.writerow(
					[args.model_name, args.scenario, args.pedestrian_vector_dim, args.batch_size,
					 training_loss, 0, 0, args.n_epochs,
					 args.learning_rate_init, args.grads_clip, args.n_mixtures])


def write_results_summary(mse, fde, avg_div, args, test_args, cvmse, cvfde):
	if not os.path.isfile(test_args.model_name + "_summary.csv"):
		with open(test_args.model_name + "_summary.csv", 'w') as csvfile:
			# Write header
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			writer.writerow(["Model name", "Dataset", "Ped Info Size",
			                 "Batch size", "MSE", "FDE","Diversity",
			                 "dt", "Prediction Horizon", "Previous Steps", "tbpt","Test Dataset","Others Info", "CV-MSE", "CV-FDE"])
			writer.writerow(
				[args.model_name+ "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size,
				 mse, fde,avg_div, args.dt,
				 args.prediction_horizon, args.prev_horizon, args.truncated_backprop_length,test_args.scenario,args.others_info, cvmse, cvfde])

	else:
		with open(test_args.model_name + "_summary.csv", 'r') as readFile:
			reader = csv.reader(readFile)
			lines = list(reader)
		with open(test_args.model_name + "_summary.csv", 'w') as csvfile:
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			if len(lines) > 3000:#args.exp_num + 1:
				for i in range(len(lines)):
					if i == args.exp_num + 1:
						writer.writerow(
							[args.model_name + "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size,
							 mse, fde, avg_div,args.dt,
							 args.prediction_horizon, args.prev_horizon, args.truncated_backprop_length, test_args.scenario,args.others_info, cvmse, cvfde])
					else:
						writer.writerow(lines[i])
			else:
				writer.writerows(lines)
				writer.writerow(
					[args.model_name + "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size,
					 mse, fde,avg_div, args.dt,
					 args.prediction_horizon, args.prev_horizon, args.truncated_backprop_length, test_args.scenario,args.others_info, cvmse, cvfde])

def write_keras_results_summary(loss, args):
	if not os.path.isfile(args.model_name + "_summary.csv"):
		with open(args.model_name + "_summary.csv", 'w') as csvfile:
			# Write header
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			writer.writerow(["Model name", "Dataset", "Ped Info Size", "Batch size", "Number of steps", "Loss"])
			writer.writerow( [args.model_name+ "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size, args.n_steps, loss] )

	else:
		with open(args.model_name + "_summary.csv", 'r') as readFile:
			reader = csv.reader(readFile)
			lines = list(reader)
		with open(args.model_name + "_summary.csv", 'w') as csvfile:
			writer = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_MINIMAL)
			if len(lines) > 3000:#args.exp_num + 1:
				for i in range(len(lines)):
					if i == args.exp_num + 1:
						writer.writerow( [args.model_name+ "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size, args.n_steps, loss] )
					else:
						writer.writerow(lines[i])
			else:
				writer.writerows(lines)
				writer.writerow( [args.model_name+ "_" + str(args.exp_num), args.scenario, args.pedestrian_vector_dim, args.batch_size, args.n_steps, loss] )


def convert_dataset_python_version(pwd,pwd_out):
	print("Loading data from: '{}'".format(pwd))
	file = open(pwd, 'rb')
	tmp_self = pkl.load(file, encoding='latin1')
	file.close()
	file = open(pwd_out, 'wb')
	pkl.dump(tmp_self, file, protocol=2)


def model_selector(args):
	if "modelKerasRNN_arbitraryAgents" in args.model_name:
		if sys.version_info[0] < 3:
			from modelKerasRNN_arbitraryAgents import NetworkModel
		else:
			from src.models.modelKerasRNN_arbitraryAgents import NetworkModel
	elif "IntNet" in args.model_name:
		if sys.version_info[0] < 3:
			from IntNet import NetworkModel
		else:
			from src.models.IntNet import NetworkModel
	elif "Grid" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_Grid import NetworkModel
		else:
			from src.models.VGDNN_Grid import NetworkModel
	elif "no_grid" in args.model_name:
		if sys.version_info[0] < 3:
			from RNN_no_grid_goal import NetworkModel
		else:
			from src.models.RNN_no_grid_goal import NetworkModel
	elif "goal" in args.model_name:
		if sys.version_info[0] < 3:
			from RNN_goal import NetworkModel
		else:
			from src.models.RNN_goal import NetworkModel
	elif "simple" in args.model_name:
		if sys.version_info[0] < 3:
			from RNN_simple import NetworkModel
		else:
			from src.models.RNN_simple import NetworkModel
	elif "RNN" in args.model_name:
		if sys.version_info[0] < 3:
			from RNN import NetworkModel
		else:
			from src.models.RNN import NetworkModel
	elif "KL" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_KL import NetworkModel
		else:
			from src.models.VGDNN_KL import NetworkModel
	elif "diversity" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_attention_diversity import NetworkModel
		else:
			from src.models.VGDNN_attention_diversity import NetworkModel
	elif "ped" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_ped_attention import NetworkModel
		else:
			from src.models.VGDNN_ped_attention import NetworkModel
	elif "attention" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_attention import NetworkModel
		else:
			from src.models.VGDNN_attention import NetworkModel
	elif "multihead" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_multihead_attention import NetworkModel
		else:
			from src.models.VGDNN_multihead_attention import NetworkModel
	elif "simple" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_simple import NetworkModel
		else:
			from src.models.VGDNN_simple import NetworkModel
	elif "pos" in args.model_name:
		if sys.version_info[0] < 3:
			from VGDNN_pos import NetworkModel
		else:
			from src.models.VGDNN_pos import NetworkModel
	else:
		if sys.version_info[0] < 3:
			from VGDNN import NetworkModel
		else:
			from src.models.VGDNN import NetworkModel

	return NetworkModel(args)

def parse_args(defaults, parse_type):
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name', help='Path to directory that comprises the model (default="model_name").', type=str, default=defaults['model_name'])
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=defaults['exp_num'])
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").', type=str, default=defaults['scenario'])
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str, default=defaults['data_path'])

	if parse_type == "Train" or parse_type == "train":
		parser.add_argument('--input1_type', help='Data type for the first input', type=str, default=defaults['input1_type'])
		parser.add_argument('--input2_type', help='Data type for the second input', type=str, default=defaults['input2_type'])
		parser.add_argument('--model_path', help='Path to directory to save the model (default=""../trained_models/"+model_name").', type=str, default=defaults['model_path'])
		parser.add_argument('--pretrained_convnet_path', help='Path to directory that comprises the pre-trained convnet model (default=" ").', type=str, default=defaults['pretrained_convnet_path'])
		parser.add_argument('--log_dir', help='Path to the log directory of the model (default=""../trained_models/"+model_name").', type=str, default=defaults['log_dir'])
		parser.add_argument('--scenario_val', help='Dataset pkl file', type=str, default= defaults['scenario_val'])
		parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool, default=defaults['real_world_data'])
		parser.add_argument('--dataset', help='Dataset pkl file', type=str, default= defaults['scenario'] + '.pkl')
		parser.add_argument('--dataset_val', help='Dataset pkl file', type=str, default= defaults['scenario_val'] + '.pkl')
		parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str, default='LSTM')
		parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool, default=defaults['warmstart_model'])
		parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool, default=defaults['warm_start_convnet'])
		parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float, default=defaults['dt'])
		parser.add_argument('--n_steps', help='Number of epochs (default=10000).', type=int, default=defaults['n_steps'])
		parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=defaults['batch_size'])
		parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float, default=defaults['regularization_weight'])
		parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float, default=defaults['keep_prob'])
		parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float, default=defaults['learning_rate_init'])
		parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float, default=defaults['beta_rate_init'])
		parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool, default=defaults['dropout'])
		parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float, default=defaults['grads_clip'])
		parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).', type=int, default=defaults['truncated_backprop_length'])
		parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int, default=defaults['prediction_horizon'])
		parser.add_argument('--prev_horizon', help='Previous seq length.', type=int, default=defaults['prev_horizon'])
		parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int, default=defaults['rnn_state_size'])
		parser.add_argument('--rnn_state_size_lstm_ped', help='Number of RNN / LSTM units of the pedestrian lstm layer (default=32).', type=int, default=defaults['rnn_state_size_lstm_ped'])
		parser.add_argument('--rnn_state_size_bilstm_ped', help='Number of RNN / LSTM units of the pedestrian bidirectional lstm layer (default=32).',type=int, default=defaults['rnn_state_size_bilstm_ped'])
		parser.add_argument('--rnn_state_size_lstm_grid', help='Number of RNN / LSTM units of the grid lstm layer (default=32).', type=int, default=defaults['rnn_state_size_lstm_grid'])
		parser.add_argument('--rnn_state_size_lstm_concat', help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).', type=int, default=defaults['rnn_state_size_lstm_concat'])
		parser.add_argument('--prior_size', help='prior_size', type=int, default=defaults['prior_size'])
		parser.add_argument('--latent_space_size', help='latent_space_size', type=int, default=defaults['latent_space_size'])
		parser.add_argument('--x_dim', help='x_dim', type=int, default=defaults['x_dim'])
		parser.add_argument('--fc_hidden_unit_size', help='Number of fully connected layer units after LSTM layer (default=64).', type=int, default=defaults['fc_hidden_unit_size'])
		parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int, default=defaults['input_state_dim'])
		parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float, default=defaults['input_dim'])
		parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float, default=defaults['output_dim'])
		parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int, default=defaults['output_pred_state_dim'])
		parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int, default=defaults['cmd_vector_dim'])
		parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int, default=defaults['n_mixtures'])
		parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int, default=defaults['pedestrian_vector_dim'])
		parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int, default=defaults['pedestrian_vector_state_dim'])
		parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float, default=defaults['max_range_ped_grid'])
		parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float, default=defaults['pedestrian_radius'])
		parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=int, default=0)
		parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int, default=defaults['print_freq'])
		parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).', type=int, default=defaults['save_freq'])
		parser.add_argument('--patience', help='Patience to stop the training loop early', type=int, default=defaults['patience'])
		parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool, default=False)
		parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool, default=defaults['agents_on_grid'])
		parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool, default=defaults['normalize_data'])
		parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool, default=defaults['rotated_grid'])
		parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool, default=defaults['centered_grid'])
		parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float, default=0)
		parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=defaults['submap_width'])
		parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=defaults['submap_height'])
		parser.add_argument('--submap_resolution', help='Map resolution.', type=float, default=defaults['submap_resolution'])
		parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
		parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
		parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
		parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool, default=False)
		parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool, default=defaults['predict_positions'])
		parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool, default=True)
		parser.add_argument('--relative_info', help='Use relative info for other agents.', type=sup.str2bool, default=True)
		parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool, default=defaults['regulate_log_loss'])
		parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool, default=defaults['diversity_update'])
		parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str, default='../config/topics.yaml')
		parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
		parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
		parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
		parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
		parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
		parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
		parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
		parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
		parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
		parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
		parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
		parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	else:
		parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
		parser.add_argument('--n_samples', help='Number of samples', type=int, default=1)
		parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool, default=True)
		parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool, default=False)
		parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float, default=0)
		parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float, default=5)
		parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool, default=False)
		parser.add_argument('--update_state', help='update_state', type=sup.str2bool, default=False)
		parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool, default=False)
		parser.add_argument('--freeze_other_agents', help='Freeze other agents.', type=sup.str2bool, default=False)
		parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,default=False)

	args = parser.parse_args()

	return args

if __name__ == '__main__':
	convert_dataset_python_version("/home/bbrito/Code/I-LSTM/data/2_agents_swap/trajs/GA3C-CADRL-10.pkl",
	                               "/home/bbrito/Code/I-LSTM/data/2_agents_swap/trajs/GA3C-CADRL-10-py27.pkl")

src/data_utils/utils.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
# import tensorflow_probability as tfp
import numpy as np
import os

def scaled_dot_product_attention(q, k, v, mask):
	"""Calculate the attention weights.
	q, k, v must have matching leading dimensions.
	k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
	The mask has different shapes depending on its type(padding or look ahead)
	but it must be broadcastable for addition.

	Args:
		q: query shape == (..., seq_len_q, depth)
		k: key shape == (..., seq_len_k, depth)
		v: value shape == (..., seq_len_v, depth_v)
		mask: Float tensor with shape broadcastable
					to (..., seq_len_q, seq_len_k). Defaults to None.

	Returns:
		output, attention_weights
	"""

	matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

	# scale matmul_qk
	dk = tf.cast(tf.shape(k)[-1], tf.float32)
	scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

	# add the mask to the scaled tensor.
	if mask is not None:
		scaled_attention_logits += (mask * -1e9)

	# softmax is normalized on the last axis (seq_len_k) so that the scores
	# add up to 1.
	attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

	output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

	return output, attention_weights

class MultiHeadAttention(tf.keras.layers.Layer):
	def __init__(self, d_model, num_heads):
		super(MultiHeadAttention, self).__init__()
		self.num_heads = num_heads
		self.d_model = d_model

		assert d_model % self.num_heads == 0

		self.depth = d_model // self.num_heads

		self.wq = tf.keras.layers.Dense(d_model)
		self.wk = tf.keras.layers.Dense(d_model)
		self.wv = tf.keras.layers.Dense(d_model)

		self.dense = tf.keras.layers.Dense(d_model)

	def split_heads(self, x, batch_size):
		"""Split the last dimension into (num_heads, depth).
		Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
		"""

		x = tf.reshape(x, (batch_size, 1, self.num_heads, self.depth))
		return tf.transpose(x, perm=[0, 2, 1, 3])

	def call(self, v, k, q, mask):
		batch_size = tf.shape(q)[0]

		q = self.wq(q)  # (batch_size, seq_len, d_model)
		k = self.wk(k)  # (batch_size, seq_len, d_model)
		v = self.wv(v)  # (batch_size, seq_len, d_model)


		q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
		k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
		v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

		# scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
		# attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
		scaled_attention, attention_weights = scaled_dot_product_attention(
			q, k, v, mask)

		scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

		concat_attention = tf.reshape(scaled_attention,
																	(batch_size, 1, self.d_model))  # (batch_size, seq_len_q, d_model)

		output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

		return output, attention_weights
src/external/attention.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_h = h_dim
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_x, state_is_tuple=True)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state

            #   Eq 5
            x_in = x[:,:x.shape[1]-self.args.pedestrian_vector_dim]
            ped_info = x[:,x.shape[1]-self.args.pedestrian_vector_dim:x.shape[1]]
            with tf.variable_scope("Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.exp(linear(prior_hidden, self.n_z))

            with tf.variable_scope("phi_x"):
                x_1 = tf.nn.relu(linear(x_in, self.n_x_1))
                #x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Encoder"):
                with tf.variable_scope("hidden"):
                    enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    #enc_hidden = linear(tf.concat(axis=1, values=(x_1, h)), self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.exp(linear(enc_hidden, self.n_z))
            eps = tf.random_normal(tf.shape(enc_mu), prior_mu, prior_sigma, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("phi_z"):
                z_1 = tf.nn.relu(linear(z, self.n_z_1))

            # propagate hidden state eq(7)
            output, state2 = self.lstm(tf.concat(axis=1, values=(x_1, z_1,ped_info)), state)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
src/external/irnn_cell.py
eliatrevisan/IA-MPPI-LBM
import os
import time
import matplotlib.pyplot as plt
import copy
import torch
import numpy as np
import torch.nn as nn
import torch.optim as opt
from tqdm import tqdm, trange
from itertools import chain
from torch.autograd import Variable
from torch.utils.data import DataLoader

# SOCIALWAYS: ====== training hyper-parameters ======
# Unrolled GAN
n_unrolling_steps = 0
# Info GAN
use_info_loss = True
loss_info_w = 0.5
n_latent_codes = 2
# L2 GAN
use_l2_loss = False
use_variety_loss = False
loss_l2_w = 0.5  # WARNING for both l2 and variety
# Learning Rate
lr_g = 1E-3
lr_d = 1E-4
# FIXME: ====== Network Size ===================
# Batch size
batch_size = 256
# LSTM hidden size
hidden_size = 64
n_epochs = 100000
num_social_features = 3
social_feature_size = hidden_size
noise_len = hidden_size // 2
n_lstm_layers = 1
use_social = False
n_next = 12
# ================================================

# Augment tensors of positions into positions+velocity
def get_traj_4d(obsv_p, pred_p):
    obsv_v = obsv_p[:, 1:] - obsv_p[:, :-1]
    obsv_v = torch.cat([obsv_v[:, 0].unsqueeze(1), obsv_v], dim=1)
    obsv_4d = torch.cat([obsv_p, obsv_v], dim=2)
    if len(pred_p) == 0: return obsv_4d
    pred_p_1 = torch.cat([obsv_p[:, -1].unsqueeze(1), pred_p[:, :-1]], dim=1)
    pred_v = pred_p - pred_p_1
    pred_4d = torch.cat([pred_p, pred_v], dim=2)
    return obsv_4d, pred_4d


# Evaluate the error between the model prediction and the true path
def calc_error(pred_hat, pred):
    N = pred.size(0)
    T = pred.size(1)
    err_all = torch.pow((pred_hat - pred) / ss, 2).sum(dim=2).sqrt()  # N x T
    FDEs = err_all.sum(dim=0).item() / N
    ADEs = torch.cumsum(FDEs)
    for ii in range(T):
        ADEs[ii] /= (ii + 1)
    return ADEs.data.cpu().numpy(), FDEs.data().cpu().numpy()


class AttentionPooling(nn.Module):
    def __init__(self, h_dim, f_dim):
        super(AttentionPooling, self).__init__()
        self.f_dim = f_dim
        self.h_dim = h_dim
        self.W = nn.Linear(h_dim, f_dim, bias=True)

    def forward(self, f, h, sub_batches):
        Wh = self.W(h)
        S = torch.zeros_like(h)
        for sb in sub_batches:
            N = sb[1] - sb[0]
            if N == 1: continue

            for ii in range(sb[0], sb[1]):
                fi = f[ii, sb[0]:sb[1]]
                sigma_i = torch.bmm(fi.unsqueeze(1), Wh[sb[0]:sb[1]]. unsqueeze(2))
                sigma_i[ii-sb[0]] = -1000

                attentions = torch.softmax(sigma_i.squeeze(), dim=0)
                S[ii] = torch.mm(attentions.view(1, N), h[sb[0]:sb[1]])

        return S


class EmbedSocialFeatures(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(EmbedSocialFeatures, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.fc = nn.Sequential(nn.Linear(input_size, 32), nn.ReLU(),
                                nn.Linear(32, 64), nn.ReLU(),
                                nn.Linear(64, hidden_size))

    def forward(self, ftr_list, sub_batches):
        embedded_features = self.fc(ftr_list)
        return embedded_features


def DCA(xA_4d, xB_4d):
    dp = xA_4d[:2] - xB_4d[:2]
    dv = xA_4d[2:] - xB_4d[2:]
    ttca = torch.dot(-dp, dv) / (torch.norm(dv) ** 2 + 1E-6)
    # ttca = torch.max(ttca, 0)
    dca = torch.norm(dp + ttca * dv)
    return dca


def Bearing(xA_4d, xB_4d):
    dp = xA_4d[:2] - xB_4d[:2]
    v = xA_4d[2:]
    cos_theta = torch.dot(dp, v) / (torch.norm(dp) * torch.norm(v) + 1E-6)
    return cos_theta


def DCA_MTX(x_4d, D_4d):
    Dp = D_4d[:, :, :2]
    Dv = D_4d[:, :, 2:]
    DOT_Dp_Dv = torch.mul(Dp[:,:,0], Dv[:,:,0]) + torch.mul(Dp[:,:,1], Dv[:,:,1])
    Dv_sq = torch.mul(Dv[:,:,0], Dv[:,:,0]) + torch.mul(Dv[:,:,1], Dv[:,:,1]) + 1E-6
    TTCA = -torch.div(DOT_Dp_Dv, Dv_sq)
    DCA = torch.zeros_like(Dp)
    DCA[:, :, 0] = Dp[:, :, 0] + TTCA * Dv[:, :, 0]
    DCA[:, :, 1] = Dp[:, :, 1] + TTCA * Dv[:, :, 1]
    DCA = torch.norm(DCA, dim=2)
    return DCA


def BearingMTX(x_4d, D_4d):
    Dp = D_4d[:, :, :2]  # NxNx2
    v = x_4d[:, 2:].unsqueeze(1).repeat(1, x_4d.shape[0], 1)  # => NxNx2
    DOT_Dp_v = Dp[:, :, 0] * v[:, :, 0] + Dp[:, :, 1] * v[:, :, 1]
    COS_THETA = torch.div(DOT_Dp_v, torch.norm(Dp, dim=2) * torch.norm(v, dim=2) + 1E-6)
    return COS_THETA


def SocialFeatures(x, sub_batches):
    N = x.shape[0]  # x is NxTx4 tensor

    x_ver_repeat = x[:, -1].unsqueeze(0).repeat(N, 1, 1)
    x_hor_repeat = x[:, -1].unsqueeze(1).repeat(1, N, 1)
    Dx_mat = x_hor_repeat - x_ver_repeat

    l2_dist_MTX = Dx_mat[:, :, :2].norm(dim=2)
    bearings_MTX = BearingMTX(x[:, -1], Dx_mat)
    dcas_MTX = DCA_MTX(x[:, -1], Dx_mat)
    sFeatures_MTX = torch.stack([l2_dist_MTX, bearings_MTX, dcas_MTX], dim=2)

    return sFeatures_MTX   # directly return the Social Features Matrix


# LSTM path encoding module
class EncoderLstm(nn.Module):
    def __init__(self, hidden_size, n_layers=2):
        # Dimension of the hidden state (h)
        self.hidden_size = hidden_size
        super(EncoderLstm, self).__init__()
        # Linear embedding 4xh
        self.embed = nn.Linear(4, self.hidden_size)
        # The LSTM cell.
        # Input dimension (observations mapped through embedding) is the same as the output
        self.lstm = nn.LSTM(self.hidden_size, self.hidden_size, num_layers=n_layers, batch_first=True)
        self.lstm_h = []
        # init_weights(self)

    def init_lstm(self, h, c):
        # Initialization of the LSTM: hidden state and cell state
        self.lstm_h = (h, c)

    def forward(self, obsv):
        # Batch size
        bs = obsv.shape[0]
        # Linear embedding
        obsv = self.embed(obsv)
        # Reshape and applies LSTM over a whole sequence or over one single step
        y, self.lstm_h = self.lstm(obsv.view(bs, -1, self.hidden_size), self.lstm_h)
        return y


class Discriminator(nn.Module):
    def __init__(self, n_next, hidden_dim, n_latent_code):
        super(Discriminator, self).__init__()
        self.lstm_dim = hidden_dim
        self.n_next = n_next
        # LSTM Encoder for the observed part
        self.obsv_encoder_lstm = nn.LSTM(4, hidden_dim, batch_first=True)
        # FC sub-network: input is hidden_dim, output is hidden_dim//2. This ouput will be part of
        # the input of the classifier.
        self.obsv_encoder_fc = nn.Sequential(nn.Linear(hidden_dim, hidden_dim // 2), nn.LeakyReLU(0.2),
                                             nn.Linear(hidden_dim // 2, hidden_dim // 2))
        # FC Encoder for the predicted part: input is n_next*4 (whole predicted trajectory), output is
        # hidden_dim//2. This ouput will also be part of the input of the classifier.
        self.pred_encoder = nn.Sequential(nn.Linear(n_next * 4, hidden_dim // 2), nn.LeakyReLU(0.2),
                                          nn.Linear(hidden_dim // 2, hidden_dim // 2))
        # Classifier: input is hidden_dim (concatenated encodings of observed and predicted trajectories), output is 1
        self.classifier = nn.Sequential(nn.Linear(hidden_dim, hidden_dim // 2), nn.LeakyReLU(0.2),
                                        nn.Linear(hidden_dim // 2, 1))
        # Latent code inference: input is hidden_dim (concatenated encodings of observed and predicted trajectories), output is n_latent_code (distribution of latent codes)
        self.latent_decoder = nn.Sequential(nn.Linear(hidden_dim, hidden_dim // 2), nn.LeakyReLU(0.2),
                                            nn.Linear(self.lstm_dim // 2, n_latent_code))

    def forward(self, obsv, pred):
        bs = obsv.size(0)
        lstm_h_c = (torch.zeros(1, bs, self.lstm_dim).cuda(),
                    torch.zeros(1, bs, self.lstm_dim).cuda())
        # Encoding of the observed sequence trhough an LSTM cell
        obsv_code, lstm_h_c = self.obsv_encoder_lstm(obsv, lstm_h_c)
        # Further encoding through a FC layer
        obsv_code = self.obsv_encoder_fc(obsv_code[:, -1])
        # Encoding of the predicted/next part of the sequence through a FC layer
        pred_code = self.pred_encoder(pred.view(-1, self.n_next * 4))
        both_codes = torch.cat([obsv_code, pred_code], dim=1)
        # Applies classifier to the concatenation of the encodings of both parts
        label = self.classifier(both_codes)
        # Inference on the latent code
        code_hat = self.latent_decoder(both_codes)
        return label, code_hat

    def load(self, backup):
        for m_from, m_to in zip(backup.modules(), self.modules()):
            if isinstance(m_to, nn.Linear):
                m_to.weight.data = m_from.weight.data.clone()
                if m_to.bias is not None:
                    m_to.bias.data = m_from.bias.data.clone()


# FC path decoding module
class DecoderFC(nn.Module):
    def __init__(self, hidden_dim):
        super(DecoderFC, self).__init__()
        # Fully connected sub-network. Input is hidden_dim, output is 2.
        self.fc1 = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim), nn.LeakyReLU(0.2),
                                       # torch.nn.Linear(64, 64), nn.LeakyReLU(0.2),
                                       torch.nn.Linear(hidden_dim, hidden_dim // 2), nn.LeakyReLU(0.2),
                                       torch.nn.Linear(hidden_dim // 2, hidden_dim // 4),
                                       torch.nn.Linear(hidden_dim // 4, 2))

    def forward(self, h, s, z):
        # For each sample in the batch, concatenate h (hidden state), s (social term) and z (noise)
        inp = torch.cat([h, s, z], dim=1)
        # Applies the fully connected layer
        out = self.fc1(inp)
        return out


# LSTM path decoding module
class DecoderLstm(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(DecoderLstm, self).__init__()
        # Decoding LSTM
        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True)
        # Fully connected sub-network. Input is hidden_size, output is 2.
        self.fc = nn.Sequential(torch.nn.Linear(hidden_size, 64), nn.Sigmoid(),
                                torch.nn.Linear(64, 64), nn.LeakyReLU(0.2),
                                torch.nn.Linear(64, 32), nn.LeakyReLU(0.2),
                                torch.nn.Linear(32, 2))

        # init_weights(self)
        self.lstm_h = []

    def init_lstm(self, h, c):
        # Initialization of the LSTM: hidden state and cell state
        self.lstm_h = (h, c)

    def forward(self, h, s, z):
        # Batch size
        bs = z.shape[0]
        # For each sample in the batch, concatenate h (hidden state), s (social term) and z (noise)
        inp = torch.cat([h, s, z], dim=1)
        # Applies a forward step.
        out, self.lstm_h = self.lstm(inp.unsqueeze(1), self.lstm_h)
        # Applies the fully connected layer to the LSTM output
        out = self.fc(out.squeeze())
        return out


# ===================================================


# =============== Training Loop ==================
def train():
    tic = time.clock()
    # Evaluation metrics (ADE/FDE)
    train_ADE, train_FDE = 0, 0
    batch_size_accum = 0;
    sub_batches = []
    # For all the training batches
    for ii, batch_i in enumerate(train_batches):
        batch_size_accum += batch_i[1] - batch_i[0]
        sub_batches.append(batch_i)

        # FIXME: Just keep it for toy dataset
        # sub_batches = the_batches
        # batch_size_accum = sub_batches[-1][1]
        # ii = train_size-1

        if ii >= train_size - 1 or \
                batch_size_accum + (the_batches[ii + 1][1] - the_batches[ii + 1][0]) > batch_size:
            # Observed partial paths
            obsv = dataset_obsv[sub_batches[0][0]:sub_batches[-1][1]]
            # Future partial paths
            pred = dataset_pred[sub_batches[0][0]:sub_batches[-1][1]]
            sub_batches = sub_batches - sub_batches[0][0]
            # May have to fill with 0
            filling_len = batch_size - int(batch_size_accum)
            #obsv = torch.cat((obsv, torch.zeros(filling_len, n_past, 2).cuda()), dim=0)
            #pred = torch.cat((pred, torch.zeros(filling_len, n_next, 2).cuda()), dim=0)

            bs = batch_size_accum

            # Completes the positional vectors with velocities (to have dimension 4)
            obsv_4d, pred_4d = get_traj_4d(obsv, pred)
            zeros = Variable(torch.zeros(bs, 1) + np.random.uniform(0, 0.1), requires_grad=False).cuda()
            ones = Variable(torch.ones(bs, 1) * np.random.uniform(0.9, 1.0), requires_grad=False).cuda()
            noise = torch.FloatTensor(torch.rand(bs, noise_len)).cuda()

            # ============== Train Discriminator ================
            for u in range(n_unrolling_steps + 1):
                # Zero the gradient buffers of all parameters
                D.zero_grad()
                with torch.no_grad():
                    pred_hat_4d = predict(obsv, noise, n_next, sub_batches)

                fake_labels, code_hat = D(obsv_4d, pred_hat_4d)  # classify fake samples
                # Evaluate the MSE loss: the fake_labels should be close to zero
                d_loss_fake = mse_loss(fake_labels, zeros)
                d_loss_info = mse_loss(code_hat.squeeze(), noise[:, :n_latent_codes])
                # Evaluate the MSE loss: the real should be close to one
                real_labels, code_hat = D(obsv_4d, pred_4d)  # classify real samples
                d_loss_real = mse_loss(real_labels, ones)

                #  FIXME: which loss functinos to use for D?
                d_loss = d_loss_fake + d_loss_real
                if use_info_loss:
                    #
                    d_loss += loss_info_w * d_loss_info
                d_loss.backward()  # update D
                D_optimizer.step()

                if u == 0 and n_unrolling_steps > 0:
                    backup = copy.deepcopy(D)

            # =============== Train Generator ================= #
            # Zero the gradient buffers of all the discriminator parameters
            D.zero_grad()
            # Zero the gradient buffers of all the generator parameters
            predictor_optimizer.zero_grad()
            # Applies a forward step of prediction
            pred_hat_4d = predict(obsv, noise, n_next, sub_batches)

            # Classify the generated fake sample
            gen_labels, code_hat = D(obsv_4d, pred_hat_4d)
            # L2 loss between the predicted paths and the true ones
            g_loss_l2 = mse_loss(pred_hat_4d[:, :, :2], pred)
            # Adversarial loss (classification labels should be close to one)
            g_loss_fooling = mse_loss(gen_labels, ones)
            # Information loss
            g_loss_info = mse_loss(code_hat.squeeze(), noise[:, :n_latent_codes])

            #  FIXME: which loss functions to use for G?
            #
            g_loss = g_loss_fooling
            # If using the info loss
            if use_info_loss:
                g_loss += loss_info_w * g_loss_info
            # If using the L2 loss
            if use_l2_loss:
                g_loss += loss_l2_w * g_loss_l2
            if use_variety_loss:
                KV = 20
                all_20_losses = []
                for k in range(KV):
                    pred_hat_4d = predict(obsv, noise, n_next, sub_batches)
                    loss_l2_k = mse_loss(pred_hat_4d[k, :, :2], pred[k])
                all_20_losses.append(loss_l2_k.unsqueeze(0))
                all_20_losses = torch.cat(all_20_losses)
                variety_loss, _ = torch.min(all_20_losses, dim=0)
                g_loss += loss_l2_w * variety_loss

            g_loss.backward()
            predictor_optimizer.step()

            if n_unrolling_steps > 0:
                D.load(backup)
                del backup

            # calculate error
            with torch.no_grad():  # TODO: use the function above
                err_all = torch.pow((pred_hat_4d[:, :, :2] - pred) / ss, 2)
                err_all = err_all.sum(dim=2).sqrt()
                e = err_all.sum().item() / n_next
                train_ADE += e
                train_FDE += err_all[:, -1].sum().item()

            batch_size_accum = 0;
            sub_batches = []

    train_ADE /= n_train_samples
    train_FDE /= n_train_samples
    toc = time.clock()
    print(" Epc=%4d, Train ADE,FDE = (%.3f, %.3f) | time = %.1f" \
          % (epoch, train_ADE, train_FDE, toc - tic))


def test(n_gen_samples=1, linear=False, write_to_file=None, just_one=False):
    # =========== Test error ============
    plt.close()
    ade_avg_12, fde_avg_12 = 0, 0
    ade_min_12, fde_min_12 = 0, 0
    counter = 0
    for ii, batch_i in enumerate(test_batches):
        if ii > 20: break
        obsv = dataset_obsv[batch_i[0]:batch_i[1]]
        pred = dataset_pred[batch_i[0]:batch_i[1]]
        current_t = dataset_t[batch_i[0]]
        bs = int(batch_i[1] - batch_i[0])
        with torch.no_grad():
            all_20_errors = []
            all_20_preds = []

            linear_preds = predict_cv(obsv, n_next)
            if linear and not write_to_file:
                all_20_preds.append(linear_preds.unsqueeze(0))
                err_all = torch.pow((linear_preds[:, :, :2] - pred) / ss, 2).sum(dim=2, keepdim=True).sqrt()
                all_20_errors.append(err_all.unsqueeze(0))
            else:
                for kk in range(n_gen_samples):
                    noise = torch.FloatTensor(torch.rand(bs, noise_len)).cuda()
                    pred_hat_4d = predict(obsv, noise, n_next)
                    all_20_preds.append(pred_hat_4d.unsqueeze(0))
                    err_all = torch.pow((pred_hat_4d[:, :, :2] - pred) / ss, 2).sum(dim=2, keepdim=True).sqrt()
                    all_20_errors.append(err_all.unsqueeze(0))

            all_20_errors = torch.cat(all_20_errors)
            if write_to_file:
                file_name = os.path.join(write_to_file, str(epoch) + '-' + str(current_t) + '.npz')
                print('saving to ', file_name)
                np_obsvs = scale.denormalize(obsv[:, :, :2].data.cpu().numpy())
                np_preds_our = scale.denormalize(torch.cat(all_20_preds)[:, :, :, :2].data.cpu().numpy())
                np_preds_gtt = scale.denormalize(pred[:, :, :2].data.cpu().numpy())
                np_preds_lnr = scale.denormalize(linear_preds[:, :, :2].data.cpu().numpy())
                np.savez(file_name, timestamp=current_t,
                         obsvs=np_obsvs, preds_our=np_preds_our, preds_gtt=np_preds_gtt, preds_lnr=np_preds_lnr)

            # =============== Prediction Errors ================
            dummy = all_20_errors.mean(2).mean(0, keepdim=True)
            fde_min_12_i, _ = all_20_errors[:, :, -1].min(0, keepdim=True)
            ade_min_12_i, _ = all_20_errors.mean(2).min(0, keepdim=True)
            fde_min_12 += fde_min_12_i.sum().item()
            ade_min_12 += ade_min_12_i.sum().item()
            fde_avg_12 += all_20_errors[:, :, -1].mean(0, keepdim=True).sum().item()
            ade_avg_12 += all_20_errors.mean(2).mean(0, keepdim=True).sum().item()
            counter +=dummy.shape[1]
            # ==================================================
        if just_one: break

    ade_avg_12 /= counter
    fde_avg_12 /= counter
    ade_min_12 /= counter
    fde_min_12 /= n_test_samples
    print('Avg ADE,FDE (12)= (%.3f, %.3f) | Min(20) ADE,FDE (12)= (%.3f, %.3f)' \
          % (ade_avg_12, fde_avg_12, ade_min_12, fde_min_12))


# =======================================================
# ===================== M A I N =========================
# =======================================================
src/external/socialways.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_h = h_dim
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_dec_hidden, state_is_tuple=True,reuse=tf.AUTO_REUSE)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state

            #   Eq 5
            with tf.variable_scope("Storn_Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.exp(linear(prior_hidden, self.n_z))

            with tf.variable_scope("Storn_phi_x"):
                x_1 = tf.nn.relu(linear(x, self.n_x_1))
                #x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Storn_Encoder"):
                with tf.variable_scope("hidden"):
                    enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    #enc_hidden = linear(tf.concat(axis=1, values=(x_1, h)), self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.exp(linear(enc_hidden, self.n_z))
            eps = tf.random_normal(tf.shape(enc_mu), 0.0, 1.0, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("Storn_phi_z"):
                z_1 = tf.nn.relu(linear(z, self.n_z_1))

            # propagate hidden state eq(7)
            concat = tf.concat(axis=1, values=(x_1, z_1))
            output, state2 = self.lstm(z_1, state)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
src/external/storn_cell.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_h = h_dim
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_x, state_is_tuple=True,reuse=tf.AUTO_REUSE)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state

            #   Eq 5
            with tf.variable_scope("Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.exp(linear(prior_hidden, self.n_z)) # exp() ensures the std.dev. is always positive

            with tf.variable_scope("phi_x"):
                x_1 = tf.nn.relu(linear(x, self.n_x_1))
                #x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Encoder"):
                with tf.variable_scope("hidden"):
                    enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    #enc_hidden = linear(tf.concat(axis=1, values=(x_1, h)), self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.exp(linear(enc_hidden, self.n_z)) # exp() ensures the std.dev. is always positive
            eps = tf.random.normal(tf.shape(enc_mu), prior_mu, prior_sigma, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("phi_z"):
                z_1 = tf.nn.relu(linear(z, self.n_z_1))

            # propagate hidden state eq(7)
            output, state2 = self.lstm(tf.concat(axis=1, values=(x_1, z_1)), state) # tf.concat(axis=1, values=(x_1, z_1)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
src/external/vrnn_cell.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_x, state_is_tuple=True,reuse=tf.AUTO_REUSE)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state
            x_in = x[:,:self.args.rnn_state_size]
            #   Eq 5
            with tf.variable_scope("Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.exp(linear(prior_hidden, self.n_z))

            with tf.variable_scope("phi_x"):
                x_1 = tf.nn.relu(linear(x, self.n_x_1))
                #x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Encoder"):
                with tf.variable_scope("hidden"):
                    enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    #enc_hidden = linear(tf.concat(axis=1, values=(x_1, h)), self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.exp(linear(enc_hidden, self.n_z))
            eps = tf.random.normal(tf.shape(enc_mu), prior_mu, prior_sigma, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("phi_z"):
                z_1 = tf.nn.relu(linear(z, self.n_z_1))

            # propagate hidden state eq(7)
            output, state2 = self.lstm(z_1, state) # tf.concat(axis=1, values=(x_1, z_1)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
src/external/vrnn_cell2.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import sys
if sys.version_info[0] < 3:
  from tf_utils import *
else:
  from src.models.tf_utils import *

class VariationalRNNCell(tf.contrib.rnn.RNNCell):
    """Variational RNN cell."""

    def __init__(self,args, x_dim, y_dim,h_dim ,z_dim = 100):
        self.n_h = h_dim
        self.n_x = x_dim
        self.n_z = z_dim
        self.n_x_1 = x_dim
        self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = y_dim
        self.n_prior_hidden = z_dim
        self.lstm = tf.contrib.rnn.LSTMCell(self.n_x, state_is_tuple=True)
        self.args = args

    @property
    def state_size(self):
        return (self.n_x, self.n_x)

    @property
    def output_size(self):
        return self.n_x

    def __call__(self, x, state, scope=None):
        with tf.variable_scope(scope or type(self).__name__):
            h, c = state
            x_in = x[:,:self.args.rnn_state_size]
            x_ped = x[:,self.args.rnn_state_size:]
            #   Eq 5
            with tf.variable_scope("Prior"):
                with tf.variable_scope("hidden"):
                    prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
                    #prior_hidden = linear(h, self.n_prior_hidden)
                with tf.variable_scope("mu"):
                    prior_mu = linear(prior_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    prior_sigma = tf.nn.softplus(linear(prior_hidden, self.n_z))

            with tf.variable_scope("phi_x"):
                #x_1 = tf.nn.relu(linear(x, self.n_x_1))
                x_1 = linear(x, self.n_x_1)

            with tf.variable_scope("Encoder"):
                with tf.variable_scope("hidden"):
                    #enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,values=(x_1, h)), self.n_enc_hidden))
                    enc_hidden = linear(x_in, self.n_enc_hidden)
                with tf.variable_scope("mu"):
                    enc_mu    = linear(enc_hidden, self.n_z)
                with tf.variable_scope("sigma"):
                    enc_sigma = tf.nn.softplus(linear(enc_hidden, self.n_z))
            eps = tf.random.normal(tf.shape(enc_mu), prior_mu, prior_sigma, dtype=tf.float32)
            # z = mu + sigma*epsilon
            z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
            with tf.variable_scope("phi_z"):
                z_1 = linear(tf.concat(axis=1,values=(z,x_ped)), self.n_z_1)

            # propagate hidden state eq(7)
            output, state2 = self.lstm(z_1, state) # tf.concat(axis=1, values=(x_1, z_1)

        return (enc_mu, enc_sigma, output, prior_mu, prior_sigma), state2
src/external/vrnn_cell3.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os
from colorama import Fore, Style

class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    self.grid_width, self.grid_height], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, 12,
		                                                        self.pedestrian_vector_dim*self.args.n_other_agents], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.diversity_placeholder = tf.placeholder(dtype=tf.float32,
		                                            shape=[self.batch_size, self.truncated_backprop_length,
		                                                   self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')
		self.seq_length = tf.placeholder(tf.int32, [None])
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped_fw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped_fw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_ped_bw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped_bw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_cell_state_current_lstm_ped_fw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped_fw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_ped_bw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped_bw = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		self.test_natch_ped_grid = np.zeros([args.batch_size, 12, self.pedestrian_vector_dim*self.args.n_other_agents])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped_fw = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped_fw')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped_fw = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped_fw')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped_fw = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped_fw, self.hidden_state_lstm_ped_fw)
		self.cell_state_lstm_ped_bw = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped_bw')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped_bw = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped_bw')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped_bw = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped_bw, self.hidden_state_lstm_ped_bw)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)
		outputs_diversity_series = tf.unstack(self.diversity_placeholder, axis=1)
		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[ped_grid_out_dim*2, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
					                                         trainable=True)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped_fw = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell_fw', trainable=True)
					self.cell_ped_bw = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell_fw',
					                                           trainable=True)
					ped_grid_feature_series , self.current_state_lstm_ped = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.cell_ped_fw,
					                                                                                        cell_bw=self.cell_ped_bw,
					                                                                                        inputs=tf.stack(ped_grid_series,axis=1),
					                                                                                        dtype=tf.float32,
					                                                                                        sequence_length=self.seq_length,
					                                                                                        initial_state_fw=self.init_state_tuple_lstm_ped_fw,
					                                                                                        initial_state_bw=self.init_state_tuple_lstm_ped_bw)

					# Embedding layer of pedestrian grid
					self.cell_outputs_series_lstm_ped = self.process_pedestrian_grid(tf.unstack(tf.stack(ped_grid_feature_series,axis=2),axis=1))

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state,
				                                               self.cell_outputs_series_lstm_grid,
				                                               self.cell_outputs_series_lstm_ped):
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

					# Concatenated LSTM layer
					with tf.variable_scope('generation_model') as scope:

						# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
						cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
						                             args.latent_space_size)
						self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
							cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

						# Output reshape
						enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

						# Decoder layers
						dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
						dec_output_layer = []
						dec_mux_series = []
						dec_muy_series = []
						dec_sigmax_series = []
						dec_sigmay_series = []
						dec_pi_series = []
						# Reusing weights for each truncated back propagation step
						for t_idx in range(len(dec_fc_layer)):
							dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, 50))
							dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100))
						for t_idx in range(len(dec_output_layer)):
							m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1)
							dec_mux_series.append(m_x)
							dec_muy_series.append(m_y)
							dec_sigmax_series.append(s_x)
							dec_sigmay_series.append(s_y)
						dec_pi_series = [dense for dense in dec_pi_series]

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
				                                                decay_rate=0.9, staircase=False)

				self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

				prediction_loss_list = []
				loss_list = []
				kl_loss_list = []
				self.div_loss = []
				self.prediction = []
				self.likelihood = []
				pred = []
				div_loss_over_truncated_back_prop = []

				# Compute Loss Function
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out, out_div in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series, dec_pi_series,
						    outputs_series, outputs_diversity_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(dec_pi, axis=1)

					self.likelihood.append(pi)
					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
						x_diversity = tf.reshape(out_div[:, idx_x], shape=(-1, 1))
						y_diversity = tf.reshape(out_div[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                 sigmax[prediction_step], sigmay[prediction_step],
						                                 x_data, y_data, pi)

						diversity_loss = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                        sigmax[prediction_step], sigmay[prediction_step],
						                                       x_diversity, y_diversity, pi)

						div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss_list.append(kl_loss)
					div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			self.prediction = tf.stack(self.prediction, axis=1)
			self.likelihood = tf.stack(self.likelihood, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
			self.total_loss = tf.reduce_mean(loss_list, axis=0)+args.diversity_update*self.div_loss + tf.reduce_mean(kl_loss_list, axis=0)*self.beta
			self.reconstruction_loss = tf.reduce_mean(loss_list, axis=0)
			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			self.loss_summary = tf.summary.scalar('loss', self.total_loss)
			self.kl_loss_summary = tf.summary.scalar('kl-loss', self.kl_loss)
			self.div_loss_summary = tf.summary.scalar('div_loss', self.div_loss)
			self.reconstruction_loss_summary = tf.summary.scalar('reconstruction_loss', self.reconstruction_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)

			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
			self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
			                                                       graph=tf.Session().graph)

			self.summary = tf.summary.merge_all()

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
			self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		n_other_agents = np.zeros([self.args.batch_size])
		for i,other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.step: kwargs['step'],
		        self.seq_length: n_other_agents,
		        self.diversity_placeholder: kwargs['batch_div'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
		        self.cell_state_lstm_ped_fw: self.cell_state_current_lstm_ped_fw,
		        self.hidden_state_lstm_ped_fw: self.hidden_state_current_lstm_ped_fw,
		        self.cell_state_lstm_ped_bw: self.cell_state_current_lstm_ped_bw,
		        self.hidden_state_lstm_ped_bw: self.hidden_state_current_lstm_ped_bw,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }


	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		n_other_agents = np.zeros([self.args.batch_size])
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		self.test_natch_ped_grid = np.roll(self.test_natch_ped_grid,shift=1,axis=1)
		self.test_natch_ped_grid[:,0] = kwargs["batch_ped_grid"][:, step]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
		        self.input_ped_grid_placeholder: self.test_natch_ped_grid,
		        self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
		        self.step: 0,
		        self.seq_length: n_other_agents,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
		        self.cell_state_lstm_ped_fw: self.test_cell_state_current_lstm_ped_fw,
		        self.hidden_state_lstm_ped_fw: self.test_hidden_state_current_lstm_ped_fw,
		        self.cell_state_lstm_ped_bw: self.test_cell_state_current_lstm_ped_bw,
		        self.hidden_state_lstm_ped_bw: self.test_hidden_state_current_lstm_ped_bw,
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def feed_val_dic(self, **kwargs):
		n_other_agents = np.zeros([self.args.batch_size])
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: kwargs["batch_vel"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.diversity_placeholder: kwargs['batch_div'],
		        self.output_placeholder: kwargs['batch_y'],
		        self.seq_length: n_other_agents,
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), kwargs["state_noise"]),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), kwargs["state_noise"]),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
		        self.cell_state_lstm_ped_fw: np.random.normal(self.test_cell_state_current_lstm_ped_fw.copy(), kwargs["ped_noise"]),
		        self.hidden_state_lstm_ped_fw: np.random.normal(self.test_hidden_state_current_lstm_ped_fw.copy(), kwargs["ped_noise"]),
		        self.cell_state_lstm_ped_bw: np.random.normal(self.test_cell_state_current_lstm_ped_bw.copy(),
		                                                   kwargs["ped_noise"]),
		        self.hidden_state_lstm_ped_bw: np.random.normal(self.test_hidden_state_current_lstm_ped_bw.copy(),
		                                                     kwargs["ped_noise"]),
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped_fw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped_fw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_ped_bw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped_bw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])


	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_cell_state_current_lstm_ped_fw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped_fw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_ped_bw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped_bw[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def update_test_hidden_state(self):
		self.test_cell_state_current = self.cell_state_current.copy()
		self.test_hidden_state_current = self.hidden_state_current.copy()
		self.test_cell_state_current_lstm_grid = self.cell_state_current_lstm_grid.copy()
		self.test_hidden_state_current_lstm_grid = self.hidden_state_current_lstm_grid.copy()
		self.test_cell_state_current_lstm_ped_fw = self.cell_state_current_lstm_ped_fw.copy()
		self.test_hidden_state_current_lstm_ped_fw = self.hidden_state_current_lstm_ped_fw.copy()
		self.test_cell_state_current_lstm_ped_bw = self.cell_state_current_lstm_ped_bw.copy()
		self.test_hidden_state_current_lstm_ped_bw = self.hidden_state_current_lstm_ped_bw.copy()
		self.test_cell_state_current_lstm_concat = self.cell_state_current_lstm_concat.copy()
		self.test_hidden_state_current_lstm_concat = self.hidden_state_current_lstm_concat.copy()

	def validation_step(self, sess, feed_dict_validation,update = True):
		batch_loss, _model_prediction,_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, summary = sess.run([self.reconstruction_loss,
		                                                self.prediction,
		                                                self.current_state,
		                                                self.current_state_lstm_grid,
		                                                self.current_state_lstm_ped,
		                                                self.current_state_lstm_concat,
		                                                self.loss_summary],
		                                                feed_dict=feed_dict_validation)

		if update:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[0]
			self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[1]
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, summary, _model_prediction

	def train_step(self, sess, feed_dict_train, step=0):

		# Diversity Training
		if step > 2000 and self.args.diversity_update:
			dict = {"batch_vel": feed_dict_train[self.input_state_placeholder],
			        "batch_grid": feed_dict_train[self.input_grid_placeholder],
			        "batch_ped_grid": feed_dict_train[self.input_ped_grid_placeholder],
			        "batch_y": feed_dict_train[self.output_placeholder],
			        "batch_div": feed_dict_train[self.diversity_placeholder],
			        "step": step,
			        "state_noise": 0.5,
			        "grid_noise": 0.0,
			        "ped_noise": 0.5
			        }
			self.update_test_hidden_state()
			feed_test_dic = self.feed_val_dic(**dict)
			loss, _, y_model_pred = self.validation_step(sess, feed_test_dic, False)
			batch_y_varible = np.zeros((self.args.batch_size, self.args.truncated_backprop_length,
			                            self.args.prediction_horizon * self.args.output_dim))
			# Grount truth
			batch_y = feed_dict_train[self.output_placeholder]

			diversity_distance = 0
			diff_traj_index = 0
			# Find most different trajectory index
			for batch_id in range(self.args.batch_size):
				for t_id in range(self.args.truncated_backprop_length):
					for mix_id in range(self.args.n_mixtures):
						for pred_id in range(self.args.prediction_horizon):
							new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures+ mix_id]
							new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + mix_id]
							batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
								[new_batch_x, new_batch_y])
						if diversity_distance < np.linalg.norm(batch_y_varible[batch_id, t_id]-batch_y[batch_id, t_id]):
							diversity_distance = np.linalg.norm(batch_y_varible[batch_id, t_id] - batch_y[batch_id, t_id])
							diff_traj_index = mix_id
					# Saving most different traj
					for pred_id in range(self.args.prediction_horizon):
						new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + diff_traj_index]
						new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + diff_traj_index]
						batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
							[new_batch_x, new_batch_y])

			# Assemble feed dict for training
			feed_dict_train[self.diversity_placeholder] = batch_y_varible

		_, batch_loss, kl_loss, recons_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		 _current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
		                                                                                                  self.total_loss,
		                                                                                                  self.kl_loss,
		                                                                                                  self.reconstruction_loss,
		                                                                                                  self.current_state,
		                                                                                                  self.current_state_lstm_grid,
		                                                                                                  self.current_state_lstm_ped,
		                                                                                                  self.current_state_lstm_concat,
		                                                                                                  self.prediction,
		                                                                                                  self.summary,
		                                                                                                  self.learning_rate,
		                                                                                                  self.beta,
		                                                                                                  self.deconv1,
		                                                                                                  self.div_loss,
		                                                                                                  self.autoencoder_loss],
		                                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[0]
		self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[1]
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		# Train autoencoder if loss larger than threshold
		""""""
		if np.mean(autoencoder_loss) > 0.05:
			self.run_autoencoder(sess, feed_dict_train)
			self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

		out = {"batch_loss": recons_loss,
		       "model_predictions": _model_prediction,
		       "summary": _summary_str,
		       "learning_rate": lr,
		       "autoencoder_loss": autoencoder_loss,
		       "log loss": batch_loss,
		       "kl_loss": kl_loss,
		       "output_decoder": output_decoder
		       }

		return out


	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)


	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, _likelihood = sess.run([self.current_state,
		                                              self.current_state_lstm_grid,
		                                              self.current_state_lstm_ped,
		                                              self.current_state_lstm_concat,
		                                              self.prediction,
		                                              self.deconv1,
                                                  self.likelihood],
		                                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[0]
			self.cell_state_current_lstm_ped_bw, self.hidden_state_current_lstm_ped_fw = _current_state_lstm_ped[1]
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, _likelihood


	def warmstart_model(self, args, sess):
		# Restore whole model
		try:
			print('Loading session from "{}"'.format(args.model_path))
			ckpt = tf.train.get_checkpoint_state(args.model_path)
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.full_saver.restore(sess, ckpt.model_checkpoint_path)
		except:
			print(Fore.RED + "Model not found..."+Style.RESET_ALL)
			sys.exit()


	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print(Fore.RED +"Error. Encoder not found...")
			sys.exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True


	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)


	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)],
				                   1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss


	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result


	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y, pi):

		# gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		# gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		loss = tf.multiply(z + tf.log(sxsy), pi)

		return tf.reduce_mean(loss)


	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):
		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #


	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors


	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector


	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations


	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation


	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/BiVRNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from irnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.irnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os


class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')

		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')

		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)


		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			with tf.variable_scope('encoder_model') as scope:

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state,
				                                                     ped_grid_series):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1,
					                             name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
				cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
				                             args.latent_space_size)
				self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
					cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# Output reshape
				enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
				enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
				dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
				prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
				prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

				# Decoder layers
				dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
				dec_fc_layer2 = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_fc_layer])))
				dec_output_layer = []
				dec_mux_series = []
				dec_muy_series = []
				dec_sigmax_series = []
				dec_sigmay_series = []
				dec_pi_series = []
				for t_idx in range(len(dec_fc_layer2)):
					dec_output_layer.append(linear(dec_fc_layer2[t_idx], self.output_vec_dim, t_idx))
				for t_idx in range(len(dec_output_layer)):
					if args.output_pred_state_dim > 4:
						m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
						dec_mux_series.append(m_x)
						dec_muy_series.append(m_y)
						dec_sigmax_series.append(s_x)
						dec_sigmay_series.append(s_y)
						dec_pi_series.append(likelihood)
					else:
						m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
						dec_mux_series.append(m_x)
						dec_muy_series.append(m_y)
						dec_sigmax_series.append(s_x)
						dec_sigmay_series.append(s_y)

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
				                                                decay_rate=0.9, staircase=False)

				self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

			prediction_loss_list = []
			loss_list = []
			kl_loss_list = []
			self.div_loss = []
			self.prediction = []
			pred = []
			div_loss_over_truncated_back_prop = []
			if args.output_pred_state_dim > 4:
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series, dec_pi_series,
						    outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
							pred = tf.concat([pred, pi[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
							pred = tf.concat([pred, pi[prediction_step]], 1)

						lossfunc = self.get_mixture_lossfunc(mux[prediction_step], muy[prediction_step],
						                                     sigmax[prediction_step], sigmay[prediction_step],
						                                     pi[prediction_step],
						                                     x_data, y_data)

						diversity_loss = self.get_diversity_loss(mux[prediction_step], muy[prediction_step],
						                                         sigmax[prediction_step], sigmay[prediction_step])

						div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
					kl_loss_list.append(kl_loss)
					div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			else:
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series,
						    outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

					sigmax = tf.maximum(tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1)), 0.2)
					sigmay = tf.maximum(tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1)), 0.2)

					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                 sigmax[prediction_step], sigmay[prediction_step],
						                                 x_data, y_data)

						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
					kl_loss_list.append(kl_loss)
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			# Reduce mean in all dimensions
			# self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
			# self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
			self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)
			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
			self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

	def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: batch_ped_grid,
		        self.step: step,
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, batch_x, batch_grid, batch_ped_grid, noise=0):
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: batch_ped_grid,
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def run(self, sess, feed_dict_train):
		_, batch_loss, kl_loss, _current_state, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, div_loss = sess.run([self.update,
		                                                                                                  self.total_loss,
		                                                                                                  self.kl_loss,
		                                                                                                  self.current_state,
		                                                                                                  self.current_state_lstm_concat,
		                                                                                                  self.prediction,
		                                                                                                  self.summary,
		                                                                                                  self.learning_rate,
		                                                                                                  self.beta,
		                                                                                                  self.div_loss],
		                                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state

		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, kl_loss, _model_prediction, _summary_str, lr, beta, 0, div_loss, 0

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state,	_current_state_lstm_concat, \
		_model_prediction = sess.run([self.current_state,
		                                              self.current_state_lstm_concat,
		                                              self.prediction],
		                                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, 0

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y):

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		# Normalization constant
		denom = tf.math.log(2 * np.pi * sxsy)
		# Final PDF calculation
		result = denom + z
		gaussian = tf.reduce_sum(result, axis=1)  # do inner summation

		return gaussian

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/CVGDNN.py
eliatrevisan/IA-MPPI-LBM
import numpy as np

class ConstantVelocity():

    def __init__(self, args):
        self.args = args

    def predict(self, input_list):
        # Initialize array to save constant velocity predictions
        all_pred = []        
        for i in range(0, len(input_list)):
            input = input_list[i]
            pred_list = np.zeros([input.shape[0]], dtype=np.ndarray)
            for step in range(input.shape[0]):
                pred = self.cv(input[step,0], input[step,1], input[step,2], input[step,3]) 
                pred_list[step] = pred
            all_pred.append(pred_list)
        return np.array(all_pred)

    def cv(self, posx, posy, vx, vy):
        pred = []

        for i in range(self.args.prediction_horizon):
            posx = posx + vx*self.args.dt
            posy = posy + vy*self.args.dt
            pred.append([posx, posy])
        return np.array(pred)




src/models/CV_model.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
	import src.data_utils.Support as sup
import numpy as np
import os


def length(sequence):
	used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
	length = tf.reduce_sum(used, 1)
	length = tf.cast(length, tf.int32)
	return length


class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_ped_size = args.rnn_state_ped_size
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.n_peds = self.args.n_other_agents
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    self.grid_width, self.grid_height], name='input_grid')
		if self.args.relative_info:
			self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
			                                                 shape=[self.batch_size, self.truncated_backprop_length,
			                                                        self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
		if "prev" in self.args.others_info:
			self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
			                                                 shape=[self.batch_size, self.truncated_backprop_length,
			                                                        self.n_peds,
			                                                        self.pedestrian_vector_dim * (self.prev_horizon + 1)],
			                                                 name='ped_grid')
		else:
			self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
			                                                 shape=[self.batch_size, self.truncated_backprop_length,
			                                                        self.n_peds,
			                                                        self.pedestrian_vector_dim * self.prediction_horizon],
			                                                 name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32, shape=[], name='global_step')
		self.seq_len = tf.placeholder(tf.int32, [None])

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.ped_gru_states_current = np.zeros([args.batch_size, self.n_peds, self.rnn_state_ped_size])
		self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_ped_gru_states_current = np.zeros([args.batch_size, self.n_peds, self.rnn_state_ped_size])
		self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.ped_gru_states = tf.placeholder(dtype=tf.float32,
		                                     shape=[self.batch_size, self.n_peds, self.rnn_state_ped_size],
		                                     name='cell_state_lstm_ped')
		self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                    name='cell_state_lstm_ped')
		self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                    name='cell_state_lstm_ped')
		# self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		# self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		# self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
					                                         trainable=train_encoder)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

				# LSTM cells to process pedestrian
				self.cell_outputs_series_lstm_ped = []
				self.current_state_lstm_ped = []
				self.other_peds_states = []
				self.other_peds_gru_outputs = []
				with tf.variable_scope('lstm_ped', reuse=tf.AUTO_REUSE) as scope:
					# TODO: try to share weights
					for i, ped in enumerate(ped_grid_series):
						# TODO: ADAPT for ped input sequence
						ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell' + str(i),
						                                  trainable=True)
						ped_input = tf.unstack(ped, axis=1)
						ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell, ped_input, dtype=tf.float32,
						                                                          initial_state=self.ped_gru_states[:, i, :])
						self.other_peds_states.append(ped_gru_state)

						if i > 0:
							for j, ped_out in enumerate(ped_gru_output):
								out_ped_concat[j] = tf.concat([out_ped_concat[j], ped_gru_output[j]], axis=1)
						else:
							out_ped_concat = ped_gru_output
					self.other_peds_gru_outputs = out_ped_concat

					self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
					                                          trainable=True)
					self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
					                                          trainable=True)

					outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
					                                                  cell_bw=self.fw_cell_ped,
					                                                  inputs=tf.stack(self.other_peds_gru_outputs, axis=1),
					                                                  dtype=tf.float32,
					                                                  initial_state_fw=self.fw_ped_states,
					                                                  initial_state_bw=self.bw_ped_states)

					self.cell_outputs_series_lstm_ped = tf.unstack(
						tf.transpose(tf.concat([outputs[0], outputs[1]], axis=2), [1, 0, 2]), axis=0)
					self.current_state_lstm_ped = states
					self.other_peds_states = tf.stack(self.other_peds_states, axis=1)

					rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
					                                                   shape=[self.rnn_state_size_lstm_ped * 2,
					                                                          self.rnn_state_size_lstm_ped],
					                                                   trainable=self.is_training)
					b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
					                                 trainable=self.is_training)

					self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
					              self.cell_outputs_series_lstm_ped]

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.fused):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out',
				                                 shape=[self.fc_hidden_unit_size, self.output_dim * self.prediction_horizon],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
				                               trainable=self.is_training)

				# Second concatenated LSTM layer
				with tf.variable_scope('lstm_concat') as scope:
					cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# FC layers
				fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
				           self.cell_outputs_series_lstm_concat]
				# fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
				#           fc_fuse]
				# fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
				# Model prediction
				self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

			# Exponential / stepwise decrease of learning rate
			global_step = tf.Variable(0, trainable=False)
			self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
			                                                decay_rate=0.9, staircase=False)

			# Loss function
			loss_list = []
			lossfunc = []
			# Iterate through tbp steps
			for pred, out in zip(self.prediction,
			                     outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
				prediction_loss_list = []
				val_prediction_loss_list = []
				for prediction_step in range(self.prediction_horizon):
					# Compute Euclidean error for each prediction step within each tbp step

					idx_x = prediction_step * self.output_dim
					idx_y = idx_x + 1

					# Calculate loss for the current ped
					lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
					                          tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

					prediction_loss_list.append(lossfunc)  # sum over number of batches

				loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

			self.prediction = tf.stack(self.prediction, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0)  # + l2 # mean over tbp steps

			# Optimizer specification
			self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			# self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

		# Only save / restore the parameters responsible for occupancy grid processing
		self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
		self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)

	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		return {self.input_state_placeholder: batch_x,
		        self.input_grid_placeholder: kwargs['batch_grid'],
		        self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
		        self.step: kwargs['step'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
		        self.ped_gru_states: self.ped_gru_states_current,
		        self.fw_ped_states: self.fw_ped_states_current,
		        self.bw_ped_states: self.bw_ped_states_current,
		        self.seq_len: [100],
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
		        self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
		        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
		        self.ped_gru_states: self.test_ped_gru_states_current,
		        self.fw_ped_states: self.test_fw_ped_states_current,
		        self.bw_ped_states: self.test_bw_ped_states_current,
		        self.seq_len: [100],
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.ped_gru_states_current = np.zeros([self.args.batch_size, self.n_peds, self.args.rnn_state_ped_size])
					self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
					self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_ped_gru_states_current = np.zeros([self.args.batch_size, self.n_peds, self.args.rnn_state_ped_size])
					self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
					self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def run(self, sess, feed_dict_train):
		_, batch_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, _other_peds_states, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
		                                                                                  self.total_loss,
		                                                                                  self.current_state,
		                                                                                  self.current_state_lstm_grid,
		                                                                                  self.current_state_lstm_ped,
		                                                                                  self.other_peds_states,
		                                                                                  self.current_state_lstm_concat,
		                                                                                  self.prediction,
		                                                                                  self.summary,
		                                                                                  self.learning_rate,
		                                                                                  self.deconv1,
		                                                                                  self.autoencoder_loss],
		                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
		self.ped_gru_states_current = _other_peds_states
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, 0, _model_prediction, _summary_str, lr, 0, output_decoder, 0, 0, 0, 0

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, _other_peds_states, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, \
		_cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
		_cell_outputs_series_lstm_concat = sess.run([self.current_state,
		                                             self.current_state_lstm_grid,
		                                             self.current_state_lstm_ped,
		                                             self.other_peds_states,
		                                             self.current_state_lstm_concat,
		                                             self.prediction,
		                                             self.deconv1,
		                                             self.cell_outputs_series_state,
		                                             self.cell_outputs_series_lstm_ped,
		                                             self.cell_outputs_series_lstm_concat
		                                             ],
		                                            feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_fw_ped_states_current, self.test_bw_ped_states_current = _current_state_lstm_ped
			self.test_ped_gru_states_current = _other_peds_states
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, output_decoder, _cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		try:
			if self.args.regulate_log_loss:
				likelihood_loss = -1.0 * tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)), 4.0)
			else:
				likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		except:
			likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/IntNet.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.prediction_horizon,
                                                              self.pedestrian_vector_dim],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []

        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_series = tf.unstack(ped, axis=1)
            ped_states = []
            # iterate over truncated back propagation
            for j, ped_t in enumerate(ped_series):
              ped_input = tf.unstack(ped_t, axis=1)
              ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
              ped_states.append(ped_gru_state)
            self.other_peds_states.append(ped_states)
          a = tf.transpose(tf.concat(tf.unstack(tf.stack(self.other_peds_states),axis=0),axis=2),[1,0,2])
          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=a,
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.squeeze(tf.stack(self.other_peds_states,axis=2),axis=0)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder, 0, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntNet2.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_size_lstm_ped])
    self.ped_dec_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_size_lstm_ped])
    self.test_ped_dec_gru_states_current = np.zeros([args.batch_size, self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.ped_dec_gru_states = tf.placeholder(dtype=tf.float32,
                                         shape=[self.batch_size, self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon],
                                         name='cell_dec_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='fw_cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='bw_cell_state_lstm_ped')
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)

    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:
      with tf.variable_scope('encoder_model') as scope:
        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_dec_states = []
        self.other_peds_gru_outputs = []
        out_ped_concat = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          w_enc_state_lstm_ped = self.get_weight_variable(name='w_enc_state_lstm_ped',
                                                          shape=[self.rnn_state_size_lstm_ped,
                                                                 self.rnn_state_ped_size],
                                                          trainable=self.is_training)
          b_enc_state_lstm_ped = self.get_bias_variable(name='b_enc_state_lstm_ped', shape=[1, self.rnn_state_ped_size],
                                                        trainable=self.is_training)

          b_dec_state_lstm_ped = self.get_bias_variable(name='b_dec_state_lstm_ped',
                                                        shape=[1, self.rnn_state_size_lstm_ped],
                                                        trainable=self.is_training)
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_enc_cell = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='ped_enc_lstm_cell'+str(i), trainable=True)
            ped_dec_cell = tf.contrib.rnn.GRUCell(self.pedestrian_vector_dim * self.prediction_horizon, name='ped_dec_lstm_cell' + str(i),trainable=True)

            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_enc_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            ped_gru_output_fc = tf.add(tf.matmul(ped_gru_output[0], w_enc_state_lstm_ped), b_enc_state_lstm_ped)

            ped_dec_fc = tf.add(tf.matmul(ped_gru_output_fc, tf.transpose(w_enc_state_lstm_ped)), b_dec_state_lstm_ped)

            ped_dec_gru_output, ped_dec_gru_state = tf.contrib.rnn.static_rnn(ped_dec_cell,[ped_dec_fc],dtype=tf.float32,initial_state=self.ped_dec_gru_states[:,i,:])

            self.other_peds_dec_states.append(ped_dec_gru_state)

            ped_loss = tf.reduce_mean(tf.sqrt(tf.square(tf.subtract(ped_dec_gru_output[0], ped_input[0]))),axis=1)

            out_ped_concat.append(ped_gru_output_fc)

          self.ped_enc_dec_loss = tf.reduce_mean(ped_loss)

          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)
          self.other_peds_dec_states = tf.stack(self.other_peds_dec_states, axis=1)

          w_birnn = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, w_birnn), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.fused):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        #fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
        #           fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction, axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      #self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)
      self.ped_dec_enc_gradients = tf.gradients(self.ped_enc_dec_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
      self.ped_dec_enc_gradients_clipped, _ = tf.clip_by_global_norm(self.ped_dec_enc_gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      self.update_ped_dec_enc = self.optimizer.apply_gradients(zip(self.ped_dec_enc_gradients_clipped, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('Model Loss', self.total_loss)
      tf.summary.scalar('Pedestrian Encode Decoder Loss', self.ped_enc_dec_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      batch_x = kwargs['batch_vel']
      batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.ped_dec_gru_states: self.ped_dec_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.ped_dec_gru_states: self.test_ped_dec_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_size_lstm_ped])
          self.ped_dec_gru_states_current = np.zeros([self.args.batch_size, self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_size_lstm_ped])
          self.test_ped_dec_gru_states_current = np.zeros([self.args.batch_size, self.n_peds, self.pedestrian_vector_dim * self.prediction_horizon])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_ped,_other_peds_states, \
    _other_peds_dec_states, _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, _ped_dec_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.other_peds_dec_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.ped_enc_dec_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.ped_dec_gru_states_current = _other_peds_dec_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, 0,0, 0, 0,_ped_dec_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)
  def run_ped_encoder(self,sess,feed_dict_train):
    _ = sess.run([self.update_ped_dec_enc],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_ped,_other_peds_states, \
    _other_peds_dec_states ,_current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.other_peds_dec_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_fw_ped_states_current, self.test_bw_ped_states_current = _current_state_lstm_ped
      self.test_ped_gru_states_current = _other_peds_states
      self.test_ped_dec_gru_states_current = _other_peds_dec_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntNet3.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')

    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)
        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        #fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
        #           fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction, axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      #self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('Model Loss', self.total_loss)
      tf.summary.scalar('Grid Encode Loss', self.autoencoder_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      tf.summary.image('Pedestrian Grid', self.deconv1)

      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      batch_x = kwargs['batch_vel']
      batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :, :], axis=0),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state,_cell_outputs_series_lstm_grid, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_grid,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_grid, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntNetGrid.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.pos_output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    pos_outputs_series = tf.unstack(self.pos_output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      pos_loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out, pos_out in zip(self.prediction,outputs_series,pos_outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        pos = 0.0
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          pos += pred[:,idx_x:idx_x+self.output_dim] * self.args.dt

          prediction_loss_list.append(lossfunc)  # sum over number of batches
        pos_loss = tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(pos,pos_out)),axis=1))
        pos_loss_list.append(tf.reduce_mean(pos_loss))  # mean over prediction horizon
        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0) +tf.reduce_mean(pos_loss_list, axis=0)*0.5 #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y, batch_pos_final_backup):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.pos_output_placeholder: batch_pos_final_backup,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntNet_end_pos_loss.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.math.abs(tf.reduce_mean(loss_list, axis=0)-0.05)+0.05
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      batch_x = kwargs['batch_x'][:,:,2:]
      batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntNet_flood.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*self.prediction_horizon], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    elif "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction,axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    batch_x = kwargs['batch_y_pred']
    batch_x[:,:,:2] = kwargs['batch_vel'][:,:,:2]
    batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntNet_future.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    if "prev" in self.args.others_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim *(self.prev_horizon+1)],
                                                       name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * self.prediction_horizon],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.fused):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        #fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
        #           fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          error_x = tf.maximum(tf.abs(tf.subtract(pred[:, idx_x], out[:, idx_x])),0.03)
          error_y = tf.maximum(tf.abs(tf.subtract(pred[:, idx_y], out[:, idx_y])),0.03)
          lossfunc = tf.sqrt(tf.add(tf.square(error_x),  # error in x
                                    tf.square(error_y)))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction, axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      #self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      batch_x = kwargs['batch_vel']
      batch_y = kwargs['batch_y']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,step,noise=0):
    return {self.input_state_placeholder: np.expand_dims(batch_x[:, step, :], axis=0),
            self.input_grid_placeholder: np.expand_dims(batch_grid[:, step, :, :], axis=0),
            self.input_ped_grid_placeholder: np.expand_dims(batch_ped_grid[:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,0, 0, autoencoder_loss,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder,_cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntNet_minloss.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    if self.args.relative_info:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')
    else:
      self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                       shape=[self.batch_size, self.truncated_backprop_length,
                                                              self.n_peds,
                                                              self.pedestrian_vector_dim * (self.prev_horizon+1)],
                                                       name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    #self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder, 0, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntNetprev.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell3 import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell3 import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    #self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
                                                     shape=[self.batch_size, self.truncated_backprop_length,
                                                            self.args.n_other_agents,
                                                            self.pedestrian_vector_dim * self.prediction_horizon],
                                                     name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    #self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    #self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    #self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    #self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    #self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    #self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    #self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    #self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    #conv_grid_feature_series = []
    #for output in self.process_grid(self.occ_grid_series):
    #  conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        #with tf.variable_scope('lstm_grid') as scope:
        #  self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
        #                                           trainable=True)
        #  self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
        #                                                                                               conv_grid_feature_series,
        #                                                                                               dtype=tf.float32,
        #                                                                                              initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped', reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.args.rnn_state_ped_size, name='ped_basic_lstm_cell' + str(i),
                                              trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell, ped_input, dtype=tf.float32,
                                                                      initial_state=self.ped_gru_states[:, i, :])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j], ped_gru_output[j]], axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                            cell_bw=self.fw_cell_ped,
                                                            inputs=tf.stack(self.other_peds_gru_outputs, axis=1),
                                                            dtype=tf.float32,
                                                            initial_state_fw=self.fw_ped_states,
                                                            initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(
            tf.transpose(tf.concat([outputs[0], outputs[1]], axis=2), [1, 0, 2]), axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states, axis=1)
        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args,args.rnn_state_size_lstm_concat,self.output_vec_dim, args.prior_size, args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step],muy[prediction_step],sigmax[prediction_step],sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    #self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    #self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            #self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            #self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,noise=0):
    return {self.input_state_placeholder: batch_x,
            #self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            #self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            #self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          #self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          #self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          #self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          #self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.args.n_other_agents, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          #self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          #self.deconv1,
                                                          #self.autoencoder_loss
                                                          ],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    #self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, 0,0, 0, 0,0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction,\
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped,\
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              #self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              #self.deconv1
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.dec_series],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      #self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
      self.ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, 0, _cell_outputs_series_state, _cell_outputs_series_lstm_ped, _cell_outputs_series_lstm_concat

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/IntVGDNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
  from utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  from src.data_utils.utils import *
import numpy as np
import os
from datetime import datetime
from colorama import Fore, Style

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)

    self.input_grid = tf.expand_dims(input=self.occ_grid_series[0], axis=3)

    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      self.prediction = tf.stack(self.prediction, axis=1)

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)

      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Tensorboard data
      self.loss_summary = tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)

      # Grid input
      sliced_grid = tf.slice(self.input_grid, begin=[0,0,0,0], size=[1, self.grid_width, self.grid_height, 1], name="slice_in_grid")
      tf.summary.image('input_grid', sliced_grid, max_outputs=1)

      # Grid output
      sliced_grid = tf.slice(self.deconv1, begin=[0,0,0,0], size=[1, self.grid_width, self.grid_height, 1], name="slice_out_grid")
      tf.summary.image('output_grid', sliced_grid, max_outputs=1)

      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
      self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'), graph=tf.Session().graph)

      self.summary = tf.summary.merge_all()

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    return {self.input_state_placeholder: kwargs['batch_vel'],
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: kwargs['batch_y'],
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_val_dic(self, **kwargs):
    return {self.input_state_placeholder: kwargs['batch_vel'],
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.output_placeholder: kwargs['batch_y'],
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, **kwargs):
    step = kwargs["step"]
    return {self.input_state_placeholder:  np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
            self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :,:], axis=1),
            self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),kwargs["state_noise"]),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),kwargs["state_noise"]),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),kwargs["grid_noise"]),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),kwargs["grid_noise"]),
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def validation_step(self,sess,feed_dict_validation):
    batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, summary = sess.run([self.total_loss,
                                           self.current_state,
                                           self.current_state_lstm_grid,
                                           self.current_state_lstm_ped,
                                           self.current_state_lstm_concat,
                                           self.loss_summary],
                                           feed_dict=feed_dict_validation)
    self.test_cell_state_current, self.test_hidden_state_current = _current_state
    self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, summary

  def train_step(self,sess,feed_dict_train,step=0):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    # Train autoencoder if loss larger than threshold
    if np.mean(autoencoder_loss) > 0.05:
      print(Fore.LIGHTMAGENTA_EX +"\n[{:s}] Training autoencoder. Model loss: {:01.2f}" .format(
          datetime.now().strftime("%d-%m-%y %H:%M:%S"), autoencoder_loss))

      self.run_autoencoder(sess, feed_dict_train)
      self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

    out = {"batch_loss": batch_loss,
            "model_predictions": _model_prediction,
            "summary": _summary_str,
            "learning_rate": lr,
            "autoencoder_loss": autoencoder_loss,
            "output_decoder": output_decoder
            }

    return out

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/RNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.goal_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, 2], name='goal')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    self.goal_series = tf.unstack(self.goal_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=True)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat+2, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(tf.concat([lstm_out,goal],axis=1), W_fc), b_fc)) for lstm_out,goal in
                   zip(self.cell_outputs_series_lstm_concat,self.goal_series)]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y,batch_goal):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.goal_placeholder: batch_goal,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_goal,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.goal_placeholder: batch_goal,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder, 0, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/RNN_goal.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))

    # Process Pedestrian occupancy grid, apply convolutional filters
    self.ped_occ_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
    ped_grid_feature_series = []
    with tf.variable_scope('ped_encoder_model') as scope:
      for output in self.process_ped_grid(self.ped_occ_grid_series):
        ped_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:

        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
                   self.cell_outputs_series_lstm_concat]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)

    self.ped_convnet_saver = tf.train.Saver(self.ped_encoder_var_list)
    self.ped_encoder_update = self.optimizer.minimize(self.ped_encoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder,output_ped_decoder,\
    autoencoder_loss,ped_encoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.ped_deconv1,
                                                          self.autoencoder_loss,
                                                          self.ped_encoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, output_decoder,output_ped_decoder, 0, autoencoder_loss,ped_encoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def run_ped_encoder(self,sess,feed_dict_train):
    _ = sess.run([self.ped_encoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)
    print('Loading Pedestrian convnet parameters from "{}"'.format(args.model_path+"/ped_encoder"))
    ckpt_conv = tf.train.get_checkpoint_state(args.model_path+"/ped_encoder")
    if ckpt_conv == None:
      print("Error. Pedestrian Convnet not found. Loading default")
      ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    print('Restoring Pedestrian convnet {}'.format(ckpt_conv.model_checkpoint_path))
    self.ped_convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope("auto_encoder") as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=True)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="auto_encoder")

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="auto_encoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors

  def process_ped_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope("auto_encoder") as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=True)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases", trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                 shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                       [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                       padding='VALID', trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                       [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                       padding='SAME', trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.ped_deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME', trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.ped_deconv1.get_shape()))

          self.ped_encoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="auto_encoder")

          self.ped_encoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.ped_deconv1)), name="auto_encoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/RNN_grid.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight

    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.goal_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, 2], name='goal')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])

    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    self.goal_series = tf.unstack(self.goal_placeholder, axis=1)
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat+2, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim*self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim*self.prediction_horizon], trainable=self.is_training)

        # Second concatenated LSTM layer
        with tf.variable_scope('lstm_concat') as scope:
          cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
          self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
            cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # FC layers
        fc_fuse = [tf.nn.relu(tf.add(tf.matmul(tf.concat([lstm_out,goal],axis=1), W_fc), b_fc)) for lstm_out,goal in
                   zip(self.cell_outputs_series_lstm_concat,self.goal_series)]
        fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.args.keep_prob) for dense in
                   fc_fuse]
        # fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
        # Model prediction
        self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

      # Exponential / stepwise decrease of learning rate
      global_step = tf.Variable(0, trainable=False)
      self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                      decay_rate=0.9, staircase=False)

      # Loss function
      loss_list = []
      lossfunc = []
      # Iterate through tbp steps
      for pred, out in zip(self.prediction,
                           outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
        prediction_loss_list = []
        val_prediction_loss_list = []
        for prediction_step in range(self.prediction_horizon):
          # Compute Euclidean error for each prediction step within each tbp step

          idx_x = prediction_step * self.output_dim
          idx_y = idx_x + 1

          # Calculate loss for the current ped
          lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
                                    tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

          prediction_loss_list.append(lossfunc)  # sum over number of batches

        loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

      # Get all trainable variables
      tvars = tf.trainable_variables()

      # L2 loss
      l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

      # Reduce mean in all dimensions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)  #+ l2 # mean over tbp steps
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y,batch_goal):
    return {self.input_state_placeholder: batch_x,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.goal_placeholder: batch_goal,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_goal,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.goal_placeholder: batch_goal,
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss,  _current_state, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr = sess.run([self.update,
                                                          self.total_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0,_model_prediction, _summary_str, lr, 0, 0, 0, 0

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction = sess.run([self.current_state,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, 0

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/RNN_no_grid_goal.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os


class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight

		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.goal_placeholder = tf.placeholder(dtype=tf.float32,
		                                       shape=[self.batch_size, self.truncated_backprop_length, args.goal_size],
		                                       name='goal')

		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32, shape=[], name='global_step')

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)
		self.goal_series = tf.unstack(self.goal_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
				# Embedding layer of pedestrian grid
				ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped,
					                                                                                           ped_grid_feature_series,
					                                                                                           dtype=tf.float32,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat+args.goal_size, self.rnn_state_size_lstm_concat+args.goal_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.rnn_state_size_lstm_concat+args.goal_size],
				                              trainable=self.is_training)
				W_fc2 = self.get_weight_variable(name='w_fc2',
				                                shape=[self.rnn_state_size_lstm_concat+args.goal_size, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc2 = self.get_bias_variable(name='b_fc2', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out',
				                                 shape=[self.fc_hidden_unit_size, self.output_dim * self.prediction_horizon],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
				                               trainable=self.is_training)

				# Second concatenated LSTM layer
				with tf.variable_scope('lstm_concat') as scope:
					cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# Concatenate the goal information
				concat_goal = []
				for lstm_out, goal_feature in zip(self.cell_outputs_series_lstm_concat, self.goal_series):
					concat_goal.append(tf.concat([lstm_out, goal_feature], axis=1, name="concatenate_lstm_out_with_goal_info"))

				# FC layers
				fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
				           concat_goal]
				#fc_drop = [tf.layers.dropout(inputs=dense, rate=1 - self.args.keep_prob) for dense in
				 #          fc_fuse]
				fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
				# Model prediction
				self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse2]

			# Exponential / stepwise decrease of learning rate
			global_step = tf.Variable(0, trainable=False)
			self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
			                                                decay_rate=0.9, staircase=False)

			# Loss function
			loss_list = []
			lossfunc = []
			# Iterate through tbp steps
			for pred, out, goal in zip(self.prediction,
			                     outputs_series,self.goal_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
				prediction_loss_list = []
				val_prediction_loss_list = []
				position = 0
				for prediction_step in range(self.prediction_horizon):
					# Compute Euclidean error for each prediction step within each tbp step

					idx_x = prediction_step * self.output_dim
					idx_y = idx_x + 1

					# Calculate loss for the current ped
					lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
					                          tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

					prediction_loss_list.append(lossfunc)  # sum over number of batches

					position += pred[:, idx_x:idx_x+2]*args.dt

				pos_lossfunc = tf.norm(tf.norm(position)-goal[:,0])+tf.norm(tf.math.atan2(position[:, 1],position[:, 0])-goal[:,1])

				loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

			self.prediction = tf.stack(self.prediction, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0)  # + l2 # mean over tbp steps

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.goal_placeholder: kwargs["batch_goal"],
		        self.step: kwargs['step'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
	        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
	        self.goal_placeholder: np.expand_dims(kwargs["batch_goal"][:, step, :], axis=1),
	        self.step: 0,
	        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
	        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
	        # self.cell_state: self.cell_state_current,
	        # self.hidden_state: self.hidden_state_current,
	        # self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
	        # self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
	        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
	        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
	        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
	        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
	        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def run(self, sess, feed_dict_train):
		_, batch_loss, _current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr = sess.run([self.update,
		                                                self.total_loss,
		                                                self.current_state,
		                                                self.current_state_lstm_ped,
		                                                self.current_state_lstm_concat,
		                                                self.prediction,
		                                                self.summary,
		                                                self.learning_rate],
		                                               feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, 0, _model_prediction, _summary_str, lr, 0, 0, 0, 0, 0, 0

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update], feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction = sess.run([self.current_state,
		                              self.current_state_lstm_ped,
		                              self.current_state_lstm_concat,
		                              self.prediction],
		                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, 0, _current_state, _current_state_lstm_ped, _current_state_lstm_concat

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		try:
			if self.args.regulate_log_loss:
				likelihood_loss = -1.0 * tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)), 4.0)
			else:
				likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		except:
			likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/RNN_plannet.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os


class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight

		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32, shape=[], name='global_step')

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
				# Embedding layer of pedestrian grid
				ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped,
					                                                                                           ped_grid_feature_series,
					                                                                                           dtype=tf.float32,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out',
				                                 shape=[self.fc_hidden_unit_size, self.output_dim * self.prediction_horizon],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
				                               trainable=self.is_training)

				# Second concatenated LSTM layer
				with tf.variable_scope('lstm_concat') as scope:
					cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# FC layers
				fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
				           self.cell_outputs_series_lstm_concat]
				fc_drop = [tf.layers.dropout(inputs=dense, rate=1 - self.args.keep_prob) for dense in
				           fc_fuse]
				# fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
				# Model prediction
				self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_drop]

			# Exponential / stepwise decrease of learning rate
			global_step = tf.Variable(0, trainable=False)
			self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
			                                                decay_rate=0.9, staircase=False)

			# Loss function
			loss_list = []
			lossfunc = []
			# Iterate through tbp steps
			for pred, out in zip(self.prediction,
			                     outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
				prediction_loss_list = []
				val_prediction_loss_list = []
				for prediction_step in range(self.prediction_horizon):
					# Compute Euclidean error for each prediction step within each tbp step

					idx_x = prediction_step * self.output_dim
					idx_y = idx_x + 1

					# Calculate loss for the current ped
					lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
					                          tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

					prediction_loss_list.append(lossfunc)  # sum over number of batches

				loss_list.append(tf.reduce_mean(prediction_loss_list))  # mean over prediction horizon

			self.prediction = tf.stack(self.prediction, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0)  # + l2 # mean over tbp steps

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.step: kwargs['step'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
	        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
	        self.step: 0,
	        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
	        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
	        # self.cell_state: self.cell_state_current,
	        # self.hidden_state: self.hidden_state_current,
	        # self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
	        # self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
	        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
	        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
	        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
	        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
	        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def run(self, sess, feed_dict_train):
		_, batch_loss, _current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr = sess.run([self.update,
		                                                self.total_loss,
		                                                self.current_state,
		                                                self.current_state_lstm_ped,
		                                                self.current_state_lstm_concat,
		                                                self.prediction,
		                                                self.summary,
		                                                self.learning_rate],
		                                               feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, 0, _model_prediction, _summary_str, lr, 0, 0, 0, 0, 0, 0

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update], feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction = sess.run([self.current_state,
		                              self.current_state_lstm_ped,
		                              self.current_state_lstm_concat,
		                              self.prediction],
		                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, 0, _current_state, _current_state_lstm_ped, _current_state_lstm_concat

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		try:
			if self.args.regulate_log_loss:
				likelihood_loss = -1.0 * tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)), 4.0)
			else:
				likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		except:
			likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/RNN_simple.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.storn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os
from colorama import Fore, Style

class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    self.grid_width, self.grid_height], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
					                                         trainable=True)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
				# Embedding layer of pedestrian grid
				ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped,
					                                                                                           ped_grid_feature_series,
					                                                                                           dtype=tf.float32,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state,
				                                               self.cell_outputs_series_lstm_grid,
				                                               self.cell_outputs_series_lstm_ped):
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

					# Concatenated LSTM layer
					with tf.variable_scope('generation_model') as scope:

						# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
						cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.rnn_state_size_lstm_concat, args.prior_size,
						                             args.latent_space_size)
						self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
							cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

						# Output reshape
						enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

						# Decoder layers
						dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
						dec_output_layer = []
						dec_mux_series = []
						dec_muy_series = []
						dec_sigmax_series = []
						dec_sigmay_series = []
						dec_pi_series = []
						# Reusing weights for each truncated back propagation step
						for t_idx in range(len(dec_fc_layer)):
							dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, 50))
							dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100))
						for t_idx in range(len(dec_output_layer)):
							m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1)
							dec_mux_series.append(m_x)
							dec_muy_series.append(m_y)
							dec_sigmax_series.append(s_x)
							dec_sigmay_series.append(s_y)
						dec_pi_series = [dense for dense in dec_pi_series]

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
				                                                decay_rate=0.9, staircase=False)

				self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

				prediction_loss_list = []
				loss_list = []
				kl_loss_list = []
				self.div_loss = []
				self.prediction = []
				self.likelihood = []
				pred = []
				div_loss_over_truncated_back_prop = []

				# Compute Loss Function
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series, dec_pi_series,
						    outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(dec_pi, axis=1)

					self.likelihood.append(pi)
					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                 sigmax[prediction_step], sigmay[prediction_step],
						                                 x_data, y_data, pi)

						# diversity_loss = self.get_diversity_loss(mux[prediction_step], muy[prediction_step],
						#                                         sigmax[prediction_step], sigmay[prediction_step])

						# div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss_list.append(kl_loss)
					# div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			self.prediction = tf.stack(self.prediction, axis=1)
			self.likelihood = tf.stack(self.likelihood, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0)# + tf.reduce_mean(kl_loss_list, axis=0)*self.beta

			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			self.loss_summary = tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)

			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
			self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
			                                                       graph=tf.Session().graph)

			self.summary = tf.summary.merge_all()

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
			self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.step: kwargs['step'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }


	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
		        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
		        self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
		        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def feed_val_dic(self, **kwargs):
		return {self.input_state_placeholder: kwargs["batch_vel"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
		        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])


	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def validation_step(self, sess, feed_dict_validation):
		batch_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, summary = sess.run([self.total_loss,
		                                                                    self.current_state,
		                                                                    self.current_state_lstm_grid,
		                                                                    self.current_state_lstm_ped,
		                                                                    self.current_state_lstm_concat,
		                                                                    self.loss_summary],
		                                                                   feed_dict=feed_dict_validation)
		self.test_cell_state_current, self.test_hidden_state_current = _current_state
		self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, summary

	def train_step(self, sess, feed_dict_train):
		_, batch_loss, kl_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
		                                                                                                  self.total_loss,
		                                                                                                  self.kl_loss,
		                                                                                                  self.current_state,
		                                                                                                  self.current_state_lstm_grid,
		                                                                                                  self.current_state_lstm_ped,
		                                                                                                  self.current_state_lstm_concat,
		                                                                                                  self.prediction,
		                                                                                                  self.summary,
		                                                                                                  self.learning_rate,
		                                                                                                  self.beta,
		                                                                                                  self.deconv1,
		                                                                                                  self.div_loss,
		                                                                                                  self.autoencoder_loss],
		                                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		# Train autoencoder if loss larger than threshold
		""""""
		if np.mean(autoencoder_loss) > 0.05:
			self.run_autoencoder(sess, feed_dict_train)
			self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

		out = {"batch_loss": batch_loss,
		       "model_predictions": _model_prediction,
		       "summary": _summary_str,
		       "learning_rate": lr,
		       "autoencoder_loss": autoencoder_loss,
		       "log loss": batch_loss,
		       "output_decoder": output_decoder
		       }

		return out


	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)


	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, _likelihood = sess.run([self.current_state,
		                                              self.current_state_lstm_grid,
		                                              self.current_state_lstm_ped,
		                                              self.current_state_lstm_concat,
		                                              self.prediction,
		                                              self.deconv1,
                                                  self.likelihood],
		                                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, _likelihood


	def warmstart_model(self, args, sess):
		# Restore whole model
		try:
			print('Loading session from "{}"'.format(args.model_path))
			ckpt = tf.train.get_checkpoint_state(args.model_path)
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.full_saver.restore(sess, ckpt.model_checkpoint_path)
		except:
			print(Fore.RED + "Model not found..."+Style.RESET_ALL)
			sys.exit()


	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print(Fore.RED +"Error. Encoder not found...")
			sys.exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True


	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)


	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)],
				                   1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss


	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result


	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y, pi):

		# gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		# gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		loss = tf.multiply(z + tf.log(sxsy), pi)

		return tf.reduce_mean(loss)


	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):
		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #


	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors


	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector


	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations


	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation


	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/STORN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
from colorama import Fore, Style
import datetime
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell2 import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell2 import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')

    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.fused):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out',
                                         shape=[self.fc_hidden_unit_size,
                                                self.output_dim * self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
                                       trainable=self.is_training)

        # cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
                                     args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
          cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series = []
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
          dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100 + t_idx))
        for t_idx in range(len(dec_output_layer)):
          m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1)
          dec_mux_series.append(m_x)
          dec_muy_series.append(m_y)
          dec_sigmax_series.append(s_x)
          dec_sigmay_series.append(s_y)
        dec_pi_series = [dense for dense in dec_pi_series]

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

        loss_list = []
        mse_loss_list = []
        kl_loss_list = []
        self.div_loss = []
        self.prediction = []
        pred = []
        div_loss_over_truncated_back_prop = []

        # Compute Loss Function
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series, dec_pi_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply probabilistic transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(dec_pi, axis=1)

          pred = []
          prediction_loss_list = []
          mse_prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            # Log-likelihood
            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                             sigmax[prediction_step], sigmay[prediction_step],
                                             x_data, y_data, pi)

            mse_lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(mux[prediction_step], x_data)),  # error in x
                                          tf.square(tf.subtract(muy[prediction_step], y_data))))  # error in y

            # diversity_loss = self.get_diversity_loss(mux[prediction_step], muy[prediction_step],
            #                                         sigmax[prediction_step], sigmay[prediction_step])

            # div_loss_over_horizon.append(diversity_loss)
            mse_prediction_loss_list.append(mse_lossfunc)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          # div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))
          mse_loss_list.append(tf.reduce_mean(mse_prediction_loss_list))

      self.prediction = tf.stack(self.prediction, axis=1)

      # Loss functions
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)

      self.mse_loss = tf.reduce_mean(mse_loss_list, axis=0)

      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)

      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)

      # Tensorboard summary
      self.loss_summary = tf.summary.scalar('loss', self.mse_loss)
      self.log_loss_summary = tf.summary.scalar('log likelihood loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)

      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
      self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
                                                             graph=tf.Session().graph)

      self.summary = tf.summary.merge_all()

      # Save / restore all variables
      self.full_saver = tf.train.Saver()
      
      # Only save / restore the parameters responsible for occupancy grid processing
      self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
      self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    # Dictionary to feed data to the model
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      #batch_x = kwargs['batch_x'][:, :, 2:]
      batch_x = kwargs['batch_vel']
      batch_y = kwargs['batch_y']
      #batch_y = kwargs['batch_pos_target']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, **kwargs):
    step = kwargs["step"]
    return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
            self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
            self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def feed_val_dic(self, **kwargs):
    return {self.input_state_placeholder: kwargs["batch_vel"],
            self.input_grid_placeholder: kwargs["batch_grid"],
            self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def validation_step(self,sess,feed_dict_validation):
    batch_loss, _mse_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _other_peds_states, _current_state_lstm_concat, summary = sess.run([self.total_loss,
                                                                        self.mse_loss,
                                                                        self.current_state,
                                                                        self.current_state_lstm_grid,
                                                                        self.current_state_lstm_ped,
                                                                        self.other_peds_states,
                                                                        self.current_state_lstm_concat,
                                                                        self.loss_summary],
                                                                        feed_dict=feed_dict_validation)
    self.test_cell_state_current, self.test_hidden_state_current = _current_state
    self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.test_fw_ped_states_current, self.test_bw_ped_states_current = _current_state_lstm_ped
    self.test_ped_gru_states_current = _other_peds_states
    self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _mse_loss, summary

  def train_step(self,sess,feed_dict_train):
    _, batch_loss, _mse_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.mse_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    # Train autoencoder if loss larger than threshold
    """
    if np.mean(autoencoder_loss) > 0.05:
      print(Fore.LIGHTMAGENTA_EX +"\n[{:s}] Training autoencoder. Model loss: {:01.2f}" .format(
          datetime.now().strftime("%d-%m-%y %H:%M:%S"), autoencoder_loss))

      self.run_autoencoder(sess, feed_dict_train)
      self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')
    """
    out = {"batch_loss": _mse_loss,
            "model_predictions": _model_prediction,
            "summary": _summary_str,
            "learning_rate": lr,
            "autoencoder_loss": autoencoder_loss,
           "log loss": batch_loss,
            "output_decoder": output_decoder
            }

    return out

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat
                                              ],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_fw_ped_states_current, self.test_bw_ped_states_current = _current_state_lstm_ped
      self.test_ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y,pi):

    #gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    #gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    normx = tf.subtract(x, dec_mux)
    normy = tf.subtract(y, dec_muy)
    # Calculate sx*sy
    sxsy = tf.multiply(dec_sigmax, dec_sigmay)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

    loss = tf.multiply(z + tf.log(sxsy),pi)

    return tf.reduce_mean(loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/SocialVDGNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
from colorama import Fore, Style
import datetime
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell2 import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell2 import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
  import src.data_utils.Support as sup
import numpy as np
import os

def length(sequence):
  used = tf.sign(tf.reduce_max(tf.abs(sequence), 2))
  length = tf.reduce_sum(used, 1)
  length = tf.cast(length, tf.int32)
  return length

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_ped_size = args.rnn_state_ped_size
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.n_peds = self.args.n_other_agents
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,self.n_peds, self.pedestrian_vector_dim], name='ped_grid')

    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,shape=[], name='global_step')
    self.seq_len = tf.placeholder(tf.int32, [None])

    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_ped_gru_states_current = np.zeros([args.batch_size,self.n_peds, self.rnn_state_ped_size])
    self.test_fw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_bw_ped_states_current = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.ped_gru_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size,self.n_peds, self.rnn_state_ped_size],
                                              name='cell_state_lstm_ped')
    self.fw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                              name='cell_state_lstm_ped')
    self.bw_ped_states = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
                                        name='cell_state_lstm_ped')
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=2)

        # LSTM cells to process pedestrian
        self.cell_outputs_series_lstm_ped = []
        self.current_state_lstm_ped = []
        self.other_peds_states = []
        self.other_peds_gru_outputs = []
        with tf.variable_scope('lstm_ped',reuse=tf.AUTO_REUSE) as scope:
          # TODO: try to share weights
          for i, ped in enumerate(ped_grid_series):
            # TODO: ADAPT for ped input sequence
            ped_cell = tf.contrib.rnn.GRUCell(self.rnn_state_ped_size, name='ped_basic_lstm_cell'+str(i), trainable=True)
            ped_input = tf.unstack(ped, axis=1)
            ped_gru_output, ped_gru_state = tf.contrib.rnn.static_rnn(ped_cell,ped_input ,dtype=tf.float32,initial_state=self.ped_gru_states[:,i,:])
            self.other_peds_states.append(ped_gru_state)

            if i > 0:
              for j, ped_out in enumerate(ped_gru_output):
                out_ped_concat[j] = tf.concat([out_ped_concat[j],ped_gru_output[j]],axis=1)
            else:
              out_ped_concat = ped_gru_output
          self.other_peds_gru_outputs = out_ped_concat

          self.fw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='fw_basic_lstm_cell' + str(i),
                                                    trainable=True)
          self.bw_cell_ped = tf.contrib.rnn.GRUCell(self.rnn_state_size_lstm_ped, name='bw_basic_lstm_cell' + str(i),
                                                    trainable=True)

          outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=self.fw_cell_ped,
                                                             cell_bw=self.fw_cell_ped,
                                                             inputs=tf.stack(self.other_peds_gru_outputs,axis=1),
                                                             dtype=tf.float32,
                                                             initial_state_fw=self.fw_ped_states,
                                                             initial_state_bw=self.bw_ped_states)

          self.cell_outputs_series_lstm_ped = tf.unstack(tf.transpose(tf.concat([outputs[0],outputs[1]],axis=2),[1,0,2]),axis=0)
          self.current_state_lstm_ped = states
          self.other_peds_states = tf.stack(self.other_peds_states,axis=1)

          rnn_state_size_lstm_ped = self.get_weight_variable(name='w_birnn',
                                          shape=[self.rnn_state_size_lstm_ped*2, self.rnn_state_size_lstm_ped],
                                          trainable=self.is_training)
          b_birnn = self.get_bias_variable(name='b_birnn', shape=[1, self.rnn_state_size_lstm_ped],
                                        trainable=self.is_training)

          self.fused = [tf.add(tf.matmul(lstm_out, rnn_state_size_lstm_ped), b_birnn) for lstm_out in
                     self.cell_outputs_series_lstm_ped]

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature, grid_feature in zip(self.cell_outputs_series_state, self.fused, self.cell_outputs_series_lstm_grid):
          concat_lstm.append(tf.concat([lstm_out, ped_feature, grid_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:
        # Hidden layer (fully connected) after concatenation LSTM
        W_fc = self.get_weight_variable(name='w_fc',
                                        shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
                                        trainable=self.is_training)
        b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
                                      trainable=self.is_training)

        # Output layer (fully connected)
        W_out = self.get_weight_variable(name='w_out',
                                         shape=[self.fc_hidden_unit_size,
                                                self.output_dim * self.prediction_horizon],
                                         trainable=self.is_training)
        b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
                                       trainable=self.is_training)

        # cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
                                     args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
          cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series = []
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
          dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100 + t_idx))
        for t_idx in range(len(dec_output_layer)):
          m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1)
          dec_mux_series.append(m_x)
          dec_muy_series.append(m_y)
          dec_sigmax_series.append(s_x)
          dec_sigmay_series.append(s_y)
        dec_pi_series = [dense for dense in dec_pi_series]

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

        loss_list = []
        mse_loss_list = []
        mse_loss_list2 = []
        kl_loss_list = []
        self.div_loss = []
        self.prediction = []
        self.likelihood = []
        pred = []
        div_loss_over_truncated_back_prop = []

        # Compute Loss Function
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series, dec_pi_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply probabilistic transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(dec_pi, axis=1)

          self.likelihood.append(pi)

          pred = []
          prediction_loss_list = []
          mse_prediction_loss_list = []
          mse_prediction_loss_list2 = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            # Log-likelihood
            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                             sigmax[prediction_step], sigmay[prediction_step],
                                             x_data, y_data, pi)

            mse_lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(tf.expand_dims(mux[prediction_step][:,0],axis=1), x_data)),
                                          tf.square(tf.subtract(tf.expand_dims(muy[prediction_step][:,0],axis=1), y_data))))
            mse_lossfunc2 = tf.sqrt(tf.add(tf.square(tf.subtract(tf.expand_dims(mux[prediction_step][:, 1],axis=1), x_data)),
                                                         tf.square(tf.subtract(tf.expand_dims(muy[prediction_step][:, 1],axis=1), y_data))))

            # diversity_loss = self.get_diversity_loss(mux[prediction_step], muy[prediction_step],
            #                                         sigmax[prediction_step], sigmay[prediction_step])

            # div_loss_over_horizon.append(diversity_loss)
            mse_prediction_loss_list.append(tf.minimum(mse_lossfunc,mse_lossfunc2))

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          # div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))
          mse_loss_list.append(tf.reduce_mean(mse_prediction_loss_list))


      self.prediction = tf.stack(self.prediction, axis=1)
      self.likelihood = tf.stack(self.likelihood, axis=1)

      # Loss functions
      self.total_loss = tf.reduce_mean(loss_list, axis=0)# + tf.reduce_mean(kl_loss_list, axis=0)

      self.mse_loss = tf.reduce_mean(mse_loss_list, axis=0)

      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)

      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)

      # Tensorboard summary
      self.loss_summary = tf.summary.scalar('loss', self.mse_loss)
      self.log_loss_summary = tf.summary.scalar('log likelihood loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)

      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
      self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
                                                             graph=tf.Session().graph)

      self.summary = tf.summary.merge_all()

      # Save / restore all variables
      self.full_saver = tf.train.Saver()
      
      # Only save / restore the parameters responsible for occupancy grid processing
      self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
      self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, **kwargs):
    # Dictionary to feed data to the model
    if self.args.rotated_grid:
      batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
      batch_x = batch_x[:,:,2:]
    else:
      #batch_x = kwargs['batch_x'][:, :, 2:]
      batch_x = kwargs['batch_vel']
      batch_y = kwargs['batch_y']
      #batch_y = kwargs['batch_pos_target']
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: kwargs['batch_grid'],
            self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
            self.step: kwargs['step'],
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.ped_gru_states: self.ped_gru_states_current,
            self.fw_ped_states: self.fw_ped_states_current,
            self.bw_ped_states: self.bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, **kwargs):
    step = kwargs["step"]
    return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
            self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
            self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def feed_val_dic(self, **kwargs):
    return {self.input_state_placeholder: kwargs["batch_vel"],
            self.input_grid_placeholder: kwargs["batch_grid"],
            self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
            self.step: 0,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            self.ped_gru_states: self.test_ped_gru_states_current,
            self.fw_ped_states: self.test_fw_ped_states_current,
            self.bw_ped_states: self.test_bw_ped_states_current,
            self.seq_len: [100],
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_ped_gru_states_current = np.zeros([self.args.batch_size,self.n_peds, self.args.rnn_state_ped_size])
          self.test_fw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_bw_ped_states_current = np.zeros([self.args.batch_size, self.args.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def validation_step(self,sess,feed_dict_validation):
    batch_loss, _mse_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _other_peds_states, _current_state_lstm_concat, summary = sess.run([self.total_loss,
                                                                        self.mse_loss,
                                                                        self.current_state,
                                                                        self.current_state_lstm_grid,
                                                                        self.current_state_lstm_ped,
                                                                        self.other_peds_states,
                                                                        self.current_state_lstm_concat,
                                                                        self.loss_summary],
                                                                        feed_dict=feed_dict_validation)
    self.test_cell_state_current, self.test_hidden_state_current = _current_state
    self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.test_fw_ped_states_current, self.test_bw_ped_states_current = _current_state_lstm_ped
    self.test_ped_gru_states_current = _other_peds_states
    self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, summary

  def train_step(self,sess,feed_dict_train):
    _, batch_loss, _mse_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, output_decoder, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.mse_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.other_peds_states,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.deconv1,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.fw_ped_states_current, self.bw_ped_states_current = _current_state_lstm_ped
    self.ped_gru_states_current = _other_peds_states
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    # Train autoencoder if loss larger than threshold
    """"""
    if np.mean(autoencoder_loss) > 0.05:
      self.run_autoencoder(sess, feed_dict_train)
      self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

    out = {"batch_loss": batch_loss,
            "model_predictions": _model_prediction,
            "summary": _summary_str,
            "learning_rate": lr,
            "autoencoder_loss": autoencoder_loss,
           "log loss": _mse_loss,
            "output_decoder": output_decoder
            }

    return out

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped,_other_peds_states, \
    _current_state_lstm_concat, \
    _model_prediction, output_decoder, \
    _cell_outputs_series_state, _cell_outputs_series_lstm_ped, \
    _cell_outputs_series_lstm_concat, _likelihood = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.other_peds_states,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                              self.deconv1,
                                              self.cell_outputs_series_state,
                                              self.cell_outputs_series_lstm_ped,
                                              self.cell_outputs_series_lstm_concat,
                                              self.likelihood],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_fw_ped_states_current, self.test_bw_ped_states_current = _current_state_lstm_ped
      self.test_ped_gru_states_current = _other_peds_states
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, _likelihood

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y,pi):

    #gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    #gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    normx = tf.subtract(x, dec_mux)
    normy = tf.subtract(y, dec_muy)
    # Calculate sx*sy
    sxsy = tf.multiply(dec_sigmax, dec_sigmay)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

    loss = tf.multiply(z + tf.log(sxsy),pi)

    return tf.reduce_mean(loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/SocialVDGNNFull.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight

		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32, shape=[], name='global_step')

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
				# Embedding layer of pedestrian grid
				ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped,
					                                                                                           ped_grid_feature_series,
					                                                                                           dtype=tf.float32,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
					concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			# Concatenated LSTM layer
			with tf.variable_scope('generation_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out',
				                                 shape=[self.fc_hidden_unit_size, self.output_dim * self.prediction_horizon],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim * self.prediction_horizon],
				                               trainable=self.is_training)

				# Concatenated LSTM layer
				with tf.variable_scope('generation_model') as scope:

					# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
					                             args.latent_space_size)
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

					# Output reshape
					enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
					enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
					self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
					prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
					prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

					# Decoder layers
					dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
					dec_output_layer = []
					dec_mux_series = []
					dec_muy_series = []
					dec_sigmax_series = []
					dec_sigmay_series = []
					dec_pi_series = []
					for t_idx in range(len(dec_fc_layer)):
						dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
						dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100 + t_idx))
					for t_idx in range(len(dec_output_layer)):
						m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1)
						dec_mux_series.append(m_x)
						dec_muy_series.append(m_y)
						dec_sigmax_series.append(s_x)
						dec_sigmay_series.append(s_y)
					dec_pi_series = [dense for dense in dec_pi_series]

					# Exponential / stepwise decrease of learning rate
					global_step = tf.Variable(0, trainable=False)
					self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
					                                                decay_rate=0.9, staircase=False)

					self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

				loss_list = []
				mse_loss_list = []
				kl_loss_list = []
				self.div_loss = []
				self.prediction = []
				pred = []
				div_loss_over_truncated_back_prop = []

				# Compute Loss Function
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series, dec_pi_series,
						    outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(dec_pi, axis=1)

					pred = []
					prediction_loss_list = []
					mse_prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						# Log-likelihood
						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                 sigmax[prediction_step], sigmay[prediction_step],
						                                 x_data, y_data, pi)

						mse_lossfunc = tf.sqrt(tf.add(tf.square(tf.subtract(mux[prediction_step][:,0], x_data)),  # error in x
						                          tf.square(tf.subtract(muy[prediction_step][:,0], y_data))))  # error in y

						# diversity_loss = self.get_diversity_loss(mux[prediction_step], muy[prediction_step],
						#                                         sigmax[prediction_step], sigmay[prediction_step])

						# div_loss_over_horizon.append(diversity_loss)
						mse_prediction_loss_list.append(mse_lossfunc)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
					kl_loss_list.append(kl_loss)
					# div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))
					mse_loss_list.append(tf.reduce_mean(mse_prediction_loss_list))

			self.prediction = tf.stack(self.prediction, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0)

			self.mse_loss = tf.reduce_mean(mse_loss_list, axis=0)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)

			# Tensorboard summary
			self.loss_summary = tf.summary.scalar('loss', self.mse_loss)
			self.log_loss_summary = tf.summary.scalar('log likelihood loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)

			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
			self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
			                                                       graph=tf.Session().graph)

			self.summary = tf.summary.merge_all()

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.step: kwargs['step'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
	        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
	        self.step: 0,
	        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
	        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
	        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
	        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
	        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
	        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
	        }

	def feed_val_dic(self, **kwargs):
		return {self.input_state_placeholder: kwargs['batch_vel'],
		        self.input_ped_grid_placeholder: kwargs['batch_ped_grid'],
		        self.output_placeholder: kwargs['batch_y'],
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def validation_step(self, sess, feed_dict_validation):
		batch_loss, _mse_loss,_current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, summary = sess.run([self.total_loss,
		                                                self.mse_loss,
		                                                self.current_state,
		                                                self.current_state_lstm_ped,
		                                                self.current_state_lstm_concat,
		                                                self.loss_summary],
		                                               feed_dict=feed_dict_validation)
		self.test_cell_state_current, self.test_hidden_state_current = _current_state
		self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _mse_loss, summary

	def train_step(self, sess, feed_dict_train):
		_, batch_loss, _mse_loss, _current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr = sess.run([self.update,
		                                                self.total_loss,
		                                                self.mse_loss,
		                                                self.current_state,
		                                                self.current_state_lstm_ped,
		                                                self.current_state_lstm_concat,
		                                                self.prediction,
		                                                self.summary,
		                                                self.learning_rate],
		                                               feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		out = {"batch_loss": _mse_loss,
		       "model_predictions": _model_prediction,
		       "summary": _summary_str,
		       "learning_rate": lr,
		       "log loss": batch_loss
		       }

		return out

	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update], feed_dict=feed_dict_train)

	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction = sess.run([self.current_state,
		                              self.current_state_lstm_ped,
		                              self.current_state_lstm_concat,
		                              self.prediction],
		                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print("Error. Encoder not found...")
			exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True

	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat(
					[tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss

	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result

	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y, pi):

		# gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		# gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		loss = tf.multiply(z + tf.log(sxsy), pi)

		return tf.reduce_mean(loss)

	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):

		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation

	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/VDGNN_simple.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os
from colorama import Fore, Style

class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.kl_weight = args.kl_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_size)
		self.grid_height = int(args.submap_size)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
													  shape=[self.batch_size, self.truncated_backprop_length,
															 self.input_state_dim * (self.prev_horizon + 1)],
													  name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
													 shape=[self.batch_size, self.truncated_backprop_length,
															self.grid_width, self.grid_height], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
														 shape=[self.batch_size, self.truncated_backprop_length,
																self.pedestrian_vector_dim*self.args.n_other_agents], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
																		  self.output_placeholder_dim], name='output')
		self.diversity_placeholder = tf.placeholder(dtype=tf.float32,
													shape=[self.batch_size, self.truncated_backprop_length,
														   self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32,
								   shape=[], name='global_step')
		self.beta = tf.placeholder(dtype=tf.float32,
								   shape=[])
		self.seq_length = tf.placeholder(tf.int32, [None])
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
										 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
										   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
												   name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
													 name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
																		self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
												  name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
													name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
													 shape=[self.batch_size, self.rnn_state_size_lstm_concat],
													 name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
													   name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
																		  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)
		outputs_diversity_series = tf.unstack(self.diversity_placeholder, axis=1)
		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))
		
		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = self.rnn_state_size_lstm_ped
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
																  shape=[ped_grid_out_dim, ped_grid_out_dim],
																  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
															 trainable=True)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
																												 conv_grid_feature_series,
																												 dtype=tf.float32,
																												 initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					ped_grid_feature_series , self.current_state_lstm_ped = tf.nn.dynamic_rnn(self.cell_ped,
																											   tf.stack(ped_grid_series,axis=1),
																											   dtype=tf.float32,
																											   sequence_length=self.seq_length,
																											   initial_state=self.init_state_tuple_lstm_ped)

					# Embedding layer of pedestrian grid
					self.cell_outputs_series_lstm_ped = self.process_pedestrian_grid(tf.unstack(ped_grid_feature_series,axis=1))

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
																								   dtype=tf.float32,
																								   initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state,
															   self.cell_outputs_series_lstm_grid,
															   self.cell_outputs_series_lstm_ped):
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))
					
					# Concatenated LSTM layer
					with tf.variable_scope('generation_model') as scope:

						#cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
						cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
						                             args.latent_space_size)
						self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
							cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

						# Output reshape
						enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

						# Decoder layers
						dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series]))) # ELU activation in decoder
						dec_output_layer = []
						dec_mux_series = []
						dec_muy_series = []
						dec_sigmax_series = []
						dec_sigmay_series = []
						dec_pi_series = []
						# Reusing weights for each truncated back propagation step
						for t_idx in range(len(dec_fc_layer)):
							dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, 50)) # Linear activation in decoder
							dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100))
						for t_idx in range(len(dec_output_layer)):
							m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1) # Prediction
							dec_mux_series.append(m_x)
							dec_muy_series.append(m_y)
							dec_sigmax_series.append(s_x)
							dec_sigmay_series.append(s_y)
						dec_pi_series = [dense for dense in dec_pi_series]

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
																decay_rate=0.9, staircase=False)

				self.KLWEIGHT = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2	# sigmoid	

				prediction_loss_list = []
				loss_list = []
				kl_loss_list = []
				self.div_loss = []
				self.prediction = []
				self.likelihood = []
				pred = []
				div_loss_over_truncated_back_prop = []

				# Compute Loss Function
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out, out_div in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
							dec_sigmax_series, dec_sigmay_series,
							prior_mu_series, prior_sigma_series, dec_pi_series,
							outputs_series, outputs_diversity_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(dec_pi, axis=1) # Prediction of Mode likelihood

					self.likelihood.append(pi)
					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
						x_diversity = tf.reshape(out_div[:, idx_x], shape=(-1, 1))
						y_diversity = tf.reshape(out_div[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
														 sigmax[prediction_step], sigmay[prediction_step],
														 x_data, y_data, pi)
						#
						#diversity_loss = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						#                                        sigmax[prediction_step], sigmay[prediction_step],
						#                                       x_diversity, y_diversity, pi)
						#
						#div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss_list.append(kl_loss)
					#
					#div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			self.prediction = tf.stack(self.prediction, axis=1)
			self.likelihood = tf.stack(self.likelihood, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars if 'bias' or 'auto_encoder' not in tvar.name)
			
			# Reduce mean in all dimensions
			#self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
			self.total_loss = tf.reduce_mean(loss_list, axis=0) + (tf.reduce_mean(kl_loss_list, axis=0) * self.KLWEIGHT) #+ l2 # Add Diversity loss: + args.diversity_update*self.div_loss
			self.reconstruction_loss = tf.reduce_mean(loss_list, axis=0)
			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

			# Optimizer specification
			#self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			self.loss_summary = tf.summary.scalar('loss', self.total_loss)
			self.kl_loss_summary = tf.summary.scalar('kl-loss', self.kl_loss)
			#self.div_loss_summary = tf.summary.scalar('div_loss', self.div_loss)
			self.reconstruction_loss_summary = tf.summary.scalar('reconstruction_loss', self.reconstruction_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			tf.summary.scalar('beta', self.beta)
			tf.summary.scalar('klweight', self.KLWEIGHT)

			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
			self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
																   graph=tf.Session().graph)

			self.summary = tf.summary.merge_all()

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
			self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		n_other_agents = np.zeros([self.args.batch_size])
		for i,other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: batch_x,
				self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
				self.input_grid_placeholder: kwargs["batch_grid"],
				self.step: kwargs['step'],
				self.beta: kwargs['beta'],
				self.seq_length: n_other_agents,
				self.diversity_placeholder: kwargs['batch_div'],
				self.output_placeholder: batch_y,
				self.cell_state: self.cell_state_current,
				self.hidden_state: self.hidden_state_current,
				self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
				self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
				self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
				self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
				self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
				self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
				}
	
	def feed_pred_dic(self, **kwargs):
		step = kwargs["step"]
		n_other_agents = np.ones([self.args.batch_size])

		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
				self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
				self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :], axis=1),
				self.step: 0,
				self.beta: 0,
				self.seq_length: n_other_agents,
				self.cell_state: self.test_cell_state_current,
				self.hidden_state: self.test_hidden_state_current,
				self.cell_state_lstm_grid: self.test_cell_state_current_lstm_grid,
				self.hidden_state_lstm_grid: self.test_hidden_state_current_lstm_grid,
				self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
				self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
				self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
				self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
				}

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		n_other_agents = np.zeros([self.args.batch_size])
		#print(kwargs.keys())
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		#print("SEQLEN INPUT: ", n_other_agents)
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
				self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
				self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
				self.step: 0,
				self.beta: 0,
				self.seq_length: n_other_agents,
				self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
				self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
				self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
				self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
				self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
				self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
				self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
				self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
				}

	def feed_val_dic(self, n_other_agent = None, **kwargs):
		n_other_agents = np.zeros([self.args.batch_size])
		if n_other_agent is None:
			for i, other in enumerate(kwargs['other_agents_pos']):
				n_other_agents[i] = len(other)
		return {self.input_state_placeholder: kwargs["batch_vel"],
				self.input_grid_placeholder: kwargs["batch_grid"],
				self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
				self.diversity_placeholder: kwargs['batch_div'],
				self.output_placeholder: kwargs['batch_y'],
				self.seq_length: n_other_agents,
				self.step: 0,
				self.beta: 0,
				self.cell_state: np.random.normal(self.test_cell_state_current.copy(), kwargs["state_noise"]),
				self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), kwargs["state_noise"]),
				self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
				self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
				self.cell_state_lstm_ped: np.random.normal(self.test_cell_state_current_lstm_ped.copy(), kwargs["ped_noise"]),
				self.hidden_state_lstm_ped: np.random.normal(self.test_hidden_state_current_lstm_ped.copy(), kwargs["ped_noise"]),
				self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
				self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
				}

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])


	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def update_test_hidden_state(self):
		self.test_cell_state_current = self.cell_state_current.copy()
		self.test_hidden_state_current = self.hidden_state_current.copy()
		self.test_cell_state_current_lstm_grid = self.cell_state_current_lstm_grid.copy()
		self.test_hidden_state_current_lstm_grid = self.hidden_state_current_lstm_grid.copy()
		self.test_cell_state_current_lstm_ped = self.cell_state_current_lstm_ped.copy()
		self.test_hidden_state_current_lstm_ped = self.hidden_state_current_lstm_ped.copy()
		self.test_cell_state_current_lstm_concat = self.cell_state_current_lstm_concat.copy()
		self.test_hidden_state_current_lstm_concat = self.hidden_state_current_lstm_concat.copy()

	def validation_step(self, sess, feed_dict_validation, update = True):
		# Added total loss (Walter)
		batch_loss, reconstr_loss, _model_prediction,_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, summary = sess.run([self.total_loss, self.reconstruction_loss,
														self.prediction,
																			self.current_state,
																			self.current_state_lstm_grid,
																			self.current_state_lstm_ped,
																			self.current_state_lstm_concat,
																			self.loss_summary],
																		   feed_dict=feed_dict_validation)

		if update:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, summary, _model_prediction

	def train_step(self, sess, feed_dict_train, step=0):

		# Diversity Training
		if step > 2000 and self.args.diversity_update:
			
			dict = {"batch_vel": feed_dict_train[self.input_state_placeholder],
					"batch_grid": feed_dict_train[self.input_grid_placeholder],
					"batch_ped_grid": feed_dict_train[self.input_ped_grid_placeholder],
					"batch_y": feed_dict_train[self.output_placeholder],
					"batch_div": feed_dict_train[self.diversity_placeholder],
					"step": step,
					"beta": feed_dict_train[self.beta],
					"state_noise": 0.2,
					"grid_noise": 0.0,
					"ped_noise": 0.2
					}
			self.update_test_hidden_state()
			feed_test_dic = self.feed_val_dic(n_other_agent=feed_dict_train[self.seq_length],**dict)
			loss, _, y_model_pred = self.validation_step(sess, feed_test_dic, False)
			batch_y_varible = np.zeros((self.args.batch_size, self.args.truncated_backprop_length,
										self.args.prediction_horizon * self.args.output_dim))
			# Grount truth
			batch_y = feed_dict_train[self.output_placeholder]

			diversity_distance = 0
			diff_traj_index = 0
			# Find most different trajectory index
			for batch_id in range(self.args.batch_size):
				for t_id in range(self.args.truncated_backprop_length):
					for mix_id in range(self.args.n_mixtures):
						for pred_id in range(self.args.prediction_horizon):
							new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures+ mix_id]
							new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + mix_id]
							batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
								[new_batch_x, new_batch_y])
						if diversity_distance < np.linalg.norm(batch_y_varible[batch_id, t_id]-batch_y[batch_id, t_id]):
							diversity_distance = np.linalg.norm(batch_y_varible[batch_id, t_id] - batch_y[batch_id, t_id])
							diff_traj_index = mix_id
					# Saving most different traj
					for pred_id in range(self.args.prediction_horizon):
						new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + diff_traj_index]
						new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + diff_traj_index]
						batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
							[new_batch_x, new_batch_y])

			# Assemble feed dict for training
			feed_dict_train[self.diversity_placeholder] = batch_y_varible

		# kl_loss, div_loss
		_, batch_loss, kl_loss, recons_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, output_decoder, autoencoder_loss = sess.run([self.update,
																										  self.total_loss,
																										  self.kl_loss,
																										  self.reconstruction_loss,
																										  self.current_state,
																										  self.current_state_lstm_grid,
																										  self.current_state_lstm_ped,
																										  self.current_state_lstm_concat,
																										  self.prediction,
																										  self.summary,
																										  self.learning_rate,
																										  self.beta,
																										  self.deconv1,
																										  #self.div_loss,
																										  self.autoencoder_loss],
																										 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		# Train autoencoder if loss larger than threshold
		"""
		if np.mean(autoencoder_loss) > 0.05:
			self.run_autoencoder(sess, feed_dict_train)
			self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')
		"""
		out = {"batch_loss": recons_loss,
			   "model_predictions": _model_prediction,
			   "summary": _summary_str,
			   "learning_rate": lr,
			   "autoencoder_loss": autoencoder_loss,
			   "log loss": batch_loss,
			   "kl_loss": kl_loss,
			   "output_decoder": output_decoder
			   }
		

		return out


	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
					 feed_dict=feed_dict_train)


	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, _likelihood = sess.run([self.current_state,
													  self.current_state_lstm_grid,
													  self.current_state_lstm_ped,
													  self.current_state_lstm_concat,
													  self.prediction,
													  self.deconv1,
												  	  self.likelihood],
													 feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, _likelihood


	def warmstart_model(self, args, sess):
		# Restore whole model
		try:
			print('Loading session from "{}"'.format(args.model_path))
			ckpt = tf.train.get_checkpoint_state(args.model_path)
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.full_saver.restore(sess, ckpt.model_checkpoint_path)
		except:
			print(Fore.RED + "Model not found..."+Style.RESET_ALL)
			sys.exit()


	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print(Fore.RED +"Error. Encoder not found...")
			sys.exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True


	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)


	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)],
								   1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss


	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result


	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y, pi):

		# gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		# gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		#loss = tf.multiply(z + tf.log(sxsy), pi)
		loss = z + tf.log(sxsy)

		return tf.reduce_mean(loss)


	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):
		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #


	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
													 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
													 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
													 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
															   name="fc_grid_weights",
															   trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
															trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
																name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
										   shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
																 [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
																 padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
																 [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
																 padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
																	  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
												 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
												 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
												 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
												 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors


	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector


	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
							  initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
																						 uniform=False, seed=None,
																						 dtype=tf.float32),
							  regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
							  initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
																						 uniform=False, seed=None,
																						 dtype=tf.float32),
							  trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
							name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations


	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
										padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation


	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		#print(input.size, weights.size)
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

	def frange_cycle_linear(self, n_iter, start=0.0, stop=1.0,  n_cycle=4, ratio=0.5):
		L = np.ones(n_iter) * stop
		period = n_iter/n_cycle
		step = (stop-start)/(period*ratio) # linear schedule

		for c in range(n_cycle):
			v, i = start, 0
			while v <= stop and (int(i+c*period) < n_iter):
				L[int(i+c*period)] = v
				v += step
				i += 1
		return L

src/models/VGDNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class VGDNNModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = args.grid_width
    self.grid_height = args.grid_height
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)

    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      with tf.variable_scope('encoder_model') as scope:
        ###################### Assemble Neural Net #####################
        # Assemble CNN
        self.createCNNWeights()

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid):
          concat_lstm.append(tf.concat([lstm_out, grid_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args,args.rnn_state_size_lstm_concat,self.output_vec_dim, args.prior_size, args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.n_mixtures > 1:
            mux, muy, sigmax, sigmay, pi = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(mux)
            dec_muy_series.append(muy)
            dec_sigmax_series.append(tf.exp(sigmax))
            dec_sigmay_series.append(tf.exp(sigmay))
            dec_pi_series.append(tf.nn.softmax(pi))
          else:
            mux, muy, sigmax, sigmay = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(mux)
            dec_muy_series.append(muy)
            dec_sigmax_series.append(tf.exp(sigmax))
            dec_sigmay_series.append(tf.exp(sigmay))

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=500,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 5000) / 1000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.prediction = []
      pred= []

      if args.n_mixtures > 1:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.split(dec_sigmax, self.prediction_horizon, axis=1)
          sigmay = tf.split(dec_sigmay, self.prediction_horizon, axis=1)
          pi = tf.split(dec_pi, self.prediction_horizon, axis=1)

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)

          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.split(dec_sigmax, self.prediction_horizon, axis=1)
          sigmay = tf.split(dec_sigmay, self.prediction_horizon, axis=1)

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)

          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.total_loss = tf.reduce_mean(loss_list, axis=0)
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
      # Only save / restore the parameters responsible for occupancy grid processing
      self.convnet_saver = tf.train.Saver(var_list={'conv1_weights': self.conv1_weights,
                                                    'conv1_biases': self.conv1_biases,
                                                    'conv2_weights': self.conv2_weights,
                                                    'conv2_biases': self.conv2_biases,
                                                    'conv3_weights:': self.conv3_weights,
                                                    'conv3_biases': self.conv3_biases,
                                                    'fc_grid_weights': self.fc_grid_weights,
                                                    'fc_grid_biases': self.fc_grid_biases,
                                                    })

  def createCNNWeights(self):
    # CNN weights for static occupancy grid processing
    train_grid_encoder_conv = False
    train_grid_encoder_fc = self.args.end_to_end
    use_summary_convnet = False
    train_encoder = self.args.end_to_end
    # Convolutional layer 1
    self.conv1_kernel_size = 5
    self.conv1_number_filters = 64
    self.conv1_stride_length = 2
    self.conv1_weights = self.get_weight_variable(name="conv1_weights",
                                                  shape=[self.conv1_kernel_size, self.conv1_kernel_size, 1,
                                                         self.conv1_number_filters], trainable=train_grid_encoder_conv,
                                                  summary=use_summary_convnet)
    self.conv1_biases = self.get_bias_variable(name="conv1_biases", shape=[self.conv1_number_filters],
                                               trainable=train_grid_encoder_conv, summary=use_summary_convnet)

    # Convolutional layer 2
    self.conv2_kernel_size = 3
    self.conv2_number_filters = 32
    self.conv2_stride_length = 2
    self.conv2_weights = self.get_weight_variable(name="conv2_weights",
                                                  shape=[self.conv2_kernel_size, self.conv2_kernel_size,
                                                         self.conv1_number_filters, self.conv2_number_filters],
                                                  trainable=train_grid_encoder_conv, summary=use_summary_convnet)
    self.conv2_biases = self.get_bias_variable(name="conv2_biases", shape=[self.conv2_number_filters],
                                               trainable=train_grid_encoder_conv, summary=use_summary_convnet)

    # Convolutional layer 3
    self.conv3_kernel_size = 3
    self.conv3_number_filters = 8
    self.conv3_stride_length = 2
    self.conv3_weights = self.get_weight_variable(name="conv3_weights",
                                                  shape=[self.conv3_kernel_size, self.conv3_kernel_size,
                                                         self.conv2_number_filters, self.conv3_number_filters],
                                                  trainable=train_grid_encoder_conv, summary=use_summary_convnet)
    self.conv3_biases = self.get_bias_variable(name="conv3_biases", shape=[self.conv3_number_filters],
                                               trainable=train_grid_encoder_conv, summary=use_summary_convnet)

    fc_grid_hidden_dim = 64
    self.fc_grid_weights = self.get_weight_variable(shape=[512, fc_grid_hidden_dim], name="fc_grid_weights",
                                                    trainable=train_grid_encoder_fc)
    self.fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                 trainable=train_grid_encoder_fc)

    # Process occupancy grid, apply convolutional filters
    occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = self.process_grid(occ_grid_series)

    # LSTM cells to process static occupancy grid
    with tf.variable_scope('lstm_grid') as scope:
      self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                               trainable=train_encoder)
      self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                   conv_grid_feature_series,
                                                                                                   dtype=tf.float32,
                                                                                                   initial_state=self.init_state_tuple_lstm_grid)

  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x[:,:,2:],
            self.input_grid_placeholder: batch_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss, _current_state, _current_state_lstm_grid, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta = sess.run([self.update,
                                                          self.total_loss,
                                                          #self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, _model_prediction, _summary_str, lr, beta

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, \
    _current_state_lstm_concat, \
    _model_prediction = sess.run([self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_concat,
                                                          self.prediction],
                                                          feed_dict=feed_dict_train)
    if update_state:
      self.cell_state_current, self.hidden_state_current = _current_state
      self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    
    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)
      conv1 = self.conv_layer(input=grid_batch, weights=self.conv1_weights, biases=self.conv1_biases, conv_stride_length=self.conv1_stride_length, name="conv1_grid")
      conv2 = self.conv_layer(input=conv1, weights=self.conv2_weights, biases=self.conv2_biases, conv_stride_length=self.conv2_stride_length, name="conv2_grid")
      conv3 = self.conv_layer(input=conv2, weights=self.conv3_weights, biases=self.conv3_biases, conv_stride_length=self.conv3_stride_length, name="conv3_grid")
      conv_grid_size = 8
      conv5_flat = tf.reshape(conv3, [-1, conv_grid_size*conv_grid_size*self.conv3_number_filters])
      fc_final = self.fc_layer(input=conv5_flat, weights=self.fc_grid_weights, biases=self.fc_grid_biases, activation = tf.nn.sigmoid, name="fc_grid")
      if idx == 0:
        tf.summary.histogram("fc_activations", fc_final)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)
    
    return conv_feature_vectors

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_Grid.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from storn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.storn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        try:
          cell_lstm_concat = vrnn_cell(args,args.x_dim,self.rnn_state_size_lstm_concat, args.prior_size, args.latent_space_size)
        except:
          cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.rnn_state_size_lstm_concat, args.prior_size,
                                       args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.maximum(tf.minimum(tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1)),0.5),0.001)
          sigmay = tf.maximum(tf.minimum(tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1)),0.5),0.001)
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step],muy[prediction_step],sigmax[prediction_step],sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step], muy[prediction_step],
                                                     sigmax[prediction_step], sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss*pow(0.9,-prediction_step))

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)# - tf.maximum(self.div_loss,1.0)
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)
      self.div_gradients = tf.gradients(-0.1*tf.minimum(self.div_loss,2.0), train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
      self.div_grads, _ = tf.clip_by_global_norm(self.div_gradients, 0.01)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      self.div_update = self.optimizer.apply_gradients(zip(self.div_grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_target):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.output_placeholder: batch_target,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, _,batch_loss, kl_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
                                                          self.div_update,
                                                          self.total_loss,
                                                          self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta,
                                                          self.deconv1,
                                                          self.div_loss,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, kl_loss,_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, loss, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                  self.total_loss,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, loss, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures-1):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)

      mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
      sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
      gaussian2 = tfd.MultivariateNormalDiag(
        loc=mu2,
        scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian2,gaussian1,allow_nan_stats=False)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")

      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_KL.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from attention import *
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.external.attention import *
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

        W_wq_concat = self.get_weight_variable(name='w_wq_concat',
                                            shape=[self.prior_hidden_size, self.prior_hidden_size],
                                            trainable=True)
        b_wq_concat = self.get_bias_variable(name='b_wq_concat', shape=[1, self.prior_hidden_size], trainable=True)

        W_wk_concat = self.get_weight_variable(name='w_wk_concat',
                                            shape=[self.prior_hidden_size, self.prior_hidden_size],
                                            trainable=True)
        b_wk_concat = self.get_bias_variable(name='b_wk_concat', shape=[1, self.prior_hidden_size], trainable=True)

        W_wv_concat = self.get_weight_variable(name='w_wv_concat',
                                            shape=[self.prior_hidden_size, self.prior_hidden_size],
                                            trainable=True)
        b_wv_concat = self.get_bias_variable(name='b_wv_concat', shape=[1, self.prior_hidden_size], trainable=True)

        wq_concat = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wq_concat, b_wq_concat) for output in concat_lstm])))
        wk_concat = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wk_concat, b_wk_concat) for output in concat_lstm])))
        wv_concat = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wv_concat, b_wv_concat) for output in concat_lstm])))

        concat_scaled_attention = []
        for state, key, value in zip(wq_concat, wk_concat, wv_concat):
          concat_scaled_output, attention_weights = scaled_dot_product_attention(tf.expand_dims(state,axis=1), tf.expand_dims(key,axis=1), tf.expand_dims(value,axis=1), mask=None)
          concat_scaled_attention.append(tf.squeeze(concat_scaled_output,axis=1))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        try:
          cell_lstm_concat = vrnn_cell(args,args.x_dim,self.rnn_state_size_lstm_concat, args.prior_size, args.latent_space_size)
        except:
          cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.rnn_state_size_lstm_concat, args.prior_size,
                                       args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_scaled_attention, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step],muy[prediction_step],sigmax[prediction_step],sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_target,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.output_placeholder: batch_target,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: np.random.normal(self.test_cell_state_current_lstm_ped.copy(),0),
            self.hidden_state_lstm_ped: np.random.normal(self.test_hidden_state_current_lstm_ped.copy(),0),
            self.cell_state_lstm_concat: np.random.normal(self.test_cell_state_current_lstm_concat.copy(),noise),
            self.hidden_state_lstm_concat: np.random.normal(self.test_hidden_state_current_lstm_concat.copy(),noise),
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss, kl_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta,
                                                          self.deconv1,
                                                          self.div_loss,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, kl_loss,_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, loss, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                  self.total_loss,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, loss, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")

      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_attention.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from attention import *
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.external.attention import *
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.diversity_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
                                                                      self.output_placeholder_dim], name='output')
    self.positions_output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
                                                                      self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    outputs_diversity_series = tf.unstack(self.diversity_placeholder, axis=1)
    positions_series = tf.unstack(self.positions_output_placeholder, axis=1)
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=True)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        try:
          cell_lstm_concat = vrnn_cell(args,args.x_dim,self.rnn_state_size_lstm_concat, args.prior_size, args.latent_space_size)
        except:
          cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.rnn_state_size_lstm_concat, args.prior_size,
                                       args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      position_loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out, out_div in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series, outputs_diversity_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
            x_diversity = tf.reshape(out_div[:, idx_x], shape=(-1, 1))
            y_diversity = tf.reshape(out_div[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_diversity,y_diversity)

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out, pos, out_div in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series, positions_series, outputs_diversity_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.minimum(tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1)),5.0)
          sigmay = tf.minimum(tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1)),5.0)

          pred = []
          prediction_loss_list = []
          div_loss_over_horizon = []
          kl_mode_loss_list = []
          pos_loss_list = []
          lossfunc = 0
          x_pos = 0
          y_pos = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
            x_pos_data = tf.reshape(pos[:, idx_x], shape=(-1, 1))
            y_pos_data = tf.reshape(pos[:, idx_y], shape=(-1, 1))
            x_diversity = tf.reshape(out_div[:, idx_x], shape=(-1, 1))
            y_diversity = tf.reshape(out_div[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)
            diversity_loss = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                       sigmax[prediction_step], sigmay[prediction_step],
                                                   x_diversity, y_diversity)

            x_pos += mux[prediction_step][:,0] * self.args.dt
            y_pos += muy[prediction_step][:,0] * self.args.dt

            pos_loss = tf.sqrt(tf.add(tf.square(tf.subtract(x_data, mux[prediction_step][:,0])),  # error in x
                                    tf.square(tf.subtract(x_data, muy[prediction_step][:,0]))))  # error in y

            pos_loss_list.append(pos_loss)

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          position_loss_list.append(tf.reduce_sum(pos_loss_list))
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.position_loss = tf.reduce_mean(position_loss_list)
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)+0.05*self.div_loss  #+ self.position_loss
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y,batch_pos,batch_div):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.positions_output_placeholder: batch_pos,
            self.diversity_placeholder: batch_div,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_target,batch_pos,noise=0.6):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.output_placeholder: batch_target,
            self.positions_output_placeholder: batch_pos,
            self.diversity_placeholder: batch_target,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: np.random.normal(self.test_cell_state_current_lstm_ped.copy(),0),
            self.hidden_state_lstm_ped: np.random.normal(self.test_hidden_state_current_lstm_ped.copy(),0),
            self.cell_state_lstm_concat: np.random.normal(self.test_cell_state_current_lstm_concat.copy(),noise),
            self.hidden_state_lstm_concat: np.random.normal(self.test_hidden_state_current_lstm_concat.copy(),noise),
            }

  def update_test_hidden_state(self):
    self.test_cell_state_current = self.cell_state_current.copy()
    self.test_hidden_state_current = self.hidden_state_current.copy()
    self.test_cell_state_current_lstm_grid = self.cell_state_current_lstm_grid.copy()
    self.test_hidden_state_current_lstm_grid = self.hidden_state_current_lstm_grid.copy()
    self.test_cell_state_current_lstm_ped = self.cell_state_current_lstm_ped.copy()
    self.test_hidden_state_current_lstm_ped = self.hidden_state_current_lstm_ped.copy()
    self.test_cell_state_current_lstm_concat = self.cell_state_current_lstm_concat.copy()
    self.test_hidden_state_current_lstm_concat = self.hidden_state_current_lstm_concat.copy()

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss, kl_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta, output_decoder, pos_loss, div_loss, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta,
                                                          self.deconv1,
                                                          self.position_loss,
                                                          self.div_loss,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, kl_loss,_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, loss, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                  self.total_loss,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, loss, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")

      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_attention_diversity.py
eliatrevisan/IA-MPPI-LBM
from tkinter import Scale
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os
from colorama import Fore, Style

class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.kl_weight = args.kl_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
													  shape=[self.batch_size, self.truncated_backprop_length,
															 self.input_state_dim * (self.prev_horizon + 1)],
													  name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
													 shape=[self.batch_size, self.truncated_backprop_length,
															self.grid_width, self.grid_height], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
														 shape=[self.batch_size, self.truncated_backprop_length,
																self.pedestrian_vector_dim*self.args.n_other_agents], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
																		  self.output_placeholder_dim], name='output')
		self.diversity_placeholder = tf.placeholder(dtype=tf.float32,
													shape=[self.batch_size, self.truncated_backprop_length,
														   self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32,
								   shape=[], name='global_step')
		self.beta = tf.placeholder(dtype=tf.float32,
								   shape=[])
		self.seq_length = tf.placeholder(tf.int32, [None])
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
										 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
										   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
												   name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
													 name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
																		self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
												  name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
													name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
													 shape=[self.batch_size, self.rnn_state_size_lstm_concat],
													 name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
													   name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
																		  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)
		outputs_diversity_series = tf.unstack(self.diversity_placeholder, axis=1)
		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
																  shape=[ped_grid_out_dim, ped_grid_out_dim],
																  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
															 trainable=True)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
																												 conv_grid_feature_series,
																												 dtype=tf.float32,
																												 initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					ped_grid_feature_series , self.current_state_lstm_ped = tf.nn.dynamic_rnn(self.cell_ped,
																											   tf.stack(ped_grid_series,axis=1),
																											   dtype=tf.float32,
																											   sequence_length=self.seq_length,
																											   initial_state=self.init_state_tuple_lstm_ped)

					# Embedding layer of pedestrian grid
					self.cell_outputs_series_lstm_ped = self.process_pedestrian_grid(tf.unstack(ped_grid_feature_series,axis=1))

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
																								   dtype=tf.float32,
																								   initial_state=self.init_state_tuple)
				
				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state,
															   self.cell_outputs_series_lstm_grid,
															   self.cell_outputs_series_lstm_ped):
					
					# Implement the diversity described in the paper
					#sigma_div = 1.0
					#lst#m_out_sampled = tfd.MultivariateNormalDiag(
						#loc=lstm_out,
						#scale_identity_multiplier=[0.2]
						#scale_diag=sigma_div
					#).sample()
					#grid_feature_sampled = tfd.MultivariateNormalDiag(
						#loc=grid_feature,
						#scale_identity_multiplier=[0.2]
						#scale_diag=sigma_div
					#).sample()
					#ped_feature_sampled = tfd.MultivariateNormalDiag(
					#	loc=ped_feature,
						#scale_diag=sigma_div
					#).sample()
					
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

					# Concatenated LSTM layer
					with tf.variable_scope('generation_model') as scope:

						# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
						cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
													 args.latent_space_size)
						self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
							cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

						# Output reshape
						enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

						# Decoder layers
						dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series]))) # ELU activation in decoder
						dec_output_layer = []
						dec_mux_series = []
						dec_muy_series = []
						dec_sigmax_series = []
						dec_sigmay_series = []
						dec_pi_series = []
						# Reusing weights for each truncated back propagation step
						for t_idx in range(len(dec_fc_layer)):
							dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, 50)) # Linear activation in decoder
							dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100))
						for t_idx in range(len(dec_output_layer)):
							m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1) # Prediction
							dec_mux_series.append(m_x)
							dec_muy_series.append(m_y)
							dec_sigmax_series.append(s_x)
							dec_sigmay_series.append(s_y)
						dec_pi_series = [dense for dense in dec_pi_series]

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
																decay_rate=0.9, staircase=False)

				#self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2	# sigmoid	

				prediction_loss_list = []
				loss_list = []
				kl_loss_list = []
				self.div_loss = []
				self.prediction = []
				self.likelihood = []
				pred = []
				div_loss_over_truncated_back_prop = []

				# Compute Loss Function
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out, out_div in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
							dec_sigmax_series, dec_sigmay_series,
							prior_mu_series, prior_sigma_series, dec_pi_series,
							outputs_series, outputs_diversity_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(dec_pi, axis=1) # Prediction of Mode likelihood

					self.likelihood.append(pi)
					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
						x_diversity = tf.reshape(out_div[:, idx_x], shape=(-1, 1))
						y_diversity = tf.reshape(out_div[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
														 sigmax[prediction_step], sigmay[prediction_step],
														 x_data, y_data, pi)
						#
						diversity_loss = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
																sigmax[prediction_step], sigmay[prediction_step],
															   x_diversity, y_diversity, pi)
						#
						div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss_list.append(kl_loss)
					#
					div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			self.prediction = tf.stack(self.prediction, axis=1)
			self.likelihood = tf.stack(self.likelihood, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
			self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)*self.beta + args.diversity_update*self.div_loss
			self.reconstruction_loss = tf.reduce_mean(loss_list, axis=0)
			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			self.loss_summary = tf.summary.scalar('loss', self.total_loss)
			self.kl_loss_summary = tf.summary.scalar('kl-loss', self.kl_loss)
			self.div_loss_summary = tf.summary.scalar('div_loss', self.div_loss)
			self.reconstruction_loss_summary = tf.summary.scalar('reconstruction_loss', self.reconstruction_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			tf.summary.scalar('beta', self.beta)

			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
			self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
																   graph=tf.Session().graph)

			self.summary = tf.summary.merge_all()

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
			self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		n_other_agents = np.zeros([self.args.batch_size])
		for i,other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: batch_x,
				self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
				self.input_grid_placeholder: kwargs["batch_grid"],
				self.step: kwargs['step'],
				self.beta: kwargs['beta'],
				self.seq_length: n_other_agents,
				self.diversity_placeholder: kwargs['batch_div'],
				self.output_placeholder: batch_y,
				self.cell_state: self.cell_state_current,
				self.hidden_state: self.hidden_state_current,
				self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
				self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
				self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
				self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
				self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
				self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
				}
	
	def feed_pred_dic(self, **kwargs):
		step = kwargs["step"]
		n_other_agents = np.zeros([self.args.batch_size])

		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"], axis=1),
				self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"], axis=1),
				self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"], axis=1),
				self.step: 0,
				self.beta: 0,
				self.seq_length: n_other_agents,
				self.cell_state: self.test_cell_state_current,
				self.hidden_state: self.test_hidden_state_current,
				self.cell_state_lstm_grid: self.test_cell_state_current_lstm_grid,
				self.hidden_state_lstm_grid: self.test_hidden_state_current_lstm_grid,
				self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
				self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
				self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
				self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
				}

	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		n_other_agents = np.zeros([self.args.batch_size])
		#print(kwargs.keys())
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
				self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
				self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
				self.step: 0,
				self.beta: 0,
				self.seq_length: n_other_agents,
				self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
				self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
				self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
				self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
				self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
				self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
				self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
				self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
				}

	def feed_val_dic(self, n_other_agent = None, **kwargs):
		n_other_agents = np.zeros([self.args.batch_size])
		if n_other_agent is None:
			for i, other in enumerate(kwargs['other_agents_pos']):
				n_other_agents[i] = len(other)
		return {self.input_state_placeholder: kwargs["batch_vel"],
				self.input_grid_placeholder: kwargs["batch_grid"],
				self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
				self.diversity_placeholder: kwargs['batch_div'],
				self.output_placeholder: kwargs['batch_y'],
				self.seq_length: n_other_agents,
				self.step: 0,
				self.beta: 0,
				self.cell_state: np.random.normal(self.test_cell_state_current.copy(), kwargs["state_noise"]),
				self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), kwargs["state_noise"]),
				self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
				self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
				self.cell_state_lstm_ped: np.random.normal(self.test_cell_state_current_lstm_ped.copy(), kwargs["ped_noise"]),
				self.hidden_state_lstm_ped: np.random.normal(self.test_hidden_state_current_lstm_ped.copy(), kwargs["ped_noise"]),
				self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
				self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
				}

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])


	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def update_test_hidden_state(self):
		self.test_cell_state_current = self.cell_state_current.copy()
		self.test_hidden_state_current = self.hidden_state_current.copy()
		self.test_cell_state_current_lstm_grid = self.cell_state_current_lstm_grid.copy()
		self.test_hidden_state_current_lstm_grid = self.hidden_state_current_lstm_grid.copy()
		self.test_cell_state_current_lstm_ped = self.cell_state_current_lstm_ped.copy()
		self.test_hidden_state_current_lstm_ped = self.hidden_state_current_lstm_ped.copy()
		self.test_cell_state_current_lstm_concat = self.cell_state_current_lstm_concat.copy()
		self.test_hidden_state_current_lstm_concat = self.hidden_state_current_lstm_concat.copy()

	def validation_step(self, sess, feed_dict_validation,update = True):
		batch_loss, _model_prediction,_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, summary = sess.run([self.reconstruction_loss,
														self.prediction,
																			self.current_state,
																			self.current_state_lstm_grid,
																			self.current_state_lstm_ped,
																			self.current_state_lstm_concat,
																			self.loss_summary],
																		   feed_dict=feed_dict_validation)

		if update:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, summary, _model_prediction

	def train_step(self, sess, feed_dict_train, step=0):

		# Diversity Training
		if step > 2000 and self.args.diversity_update:
			
			dict = {"batch_vel": feed_dict_train[self.input_state_placeholder],
					"batch_grid": feed_dict_train[self.input_grid_placeholder],
					"batch_ped_grid": feed_dict_train[self.input_ped_grid_placeholder],
					"batch_y": feed_dict_train[self.output_placeholder],
					"batch_div": feed_dict_train[self.diversity_placeholder],
					"step": step,
					"beta": feed_dict_train[self.beta],
					"state_noise": 0.5,
					"grid_noise": 0.0,
					"ped_noise": 0.5
					}
			self.update_test_hidden_state()
			feed_test_dic = self.feed_val_dic(n_other_agent=feed_dict_train[self.seq_length],**dict)
			loss, _, y_model_pred = self.validation_step(sess, feed_test_dic, False)
			batch_y_varible = np.zeros((self.args.batch_size, self.args.truncated_backprop_length,
										self.args.prediction_horizon * self.args.output_dim))
			# Grount truth
			batch_y = feed_dict_train[self.output_placeholder]

			diversity_distance = 0
			diff_traj_index = 0
			# Find most different trajectory index
			for batch_id in range(self.args.batch_size):
				for t_id in range(self.args.truncated_backprop_length):
					for mix_id in range(self.args.n_mixtures):
						for pred_id in range(self.args.prediction_horizon):
							new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures+ mix_id]
							new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + mix_id]
							batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
								[new_batch_x, new_batch_y])
						if diversity_distance < np.linalg.norm(batch_y_varible[batch_id, t_id]-batch_y[batch_id, t_id]):
							diversity_distance = np.linalg.norm(batch_y_varible[batch_id, t_id] - batch_y[batch_id, t_id])
							diff_traj_index = mix_id
					# Saving most different traj
					for pred_id in range(self.args.prediction_horizon):
						new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + diff_traj_index]
						new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + diff_traj_index]
						batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
							[new_batch_x, new_batch_y])

			# Assemble feed dict for training
			feed_dict_train[self.diversity_placeholder] = batch_y_varible
		# kl_loss, div_loss
		_, batch_loss, kl_loss, recons_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
																										  self.total_loss,
																										  self.kl_loss,
																										  self.reconstruction_loss,
																										  self.current_state,
																										  self.current_state_lstm_grid,
																										  self.current_state_lstm_ped,
																										  self.current_state_lstm_concat,
																										  self.prediction,
																										  self.summary,
																										  self.learning_rate,
																										  self.beta,
																										  self.deconv1,
																										  self.div_loss,
																										  self.autoencoder_loss],
																										 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		# Train autoencoder if loss larger than threshold
		""""""
		if np.mean(autoencoder_loss) > 0.05:
			self.run_autoencoder(sess, feed_dict_train)
			self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

		out = {"batch_loss": recons_loss,
			   "model_predictions": _model_prediction,
			   "summary": _summary_str,
			   "learning_rate": lr,
			   "autoencoder_loss": autoencoder_loss,
			   "log loss": batch_loss,
			   "kl_loss": kl_loss,
			   "output_decoder": output_decoder
			   }

		return out


	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
					 feed_dict=feed_dict_train)


	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, _likelihood = sess.run([self.current_state,
													  self.current_state_lstm_grid,
													  self.current_state_lstm_ped,
													  self.current_state_lstm_concat,
													  self.prediction,
													  self.deconv1,
												  	  self.likelihood],
													 feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, _likelihood


	def warmstart_model(self, args, sess):
		# Restore whole model
		try:
			print('Loading session from "{}"'.format(args.model_path))
			ckpt = tf.train.get_checkpoint_state(args.model_path)
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.full_saver.restore(sess, ckpt.model_checkpoint_path)
		except:
			print(Fore.RED + "Model not found..."+Style.RESET_ALL)
			sys.exit()


	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print(Fore.RED +"Error. Encoder not found...")
			sys.exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True


	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)


	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures = 3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)],
								   1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss


	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result


	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y, pi):

		# gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		# gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		#loss = tf.multiply(z + tf.log(sxsy), pi)
		loss = z + tf.log(sxsy)

		return tf.reduce_mean(loss)


	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):
		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #


	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
													 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
													 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
													 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
															   name="fc_grid_weights",
															   trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
															trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
																name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
										   shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
																 [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
																 padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
																 [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
																 padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
																	  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
												 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
												 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
												 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
												 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors


	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector


	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
							  initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
																						 uniform=False, seed=None,
																						 dtype=tf.float32),
							  regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
							  initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
																						 uniform=False, seed=None,
																						 dtype=tf.float32),
							  trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
							name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations


	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
										padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation


	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		#print(input.size, weights.size)
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

	def frange_cycle_linear(self, n_iter, start=0.0, stop=1.0,  n_cycle=4, ratio=0.5):
		L = np.ones(n_iter) * stop
		period = n_iter/n_cycle
		step = (stop-start)/(period*ratio) # linear schedule

		for c in range(n_cycle):
			v, i = start, 0
			while v <= stop and (int(i+c*period) < n_iter):
				L[int(i+c*period)] = v
				v += step
				i += 1
		return L

src/models/VGDNN_diversity.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = self.process_grid(self.occ_grid_series)
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      # CNN weights for static occupancy grid processing
      train_grid_encoder_conv = False
      train_grid_encoder_fc = self.args.end_to_end
      use_summary_convnet = False
      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args,args.rnn_state_size_lstm_concat,self.output_vec_dim, args.prior_size, args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step],muy[prediction_step],sigmax[prediction_step],sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_sum(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0) - self.div_loss
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x[:,:,2:],
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_target):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.output_placeholder: batch_target,
            self.cell_state: self.test_cell_state_current,
            self.hidden_state: self.test_hidden_state_current,
            self.cell_state_lstm_grid: self.test_cell_state_current_lstm_grid,
            self.hidden_state_lstm_grid: self.test_hidden_state_current_lstm_grid,
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta, output_decoder, div_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          #self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta,
                                                          self.deconv1,
                                                          self.div_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, 0, _model_prediction, _summary_str, lr, beta, output_decoder, div_loss

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, loss = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                  self.total_loss],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, loss

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    normx = tf.subtract(x, dec_mux)
    normy = tf.subtract(y, dec_muy)

    x = tf.square(tf.div(normx, tf.square(dec_sigmax)))
    y = tf.square(tf.div(normy, tf.square(dec_sigmay)))

    sxsy = tf.multiply(dec_sigmax, dec_sigmay)

    xy = 2 * normx * normy / sxsy

    likelihood_loss = tf.math.log(tf.maximum(2 * np.pi * sxsy, 1e-6)) + x +y - xy
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    normx = tf.subtract(x, dec_mux)
    normy = tf.subtract(y, dec_muy)

    x = tf.square(tf.div(normx, tf.maximum(tf.square(dec_sigmax),1e-6)))
    y = tf.square(tf.div(normy, tf.maximum(tf.square(dec_sigmay),1e-6)))

    sxsy = tf.maximum(tf.multiply(dec_sigmax, dec_sigmay),1e-6)

    xy = 2 * normx * normy / sxsy

    likelihood_loss = -tf.math.log(tf.maximum(dec_pi, 1e-6)) +tf.math.log(tf.maximum(2 * np.pi*sxsy, 1e-6)) +x +y - xy
    return tf.reduce_sum(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=False)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=False)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=False)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=False)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=False)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=False)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=False)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=False)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=False, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=False, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=False, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_kl_distance.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from attention import *
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.external.attention import *
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.rnn_concat_size = args.rnn_state_size + args.rnn_state_size_lstm_grid + args.rnn_state_size_lstm_ped
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Multi-head attention implementation

        W_wq_state = self.get_weight_variable(name='w_wq_state',
                                                          shape=[self.rnn_state_size, self.rnn_state_size],
                                                          trainable=True)
        b_wq_state = self.get_bias_variable(name='b_wq_state', shape=[1, self.rnn_state_size], trainable=True)

        W_wk_state = self.get_weight_variable(name='w_wk_state',
                                                          shape=[self.rnn_state_size, self.rnn_state_size],
                                                          trainable=True)
        b_wk_state = self.get_bias_variable(name='b_wk_state', shape=[1, self.rnn_state_size], trainable=True)

        W_wv_state = self.get_weight_variable(name='w_wv_state',
                                                          shape=[self.rnn_state_size, self.rnn_state_size],
                                                          trainable=True)
        b_wv_state = self.get_bias_variable(name='b_wv_state', shape=[1, self.rnn_state_size], trainable=True)

        wq_state = tf.unstack(tf.stack(([self.fc_layer(output,W_wq_state,b_wq_state) for output in self.cell_outputs_series_state])))
        wk_state = tf.unstack(tf.stack(([self.fc_layer(output,W_wk_state,b_wk_state) for output in self.cell_outputs_series_state])))
        wv_state = tf.unstack(tf.stack(([self.fc_layer(output,W_wv_state,b_wv_state) for output in self.cell_outputs_series_state])))

        state_scaled_attention = []
        for state, key, value in zip(wq_state,wk_state,wv_state):
          state_scaled_output, attention_weights = scaled_dot_product_attention(state, key, value,mask=None)
          state_scaled_attention.append(state_scaled_output)

        W_wq_ped = self.get_weight_variable(name='w_wq_ped',
                                              shape=[self.rnn_state_size_lstm_ped, self.rnn_state_size_lstm_ped],
                                              trainable=True)
        b_wq_ped = self.get_bias_variable(name='b_wq_ped', shape=[1, self.rnn_state_size_lstm_ped], trainable=True)

        W_wk_ped = self.get_weight_variable(name='w_wk_ped',
                                              shape=[self.rnn_state_size_lstm_ped, self.rnn_state_size_lstm_ped],
                                              trainable=True)
        b_wk_ped = self.get_bias_variable(name='b_wk_ped', shape=[1, self.rnn_state_size_lstm_ped], trainable=True)

        W_wv_ped = self.get_weight_variable(name='w_wv_ped',
                                              shape=[self.rnn_state_size_lstm_ped, self.rnn_state_size_lstm_ped],
                                              trainable=True)
        b_wv_ped = self.get_bias_variable(name='b_wv_ped', shape=[1, self.rnn_state_size_lstm_ped], trainable=True)

        wq_ped = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wq_ped, b_wq_ped) for output in self.cell_outputs_series_lstm_ped])))
        wk_ped = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wk_ped, b_wk_ped) for output in self.cell_outputs_series_lstm_ped])))
        wv_ped = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wv_ped, b_wv_ped) for output in self.cell_outputs_series_lstm_ped])))

        ped_scaled_attention = []
        for state, key, value in zip(wq_ped,wk_ped,wv_ped):
          ped_scaled_output, attention_weights = scaled_dot_product_attention(state, key, value,mask=None)
          ped_scaled_attention.append(ped_scaled_output)

        W_wq_grid = self.get_weight_variable(name='w_wq_grid',
                                            shape=[self.rnn_state_size_lstm_grid, self.rnn_state_size_lstm_grid],
                                            trainable=True)
        b_wq_grid = self.get_bias_variable(name='b_wq_grid', shape=[1, self.rnn_state_size_lstm_grid], trainable=True)

        W_wk_grid = self.get_weight_variable(name='w_wk_grid',
                                            shape=[self.rnn_state_size_lstm_grid, self.rnn_state_size_lstm_grid],
                                            trainable=True)
        b_wk_grid = self.get_bias_variable(name='b_wk_grid', shape=[1, self.rnn_state_size_lstm_grid], trainable=True)

        W_wv_grid = self.get_weight_variable(name='w_wv_grid',
                                            shape=[self.rnn_state_size_lstm_grid, self.rnn_state_size_lstm_grid],
                                            trainable=True)
        b_wv_grid = self.get_bias_variable(name='b_wv_grid', shape=[1, self.rnn_state_size_lstm_grid], trainable=True)

        wq_grid = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wq_grid, b_wq_grid) for output in self.cell_outputs_series_lstm_grid])))
        wk_grid = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wk_grid, b_wk_grid) for output in self.cell_outputs_series_lstm_grid])))
        wv_grid = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wv_grid, b_wv_grid) for output in self.cell_outputs_series_lstm_grid])))

        grid_scaled_attention = []
        for state, key, value in zip(wq_grid,wk_grid,wv_grid):
          grid_scaled_output, attention_weights = scaled_dot_product_attention(state, key, value,mask=None)
          grid_scaled_attention.append(grid_scaled_output)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(state_scaled_attention, grid_scaled_attention, ped_scaled_attention):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

        W_linear_concat = self.get_weight_variable(name='w_linear_concat',
                                            shape=[self.rnn_concat_size, self.rnn_concat_size],
                                            trainable=True)
        b_linear_concat = self.get_bias_variable(name='b_linear_concat', shape=[1, self.rnn_concat_size], trainable=True)

        output_linear_concat = tf.unstack(
          tf.stack(([self.fc_layer(output, W_linear_concat, b_linear_concat) for output in concat_lstm])))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args,args.rnn_state_size_lstm_concat,self.output_vec_dim, args.prior_size, args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, output_linear_concat, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step],muy[prediction_step],sigmax[prediction_step],sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_target):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.output_placeholder: batch_target,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss, kl_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta,
                                                          self.deconv1,
                                                          self.div_loss,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, kl_loss,_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, loss, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                  self.total_loss,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, loss, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_multihead_attention.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class VGDNNModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = args.grid_width
    self.grid_height = args.grid_height
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim], name='input_state')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)
      
        ###################### Assemble Neural Net #####################

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args,args.rnn_state_size_lstm_concat,self.output_vec_dim, args.prior_size, args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.n_mixtures > 1:
            mux, muy, sigmax, sigmay, pi = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(mux)
            dec_muy_series.append(muy)
            dec_sigmax_series.append(tf.exp(sigmax))
            dec_sigmay_series.append(tf.exp(sigmay))
            dec_pi_series.append(tf.nn.softmax(pi))
          else:
            mux, muy, sigmax, sigmay = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(mux)
            dec_muy_series.append(muy)
            dec_sigmax_series.append(tf.exp(sigmax))
            dec_sigmay_series.append(tf.exp(sigmay))

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=500,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 5000) / 1000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.prediction = []
      pred= []

      if args.n_mixtures > 1:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.split(dec_sigmax, self.prediction_horizon, axis=1)
          sigmay = tf.split(dec_sigmay, self.prediction_horizon, axis=1)
          pi = tf.split(dec_pi, self.prediction_horizon, axis=1)

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)

          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.split(dec_sigmax, self.prediction_horizon, axis=1)
          sigmay = tf.split(dec_sigmay, self.prediction_horizon, axis=1)

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)

          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.total_loss = tf.reduce_mean(loss_list, axis=0)
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()

  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x[:,:,2:],
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid):
    return {self.input_state_placeholder: batch_x,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss, _current_state, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta = sess.run([self.update,
                                                          self.total_loss,
                                                          #self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, _model_prediction, _summary_str, lr, beta

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction = sess.run([self.current_state,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction],
                                                          feed_dict=feed_dict_train)
    if update_state:
      self.cell_state_current, self.hidden_state_current = _current_state
      self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    
    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)
      conv1 = self.conv_layer(input=grid_batch, weights=self.conv1_weights, biases=self.conv1_biases, conv_stride_length=self.conv1_stride_length, name="conv1_grid")
      conv2 = self.conv_layer(input=conv1, weights=self.conv2_weights, biases=self.conv2_biases, conv_stride_length=self.conv2_stride_length, name="conv2_grid")
      conv3 = self.conv_layer(input=conv2, weights=self.conv3_weights, biases=self.conv3_biases, conv_stride_length=self.conv3_stride_length, name="conv3_grid")
      conv_grid_size = 8
      conv5_flat = tf.reshape(conv3, [-1, conv_grid_size*conv_grid_size*self.conv3_number_filters])
      fc_final = self.fc_layer(input=conv5_flat, weights=self.fc_grid_weights, biases=self.fc_grid_biases, activation = tf.nn.sigmoid, name="fc_grid")
      if idx == 0:
        tf.summary.histogram("fc_activations", fc_final)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)
    
    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_noGrid.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from storn_cell import VariationalRNNCell as vrnn_cell
  from attention import *
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.external.attention import *
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=train_encoder)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        W_wq_ped = self.get_weight_variable(name='w_wq_ped',
                                            shape=[self.rnn_state_size_lstm_ped, self.rnn_state_size_lstm_ped],
                                            trainable=True)
        b_wq_ped = self.get_bias_variable(name='b_wq_ped', shape=[1, self.rnn_state_size_lstm_ped], trainable=True)

        W_wk_ped = self.get_weight_variable(name='w_wk_ped',
                                            shape=[self.rnn_state_size_lstm_ped, self.rnn_state_size_lstm_ped],
                                            trainable=True)
        b_wk_ped = self.get_bias_variable(name='b_wk_ped', shape=[1, self.rnn_state_size_lstm_ped], trainable=True)

        W_wv_ped = self.get_weight_variable(name='w_wv_ped',
                                            shape=[self.rnn_state_size_lstm_ped, self.rnn_state_size_lstm_ped],
                                            trainable=True)
        b_wv_ped = self.get_bias_variable(name='b_wv_ped', shape=[1, self.rnn_state_size_lstm_ped], trainable=True)

        wq_ped = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wq_ped, b_wq_ped) for output in ped_grid_feature_series])))
        wk_ped = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wk_ped, b_wk_ped) for output in ped_grid_feature_series])))
        wv_ped = tf.unstack(
          tf.stack(([self.fc_layer(output, W_wv_ped, b_wv_ped) for output in ped_grid_feature_series])))

        ped_scaled_attention = []
        for state, key, value in zip(wq_ped, wk_ped, wv_ped):
          ped_scaled_output, attention_weights = scaled_dot_product_attention(tf.expand_dims(state, axis=1),
                                                                              tf.expand_dims(key, axis=1),
                                                                              tf.expand_dims(value, axis=1), mask=None)
          ped_scaled_attention.append(tf.squeeze(ped_scaled_output, axis=1))

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_scaled_attention, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        try:
          cell_lstm_concat = vrnn_cell(args,args.x_dim,self.rnn_state_size_lstm_concat, args.prior_size, args.latent_space_size)
        except:
          cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.rnn_state_size_lstm_concat, args.prior_size,
                                       args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step],muy[prediction_step],sigmax[prediction_step],sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0)
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_target,noise=1.0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.output_placeholder: batch_target,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),noise),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),noise),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),noise),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),noise),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: np.random.normal(self.test_cell_state_current_lstm_ped.copy(),0),
            self.hidden_state_lstm_ped: np.random.normal(self.test_hidden_state_current_lstm_ped.copy(),0),
            self.cell_state_lstm_concat: np.random.normal(self.test_cell_state_current_lstm_concat.copy(),0),
            self.hidden_state_lstm_concat: np.random.normal(self.test_hidden_state_current_lstm_concat.copy(),0),
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss, kl_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta,
                                                          self.deconv1,
                                                          self.div_loss,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, kl_loss,_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, loss, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                  self.total_loss,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, loss, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")

      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_ped_attention.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys
if sys.version_info[0] < 3:
  sys.path.append('../src/external')
  from vrnn_cell import VariationalRNNCell as vrnn_cell
  from tf_utils import *
else:
  from src.external.vrnn_cell import VariationalRNNCell as vrnn_cell
  from src.models.tf_utils import *
import numpy as np
import os

class NetworkModel():
  
  def __init__(self, args,is_training=True,batch_size = None):
    self.log_dir = args.log_dir
    self.args = args
    # Input- / Output dimensions
    self.input_dim = args.input_dim  # [x, y, vx, vy]
    self.input_state_dim = args.input_state_dim  # [vx, vy]
    self.output_dim = args.output_dim
    self.output_pred_state_dim = self.args.output_pred_state_dim
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prev_horizon = args.prev_horizon
    self.n_mixtures = args.n_mixtures
    self.prediction_horizon = args.prediction_horizon
    self.output_placeholder_dim = self.prediction_horizon * self.output_dim
    self.output_vec_dim = self.prediction_horizon *self.output_pred_state_dim * self.n_mixtures
    self.pedestrian_vector_dim = args.pedestrian_vector_dim
    self.learning_rate_init = args.learning_rate_init
    self.beta_rate_init = args.beta_rate_init
    # Network attributes
    self.rnn_state_size = args.rnn_state_size
    self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
    self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
    try:
      self.sigma_bias = args.sigma_bias
    except:
      self.sigma_bias = 0.0
    # Training parameters
    self.lambda_ = args.regularization_weight
    self.batch_size = None
    self.is_training = is_training
    self.grads_clip = args.grads_clip
    self.regularization_weight = args.regularization_weight
    self.grid_width = int(args.submap_width / args.submap_resolution)
    self.grid_height = int(args.submap_height / args.submap_resolution)
    # Specify placeholders
    self.input_state_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.input_state_dim*(self.prev_horizon+1)], name='input_state')
    self.input_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.grid_width, self.grid_height], name='input_grid')
    self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.pedestrian_vector_dim], name='ped_grid')
    self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length, self.output_placeholder_dim], name='output')
    self.positions_output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
                                                                      self.output_placeholder_dim], name='output')
    self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
    self.step = tf.placeholder(dtype=tf.float32,
                                                  shape=[], name='global_step')
    # Initialize hidden states with zeros
    self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

    # Cells used for testing
    # Initialize hidden states with zeros
    self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
    self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
    self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
    self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
    
    # Pedestrian state LSTM
    self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size], name='cell_state')  # internal state of the cell (before output gate)
    self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size], name='hidden_state')  # output of the cell (after output gate)
    self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)
    
    # Static occupancy grid LSTM
    self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid], name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid], name='hidden_state_lstm_grid')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid, self.hidden_state_lstm_grid)
    
    # Pedestrian LSTM
    self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped], name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped], name='hidden_state_lstm_ped')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)
    
    # Concatenation LSTM (state, grid, pedestrians)
    self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_concat], name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
    self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat], name='hidden_state_lstm_concat')  # output of the cell (after output gate)
    self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat, self.hidden_state_lstm_concat)

    inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
    outputs_series = tf.unstack(self.output_placeholder, axis=1)
    positions_series = tf.unstack(self.positions_output_placeholder, axis=1)
    # Print network info
    print("Input list length: {}".format(len(inputs_series)))
    print("Output list length: {}".format(len(outputs_series)))
    print("Single input shape: {}".format(inputs_series[0].get_shape()))
    print("Single output shape: {}".format(outputs_series[0].get_shape()))

    ###################### Assemble CNN #####################
    # Process occupancy grid, apply convolutional filters
    self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
    conv_grid_feature_series = []
    for output in self.process_grid(self.occ_grid_series):
      conv_grid_feature_series.append(tf.stop_gradient(output))
    
    # Set up model with variables
    with tf.variable_scope("model") as scope:

      train_encoder = self.args.end_to_end

      with tf.variable_scope('encoder_model') as scope:
        # Pedestrian grid embedding layer
        ped_grid_out_dim = 128
        self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
                                                          shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
                                                          trainable=True)
        self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)


        # LSTM cells to process static occupancy grid
        with tf.variable_scope('lstm_grid') as scope:
          self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
                                                   trainable=True)
          self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
                                                                                                       conv_grid_feature_series,
                                                                                                       dtype=tf.float32,
                                                                                                       initial_state=self.init_state_tuple_lstm_grid)

        # Process pedestrian grid
        ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
        # Embedding layer of pedestrian grid
        ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

        # LSTM cells to process pedestrian grid
        with tf.variable_scope('lstm_ped') as scope:
          self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(self.cell_ped, ped_grid_feature_series, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_ped)

        # LSTM cells to process pedestrian state
        with tf.variable_scope('lstm_state') as scope:
          self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size,name='basic_lstm_cell',trainable = True)
          self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series, dtype=tf.float32, initial_state=self.init_state_tuple)

        # Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
        concat_lstm = []
        for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid, self.cell_outputs_series_lstm_ped):
          concat_lstm.append(tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

      # Concatenated LSTM layer
      with tf.variable_scope('generation_model') as scope:

        #cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
        cell_lstm_concat = vrnn_cell(args,args.rnn_state_size_lstm_concat,self.output_vec_dim, args.prior_size, args.latent_space_size)
        self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

        # Output reshape
        enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
        prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]),axis=0)
        prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

        # Decoder layers
        dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in dec_series])))
        dec_output_layer = []
        dec_mux_series = []
        dec_muy_series = []
        dec_sigmax_series = []
        dec_sigmay_series = []
        dec_pi_series =[]
        for t_idx in range(len(dec_fc_layer)):
          dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, t_idx))
        for t_idx in range(len(dec_output_layer)):
          if args.output_pred_state_dim > 4:
            m_x, m_y, s_x, s_y, likelihood = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)
            dec_pi_series.append(likelihood)
          else:
            m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], self.args.output_pred_state_dim, axis=1)
            dec_mux_series.append(m_x)
            dec_muy_series.append(m_y)
            dec_sigmax_series.append(s_x)
            dec_sigmay_series.append(s_y)

        # Exponential / stepwise decrease of learning rate
        global_step = tf.Variable(0, trainable=False)
        self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
                                                        decay_rate=0.9, staircase=False)

        self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

      prediction_loss_list = []
      loss_list = []
      position_loss_list = []
      kl_loss_list = []
      self.div_loss = []
      self.prediction = []
      pred= []
      div_loss_over_truncated_back_prop = []
      if args.output_pred_state_dim > 4:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax,dec_sigmay,prior_mu, prior_sigma, dec_pi, out in \
            zip(enc_mu_series,enc_sigma_series,dec_mux_series,dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series,prior_sigma_series, dec_pi_series,outputs_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          # Apply transformations
          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
          sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
          sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
          pi = tf.nn.softmax(tf.split(dec_pi, self.prediction_horizon, axis=1))

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          div_loss_over_horizon = []
          lossfunc = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))

            if prediction_step==0:
              pred = tf.concat([mux[prediction_step],muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)
            else:
              pred = tf.concat([pred,mux[prediction_step]],1)
              pred = tf.concat([pred,muy[prediction_step]],1)
              pred = tf.concat([pred, sigmax[prediction_step]],1)
              pred = tf.concat([pred, sigmay[prediction_step]],1)
              pred = tf.concat([pred, pi[prediction_step]], 1)

            lossfunc = self.get_mixture_lossfunc(mux[prediction_step],muy[prediction_step],
                                         sigmax[prediction_step],sigmay[prediction_step],
                                         pi[prediction_step],
                                         x_data,y_data)

            diversity_loss = self.get_diversity_loss(mux[prediction_step],muy[prediction_step],sigmax[prediction_step],sigmay[prediction_step])

            div_loss_over_horizon.append(diversity_loss)
            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      else:
        for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, out, pos in \
            zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
                dec_sigmax_series, dec_sigmay_series,
                prior_mu_series, prior_sigma_series,
                outputs_series, positions_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

          mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
          muy = tf.split(dec_muy, self.prediction_horizon, axis=1)

          sigmax = tf.minimum(tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1)),5.0)
          sigmay = tf.minimum(tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1)),5.0)

          pred = []
          prediction_loss_list = []
          kl_mode_loss_list = []
          pos_loss_list = []
          lossfunc = 0
          x_pos = 0
          y_pos = 0
          for prediction_step in range(self.prediction_horizon):
            # Compute Euclidean error for each prediction step within each tbp step
            idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
            idx_x = prediction_step * self.output_dim
            idx_y = idx_x + 1

            # ground truth data
            x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
            y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
            x_pos_data = tf.reshape(pos[:, idx_x], shape=(-1, 1))
            y_pos_data = tf.reshape(pos[:, idx_y], shape=(-1, 1))

            if prediction_step == 0:
              pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)
            else:
              pred = tf.concat([pred, mux[prediction_step]], 1)
              pred = tf.concat([pred, muy[prediction_step]], 1)
              pred = tf.concat([pred, sigmax[prediction_step]], 1)
              pred = tf.concat([pred, sigmay[prediction_step]], 1)

            lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
                                                 sigmax[prediction_step], sigmay[prediction_step],
                                                 x_data, y_data)

            x_pos += mux[prediction_step][:,0] * self.args.dt
            y_pos += muy[prediction_step][:,0] * self.args.dt

            pos_loss = tf.sqrt(tf.add(tf.square(tf.subtract(x_pos, tf.squeeze(x_pos_data,axis=1))),  # error in x
                                    tf.square(tf.subtract(y_pos, tf.squeeze(y_pos_data,axis=1)))))  # error in y

            pos_loss_list.append(pos_loss)

            prediction_loss_list.append(lossfunc)

          self.prediction.append(pred)
          position_loss_list.append(tf.reduce_sum(pos_loss_list))
          kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
          kl_loss = tf.squeeze(tf.multiply(kl_loss, self.beta))
          kl_loss_list.append(kl_loss)
          loss_list.append(tf.reduce_mean(prediction_loss_list))

      # Reduce mean in all dimensions
      #self.kl_loss = tf.reduce_mean(tf.reduce_mean(kl_mode_loss_list, axis=0), axis=0)
      #self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.position_loss = tf.reduce_mean(position_loss_list)
      self.total_loss = tf.reduce_mean(loss_list, axis=0) + tf.reduce_mean(kl_loss_list, axis=0) + 10*self.position_loss
      self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)
      self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
      
      # Optimizer specification
      #self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
      self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
      # Compute gradients
      train_vars = tf.trainable_variables()
      self.gradients = tf.gradients(self.total_loss, train_vars)

      # Clip gradients
      self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)
    
      # Apply gradients to update model
      self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
      #self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
      # Tensorboard summary
      
      ## Loss
      tf.summary.scalar('loss', self.total_loss)
      tf.summary.scalar('learning_rate', self.learning_rate)
      self.summary = tf.summary.merge_all()
      self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)

      # Save / restore all variables 
      self.full_saver = tf.train.Saver()
      
    # Only save / restore the parameters responsible for occupancy grid processing
    self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
    self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


  def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y,batch_pos):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: step,
            self.output_placeholder: batch_y,
            self.positions_output_placeholder: batch_pos,
            self.cell_state: self.cell_state_current,
            self.hidden_state: self.hidden_state_current,
            self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid,0),
            self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid,0),
            self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
            }

  def feed_test_dic(self, batch_x,batch_grid, batch_ped_grid,batch_target,batch_pos,noise=0):
    return {self.input_state_placeholder: batch_x,
            self.input_grid_placeholder: batch_grid,
            self.input_ped_grid_placeholder: batch_ped_grid,
            self.step: 0,
            self.output_placeholder: batch_target,
            self.positions_output_placeholder: batch_pos,
            self.cell_state: np.random.normal(self.test_cell_state_current.copy(),0),
            self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(),0),
            self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(),0),
            self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(),0),
            #self.cell_state: self.cell_state_current,
            #self.hidden_state: self.hidden_state_current,
            #self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
            #self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
            self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
            self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
            self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
            self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
            }

  def reset_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def reset_test_cells(self, sequence_reset):
    # Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
    if np.any(sequence_reset):
      for sequence_idx in range(sequence_reset.shape[0]):
        if sequence_reset[sequence_idx] == 1:
          # Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
          self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
          self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
          self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
          self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
          self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

  def run(self,sess,feed_dict_train):
    _, batch_loss, kl_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat,\
    _model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
                                                          self.total_loss,
                                                          self.kl_loss,
                                                          self.current_state,
                                                          self.current_state_lstm_grid,
                                                          self.current_state_lstm_ped,
                                                          self.current_state_lstm_concat,
                                                          self.prediction,
                                                          self.summary,
                                                          self.learning_rate,
                                                          self.beta,
                                                          self.deconv1,
                                                          self.position_loss,
                                                          self.autoencoder_loss],
                                                          feed_dict=feed_dict_train)
    self.cell_state_current, self.hidden_state_current = _current_state
    self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
    self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
    self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

    return batch_loss, kl_loss,_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss

  def run_autoencoder(self,sess,feed_dict_train):
    _ = sess.run([self.autoencoder_update],
                                                          feed_dict=feed_dict_train)

  def predict(self,sess,feed_dict_train,update_state=True):
    _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
    _current_state_lstm_concat, \
    _model_prediction, loss, output_decoder = sess.run([self.current_state,
                                              self.current_state_lstm_grid,
                                              self.current_state_lstm_ped,
                                              self.current_state_lstm_concat,
                                              self.prediction,
                                  self.total_loss,
                                        self.deconv1],
                                              feed_dict=feed_dict_train)
    if update_state:
      self.test_cell_state_current, self.test_hidden_state_current = _current_state
      self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
      self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
      self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

    return _model_prediction, loss, output_decoder

  def warmstart_model(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.model_path))
    ckpt = tf.train.get_checkpoint_state(args.model_path)
    print('Restoring model {}'.format(ckpt.model_checkpoint_path))
    self.full_saver.restore(sess, ckpt.model_checkpoint_path)

  def warmstart_model_encoder(self, args, sess):
    # Restore whole model
    print('Loading session from "{}"'.format(args.pretrained_encoder_path))
    ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
    if ckpt == None:
      print("Error. Encoder not found...")
      exit()
    else:
      print('Restoring model {}'.format(ckpt.model_checkpoint_path))
      self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

    return True

  def warmstart_convnet(self, args, sess):
    # Restore convnet parameters from pretrained ones (autoencoder encoding part)
    print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
    ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
    if ckpt_conv == None:
      print("Error. Convnet not found...")
      exit()
    else:
      print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
      self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

  def get_diversity_loss(self,mux,muy,sigmax,sigmay):
    # It only works for n_mixtures =3
    div_loss=0
    for mix_id in range(self.args.n_mixtures):
      mu1 = tf.concat([tf.expand_dims(mux[:,mix_id],axis=1), tf.expand_dims(muy[:,mix_id],axis=1)],1)
      sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
      gaussian1 = tfd.MultivariateNormalDiag(
        loc=mu1,
        scale_diag=sigma1)
      if mix_id+1 < self.args.n_mixtures:
        mu2 = tf.concat([tf.expand_dims(mux[:, mix_id+1], axis=1), tf.expand_dims(muy[:, mix_id+1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id+1], axis=1), tf.expand_dims(sigmay[:, mix_id+1], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)
      else:
        mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
        sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
        gaussian2 = tfd.MultivariateNormalDiag(
          loc=mu2,
          scale_diag=sigma2)

      div_loss += tfd.kl_divergence(gaussian1,gaussian2)


      return div_loss

  def tf_2d_normal(self, x, y, mux, muy, sx, sy):
    '''
    Function that implements the PDF of a 2D normal distribution
    params:
    x : input x points
    y : input y points
    mux : mean of the distribution in x
    muy : mean of the distribution in y
    sx : std dev of the distribution in x
    sy : std dev of the distribution in y
    rho : Correlation factor of the distribution
    '''
    # eq 3 in the paper
    # and eq 24 & 25 in Graves (2013)
    # Calculate (x - mux) and (y-muy)
    rho = tf.constant(0.0) # hack because we dont want correlation now
    normx = tf.subtract(x, mux)
    normy = tf.subtract(y, muy)
    # Calculate sx*sy
    sxsy = tf.multiply(sx, sy)
    # Calculate the exponential factor
    z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(tf.multiply(rho, tf.multiply(normx, normy)),
                                                                                 sxsy)
    negRho = 1 - tf.square(rho)
    # Numerator
    result = tf.exp(tf.div(-z, 2 * negRho))
    # Normalization constant
    denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
    # Final PDF calculation
    result = tf.div(result, denom)
    return result

  def tf_kl_gaussgauss(self,mu_1, sigma_1, mu_2, sigma_2):
    with tf.variable_scope("kl_gaussgauss"):
      #return tf.reduce_sum(0.5 * (
      #    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
      #), 1)
      #return tf.reduce_sum(0.5 * (
      #    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
      #    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
      #    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
      #), 1)
      kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
      return tf.reduce_mean(kl)

  def get_log_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    try:
      if self.args.regulate_log_loss:
        likelihood_loss = -1.0*tf.minimum(tf.math.log(tf.maximum(gaussian, 1e-10)),4.0)
      else:
        likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    except:
      likelihood_loss = -1.0 * tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)

  def get_mixture_lossfunc(self,dec_mux,dec_muy, dec_sigmax,dec_sigmay,dec_pi,x,y):

    gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
    gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
    gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

    likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
    return tf.reduce_mean(likelihood_loss)#
      
  def process_grid(self, occ_grid_series):
    """
    Process occupancy grid with series of convolutional and fc layers.
    input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
    output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
    """
    # Convolutional layer 1
    conv1_kernel_size = 5
    conv1_number_filters = 64
    conv1_stride_length = 2

    # Convolutional layer 2
    conv2_kernel_size = 3
    conv2_number_filters = 32
    conv2_stride_length = 2

    # Convolutional layer 3
    conv3_kernel_size = 3
    conv3_number_filters = 8
    conv3_stride_length = 2

    # Flatten convolutional filter vector such that it can be fed to a FC layer
    conv_grid_size = 8
    fc_grid_hidden_dim = 64

    conv_feature_vectors = []
    print("Length occ grid series: {}".format(len(occ_grid_series)))
    for idx, grid_batch in enumerate(occ_grid_series):
      grid_batch = tf.expand_dims(input=grid_batch, axis=3)

      if idx == 0:
        with tf.variable_scope('auto_encoder') as scope:
          conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                           conv1_stride_length, trainable=True)
          conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                           conv2_stride_length, trainable=True)
          conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                           conv3_stride_length, padding='VALID', trainable=False)

          conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

          fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
                                                     name="fc_grid_weights",
                                                     trainable=True)
          fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
                                                  trainable=True)

          fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

          tf.summary.histogram("fc_activations", fc_final)
          fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
                                                      name="fc_grid_decode_biases",trainable=True)
          fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
          print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

          # Reshaping into known convolutional grid size
          conv3_dec = tf.reshape(fc_grid_dec,
                                      shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
          print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

          # From output of conv3 to input of conv3 (output of conv2)
          deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
                                                            [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
                                                            padding='VALID',trainable=True)
          print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

          # From output of conv2 to input of conv2 (output of conv1)
          deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
                                                            [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
                                                            padding='SAME',trainable=True)
          print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

          # From output of conv1 to input of conv1 (raw grid input)
          self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
                                                            conv1_stride_length, padding='SAME',trainable=True)
          print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

          self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

          self.autoencoder_loss = tf.reduce_mean(
            tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


      else:
        conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
                                         conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
        conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
                                         conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
        conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
                                         conv3_stride_length, padding='VALID', trainable=True, reuse=True,
                                         scope='auto_encoder/Conv_2')

        conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
        fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)
      
      # Flatten to obtain feature vector
      conv_features = tf.contrib.layers.flatten(fc_final)
      conv_feature_vectors.append(conv_features)

    return conv_feature_vectors
    
  def process_pedestrian_grid(self, ped_grid_series):
    """
    Process pedestrian grid with a FC layer
    """
    fc_feature_vector = []
    for idx, grid_batch in enumerate(ped_grid_series):
      grid_batch = tf.contrib.layers.flatten(grid_batch)
      #fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
      fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
      if idx == 0:
        tf.summary.histogram("fc_ped_activations", fc_feature_vector)
#       fc_feature_vector.append(grid_batch)
    return fc_feature_vector

  def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
    #                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var
  
  def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
    #var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
    #                      trainable=trainable)
    var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN', uniform=False, seed=None, dtype=tf.float32),
                          trainable=trainable)
    if summary: 
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get weight variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var

  def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
    """
    Get bias variable with specific initializer.
    """
    var = tf.Variable(tf.truncated_normal(shape, mean=0.1,stddev=0.1), dtype=tf.float32,name=name)
    if summary:
      tf.summary.histogram(name, var)
    return var
    
  def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
    """
    Convolutional layer including a ReLU activation but excluding pooling.
    """
    conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding, name=name)
    activations = tf.nn.relu(conv + biases)
    if summary: 
      tf.summary.histogram(name, activations)
    return activations

  def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
    #
      input_shape = input.get_shape()
      deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1], padding=padding)
      activation = tf.nn.relu(deconv + biases)
      return activation
  
  def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
    """
    Fully connected layer with given weights and biases. 
    Activation and summary can be activated with the arguments.
    """
    affine_result = tf.matmul(input, weights) + biases
    if activation is None:
      activations = affine_result
    else:
      activations = activation(affine_result)
    if summary: 
      tf.summary.histogram(name + "_activations", activations)
    return activations

src/models/VGDNN_pos.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os
from colorama import Fore, Style

class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    self.grid_width, self.grid_height], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim*self.args.n_other_agents], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.diversity_placeholder = tf.placeholder(dtype=tf.float32,
		                                            shape=[self.batch_size, self.truncated_backprop_length,
		                                                   self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')
		self.seq_length = tf.placeholder(tf.int32, [None])
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)
		outputs_diversity_series = tf.unstack(self.diversity_placeholder, axis=1)
		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[ped_grid_out_dim, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
					                                         trainable=True)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					ped_grid_feature_series , self.current_state_lstm_ped = tf.nn.dynamic_rnn(self.cell_ped,
					                                                                                           tf.stack(ped_grid_series,axis=1),
					                                                                                           dtype=tf.float32,
					                                                                                           sequence_length=self.seq_length,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

					# Embedding layer of pedestrian grid
					self.cell_outputs_series_lstm_ped = self.process_pedestrian_grid(tf.unstack(ped_grid_feature_series,axis=1))

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state,
				                                               self.cell_outputs_series_lstm_grid,
				                                               self.cell_outputs_series_lstm_ped):
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

					# Concatenated LSTM layer
					with tf.variable_scope('generation_model') as scope:

						# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
						cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
						                             args.latent_space_size)
						self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
							cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

						# Output reshape
						enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

						# Decoder layers
						dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
						dec_output_layer = []
						dec_mux_series = []
						dec_muy_series = []
						dec_sigmax_series = []
						dec_sigmay_series = []
						dec_pi_series = []
						# Reusing weights for each truncated back propagation step
						for t_idx in range(len(dec_fc_layer)):
							dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, 50))
							dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100))
						for t_idx in range(len(dec_output_layer)):
							m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1)
							dec_mux_series.append(m_x)
							dec_muy_series.append(m_y)
							dec_sigmax_series.append(s_x)
							dec_sigmay_series.append(s_y)
						dec_pi_series = [dense for dense in dec_pi_series]

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
				                                                decay_rate=0.9, staircase=False)

				self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

				prediction_loss_list = []
				loss_list = []
				kl_loss_list = []
				self.div_loss = []
				self.prediction = []
				self.likelihood = []
				pred = []
				div_loss_over_truncated_back_prop = []

				# Compute Loss Function
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out, out_div in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series, dec_pi_series,
						    outputs_series, outputs_diversity_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(dec_pi, axis=1)

					self.likelihood.append(pi)
					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
						x_diversity = tf.reshape(out_div[:, idx_x], shape=(-1, 1))
						y_diversity = tf.reshape(out_div[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                 sigmax[prediction_step], sigmay[prediction_step],
						                                 x_data, y_data, pi)

						diversity_loss = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                        sigmax[prediction_step], sigmay[prediction_step],
						                                       x_diversity, y_diversity, pi)

						div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss_list.append(kl_loss)
					div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			self.prediction = tf.stack(self.prediction, axis=1)
			self.likelihood = tf.stack(self.likelihood, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
			self.total_loss = tf.reduce_mean(loss_list, axis=0)+args.diversity_update*self.div_loss + tf.reduce_mean(kl_loss_list, axis=0)*self.beta
			self.reconstruction_loss = tf.reduce_mean(loss_list, axis=0)
			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			self.loss_summary = tf.summary.scalar('loss', self.total_loss)
			self.kl_loss_summary = tf.summary.scalar('kl-loss', self.kl_loss)
			self.div_loss_summary = tf.summary.scalar('div_loss', self.div_loss)
			self.reconstruction_loss_summary = tf.summary.scalar('reconstruction_loss', self.reconstruction_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)

			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
			self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
			                                                       graph=tf.Session().graph)

			self.summary = tf.summary.merge_all()

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
			self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		n_other_agents = np.zeros([self.args.batch_size])
		for i,other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.step: kwargs['step'],
		        self.seq_length: n_other_agents,
		        self.diversity_placeholder: kwargs['batch_div'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }


	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		n_other_agents = np.zeros([self.args.batch_size])
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
		        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
		        self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
		        self.step: 0,
		        self.seq_length: n_other_agents,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
		        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def feed_val_dic(self, **kwargs):
		n_other_agents = np.zeros([self.args.batch_size])
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: kwargs["batch_vel"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.diversity_placeholder: kwargs['batch_div'],
		        self.output_placeholder: kwargs['batch_y'],
		        self.seq_length: n_other_agents,
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), kwargs["state_noise"]),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), kwargs["state_noise"]),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
		        self.cell_state_lstm_ped: np.random.normal(self.test_cell_state_current_lstm_ped.copy(), kwargs["ped_noise"]),
		        self.hidden_state_lstm_ped: np.random.normal(self.test_hidden_state_current_lstm_ped.copy(), kwargs["ped_noise"]),
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])


	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def update_test_hidden_state(self):
		self.test_cell_state_current = self.cell_state_current.copy()
		self.test_hidden_state_current = self.hidden_state_current.copy()
		self.test_cell_state_current_lstm_grid = self.cell_state_current_lstm_grid.copy()
		self.test_hidden_state_current_lstm_grid = self.hidden_state_current_lstm_grid.copy()
		self.test_cell_state_current_lstm_ped = self.cell_state_current_lstm_ped.copy()
		self.test_hidden_state_current_lstm_ped = self.hidden_state_current_lstm_ped.copy()
		self.test_cell_state_current_lstm_concat = self.cell_state_current_lstm_concat.copy()
		self.test_hidden_state_current_lstm_concat = self.hidden_state_current_lstm_concat.copy()

	def validation_step(self, sess, feed_dict_validation,update = True):
		batch_loss, _model_prediction,_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, summary = sess.run([self.reconstruction_loss,
		                                                self.prediction,
		                                                                    self.current_state,
		                                                                    self.current_state_lstm_grid,
		                                                                    self.current_state_lstm_ped,
		                                                                    self.current_state_lstm_concat,
		                                                                    self.loss_summary],
		                                                                   feed_dict=feed_dict_validation)

		if update:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, summary, _model_prediction

	def train_step(self, sess, feed_dict_train, step=0):

		# Diversity Training
		if step > 2000 and self.args.diversity_update:
			dict = {"batch_vel": feed_dict_train[self.input_state_placeholder],
			        "batch_grid": feed_dict_train[self.input_grid_placeholder],
			        "batch_ped_grid": feed_dict_train[self.input_ped_grid_placeholder],
			        "batch_y": feed_dict_train[self.output_placeholder],
			        "batch_div": feed_dict_train[self.diversity_placeholder],
			        "step": step,
			        "state_noise": 0.5,
			        "grid_noise": 0.0,
			        "ped_noise": 0.5
			        }
			self.update_test_hidden_state()
			feed_test_dic = self.feed_val_dic(**dict)
			loss, _, y_model_pred = self.validation_step(sess, feed_test_dic, False)
			batch_y_varible = np.zeros((self.args.batch_size, self.args.truncated_backprop_length,
			                            self.args.prediction_horizon * self.args.output_dim))
			# Grount truth
			batch_y = feed_dict_train[self.output_placeholder]

			diversity_distance = 0
			diff_traj_index = 0
			# Find most different trajectory index
			for batch_id in range(self.args.batch_size):
				for t_id in range(self.args.truncated_backprop_length):
					for mix_id in range(self.args.n_mixtures):
						for pred_id in range(self.args.prediction_horizon):
							new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures+ mix_id]
							new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + mix_id]
							batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
								[new_batch_x, new_batch_y])
						if diversity_distance < np.linalg.norm(batch_y_varible[batch_id, t_id]-batch_y[batch_id, t_id]):
							diversity_distance = np.linalg.norm(batch_y_varible[batch_id, t_id] - batch_y[batch_id, t_id])
							diff_traj_index = mix_id
					# Saving most different traj
					for pred_id in range(self.args.prediction_horizon):
						new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + diff_traj_index]
						new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + diff_traj_index]
						batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
							[new_batch_x, new_batch_y])

			# Assemble feed dict for training
			feed_dict_train[self.diversity_placeholder] = batch_y_varible

		_, batch_loss, kl_loss, recons_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
		                                                                                                  self.total_loss,
		                                                                                                  self.kl_loss,
		                                                                                                  self.reconstruction_loss,
		                                                                                                  self.current_state,
		                                                                                                  self.current_state_lstm_grid,
		                                                                                                  self.current_state_lstm_ped,
		                                                                                                  self.current_state_lstm_concat,
		                                                                                                  self.prediction,
		                                                                                                  self.summary,
		                                                                                                  self.learning_rate,
		                                                                                                  self.beta,
		                                                                                                  self.deconv1,
		                                                                                                  self.div_loss,
		                                                                                                  self.autoencoder_loss],
		                                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		# Train autoencoder if loss larger than threshold
		""""""
		if np.mean(autoencoder_loss) > 0.05:
			self.run_autoencoder(sess, feed_dict_train)
			self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

		out = {"batch_loss": recons_loss,
		       "model_predictions": _model_prediction,
		       "summary": _summary_str,
		       "learning_rate": lr,
		       "autoencoder_loss": autoencoder_loss,
		       "log loss": batch_loss,
		       "kl_loss": kl_loss,
		       "output_decoder": output_decoder
		       }

		return out


	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)


	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, _likelihood = sess.run([self.current_state,
		                                              self.current_state_lstm_grid,
		                                              self.current_state_lstm_ped,
		                                              self.current_state_lstm_concat,
		                                              self.prediction,
		                                              self.deconv1,
                                                  self.likelihood],
		                                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, _likelihood


	def warmstart_model(self, args, sess):
		# Restore whole model
		try:
			print('Loading session from "{}"'.format(args.model_path))
			ckpt = tf.train.get_checkpoint_state(args.model_path)
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.full_saver.restore(sess, ckpt.model_checkpoint_path)
		except:
			print(Fore.RED + "Model not found..."+Style.RESET_ALL)
			sys.exit()


	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print(Fore.RED +"Error. Encoder not found...")
			sys.exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True


	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)


	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)],
				                   1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss


	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result


	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y, pi):

		# gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		# gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		#loss = tf.multiply(z + tf.log(sxsy), pi)
		loss = z + tf.log(sxsy)

		return tf.reduce_mean(loss)


	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):
		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #


	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors


	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector


	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations


	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation


	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/VRNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
from tensorflow_probability import distributions as tfd
import sys

if sys.version_info[0] < 3:
	sys.path.append('../src/external')
	from vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from tf_utils import *
else:
	from src.external.vrnn_cell2 import VariationalRNNCell as vrnn_cell
	from src.models.tf_utils import *
import numpy as np
import os
from colorama import Fore, Style

class NetworkModel():

	def __init__(self, args, is_training=True, batch_size=None):
		self.log_dir = args.log_dir
		self.args = args
		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_dim = args.output_dim
		self.output_pred_state_dim = self.args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.prev_horizon = args.prev_horizon
		self.n_mixtures = args.n_mixtures
		self.prediction_horizon = args.prediction_horizon
		self.output_placeholder_dim = self.prediction_horizon * self.output_dim
		self.output_vec_dim = self.prediction_horizon * self.output_pred_state_dim * self.n_mixtures
		self.pedestrian_vector_dim = args.pedestrian_vector_dim
		self.learning_rate_init = args.learning_rate_init
		self.beta_rate_init = args.beta_rate_init
		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size
		self.prior_hidden_size = self.rnn_state_size + self.rnn_state_size_lstm_grid + self.rnn_state_size_lstm_ped
		try:
			self.sigma_bias = args.sigma_bias
		except:
			self.sigma_bias = 0.0
		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.grid_width = int(args.submap_width / args.submap_resolution)
		self.grid_height = int(args.submap_height / args.submap_resolution)
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim * (self.prev_horizon + 1)],
		                                              name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    self.grid_width, self.grid_height], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim*self.args.n_other_agents], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')
		self.diversity_placeholder = tf.placeholder(dtype=tf.float32,
		                                            shape=[self.batch_size, self.truncated_backprop_length,
		                                                   self.output_placeholder_dim], name='output')
		self.dropout_placeholder = tf.placeholder(dtype=tf.float32, name='dropout')
		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')
		self.seq_length = tf.placeholder(tf.int32, [None])
		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Cells used for testing
		# Initialize hidden states with zeros
		self.test_cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.test_cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.test_cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.test_cell_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])
		self.test_hidden_state_current_lstm_concat = np.zeros([args.batch_size, self.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)
		outputs_diversity_series = tf.unstack(self.diversity_placeholder, axis=1)
		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		###################### Assemble CNN #####################
		# Process occupancy grid, apply convolutional filters
		self.occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
		conv_grid_feature_series = []
		for output in self.process_grid(self.occ_grid_series):
			conv_grid_feature_series.append(tf.stop_gradient(output))

		# Set up model with variables
		with tf.variable_scope("model") as scope:

			train_encoder = self.args.end_to_end

			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[ped_grid_out_dim, ped_grid_out_dim],
				                                                  trainable=True)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim], trainable=True)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					self.cell_grid = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell',
					                                         trainable=True)
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(self.cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					self.cell_ped = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell', trainable=True)
					ped_grid_feature_series , self.current_state_lstm_ped = tf.nn.dynamic_rnn(self.cell_ped,
					                                                                                           tf.stack(ped_grid_series,axis=1),
					                                                                                           dtype=tf.float32,
					                                                                                           sequence_length=self.seq_length,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

					# Embedding layer of pedestrian grid
					self.cell_outputs_series_lstm_ped = self.process_pedestrian_grid(tf.unstack(ped_grid_feature_series,axis=1))

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					self.cell = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size, name='basic_lstm_cell', trainable=True)
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(self.cell, inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state,
				                                               self.cell_outputs_series_lstm_grid,
				                                               self.cell_outputs_series_lstm_ped):
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

					# Concatenated LSTM layer
					with tf.variable_scope('generation_model') as scope:

						# cell_lstm_concat = tf.nn.rnn_cell.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
						cell_lstm_concat = vrnn_cell(args, args.rnn_state_size_lstm_concat, self.output_vec_dim, args.prior_size,
						                             args.latent_space_size)
						self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
							cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

						# Output reshape
						enc_mu_series = tf.unstack(tf.stack([o[0] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						enc_sigma_series = tf.unstack(tf.stack([o[1] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						self.dec_series = tf.unstack(tf.stack([o[2] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_mu_series = tf.unstack(tf.stack([o[3] for o in self.cell_outputs_series_lstm_concat]), axis=0)
						prior_sigma_series = tf.unstack(tf.stack([o[4] for o in self.cell_outputs_series_lstm_concat]), axis=0)

						# Decoder layers
						dec_fc_layer = tf.unstack(tf.stack(tf.nn.elu([output for output in self.dec_series])))
						dec_output_layer = []
						dec_mux_series = []
						dec_muy_series = []
						dec_sigmax_series = []
						dec_sigmay_series = []
						dec_pi_series = []
						# Reusing weights for each truncated back propagation step
						for t_idx in range(len(dec_fc_layer)):
							dec_output_layer.append(linear(dec_fc_layer[t_idx], self.output_vec_dim, 50))
							dec_pi_series.append(linear(dec_fc_layer[t_idx], self.args.n_mixtures, 100))
						for t_idx in range(len(dec_output_layer)):
							m_x, m_y, s_x, s_y = tf.split(dec_output_layer[t_idx], 4, axis=1)
							dec_mux_series.append(m_x)
							dec_muy_series.append(m_y)
							dec_sigmax_series.append(s_x)
							dec_sigmay_series.append(s_y)
						dec_pi_series = [dense for dense in dec_pi_series]

				# Exponential / stepwise decrease of learning rate
				global_step = tf.Variable(0, trainable=False)
				self.learning_rate = tf.train.exponential_decay(self.learning_rate_init, self.step, decay_steps=10000,
				                                                decay_rate=0.9, staircase=False)

				self.beta = (tf.tanh((tf.to_float(self.step) - 20000) / 5000) + 1) / 2

				prediction_loss_list = []
				loss_list = []
				kl_loss_list = []
				self.div_loss = []
				self.prediction = []
				self.likelihood = []
				pred = []
				div_loss_over_truncated_back_prop = []

				# Compute Loss Function
				for enc_mu, enc_sigma, dec_mux, dec_muy, dec_sigmax, dec_sigmay, prior_mu, prior_sigma, dec_pi, out, out_div in \
						zip(enc_mu_series, enc_sigma_series, dec_mux_series, dec_muy_series,
						    dec_sigmax_series, dec_sigmay_series,
						    prior_mu_series, prior_sigma_series, dec_pi_series,
						    outputs_series, outputs_diversity_series):  # prediction dimension: prediction_horizon * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])

					# Apply transformations
					mux = tf.split(dec_mux, self.prediction_horizon, axis=1)
					muy = tf.split(dec_muy, self.prediction_horizon, axis=1)
					sigmax = tf.exp(tf.split(dec_sigmax, self.prediction_horizon, axis=1))
					sigmay = tf.exp(tf.split(dec_sigmay, self.prediction_horizon, axis=1))
					pi = tf.nn.softmax(dec_pi, axis=1)

					self.likelihood.append(pi)
					pred = []
					prediction_loss_list = []
					kl_mode_loss_list = []
					div_loss_over_horizon = []
					lossfunc = 0
					for prediction_step in range(self.prediction_horizon):
						# Compute Euclidean error for each prediction step within each tbp step
						idx_x = prediction_step * self.output_dim
						idx_y = idx_x + 1

						# ground truth data
						x_data = tf.reshape(out[:, idx_x], shape=(-1, 1))
						y_data = tf.reshape(out[:, idx_y], shape=(-1, 1))
						x_diversity = tf.reshape(out_div[:, idx_x], shape=(-1, 1))
						y_diversity = tf.reshape(out_div[:, idx_y], shape=(-1, 1))

						if prediction_step == 0:
							pred = tf.concat([mux[prediction_step], muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)
						else:
							pred = tf.concat([pred, mux[prediction_step]], 1)
							pred = tf.concat([pred, muy[prediction_step]], 1)
							pred = tf.concat([pred, sigmax[prediction_step]], 1)
							pred = tf.concat([pred, sigmay[prediction_step]], 1)

						lossfunc = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                 sigmax[prediction_step], sigmay[prediction_step],
						                                 x_data, y_data, pi)

						diversity_loss = self.get_log_lossfunc(mux[prediction_step], muy[prediction_step],
						                                        sigmax[prediction_step], sigmay[prediction_step],
						                                       x_diversity, y_diversity, pi)

						div_loss_over_horizon.append(diversity_loss)
						prediction_loss_list.append(lossfunc)

					self.prediction.append(pred)
					kl_loss = self.tf_kl_gaussgauss(enc_mu, enc_sigma, prior_mu, prior_sigma)
					kl_loss_list.append(kl_loss)
					div_loss_over_truncated_back_prop.append(tf.reduce_mean(div_loss_over_horizon))
					loss_list.append(tf.reduce_mean(prediction_loss_list))

			self.prediction = tf.stack(self.prediction, axis=1)
			self.likelihood = tf.stack(self.likelihood, axis=1)

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.div_loss = tf.reduce_mean(div_loss_over_truncated_back_prop)
			self.total_loss = tf.reduce_mean(loss_list, axis=0)+args.diversity_update*self.div_loss + tf.reduce_mean(kl_loss_list, axis=0)*self.beta
			self.reconstruction_loss = tf.reduce_mean(loss_list, axis=0)
			self.kl_loss = tf.reduce_mean(kl_loss_list, axis=0)

			# Optimizer specification
			# self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			self.loss_summary = tf.summary.scalar('loss', self.total_loss)
			self.kl_loss_summary = tf.summary.scalar('kl-loss', self.kl_loss)
			self.div_loss_summary = tf.summary.scalar('div_loss', self.div_loss)
			self.reconstruction_loss_summary = tf.summary.scalar('reconstruction_loss', self.reconstruction_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)

			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'train'), graph=tf.Session().graph)
			self.validation_summary_writer = tf.summary.FileWriter(logdir=os.path.join(self.log_dir, 'validation'),
			                                                       graph=tf.Session().graph)

			self.summary = tf.summary.merge_all()

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(self.autoencoder_var_list)
			self.autoencoder_update = self.optimizer.minimize(self.autoencoder_loss, global_step=global_step)


	def feed_dic(self, **kwargs):
		if self.args.rotated_grid:
			batch_x, batch_y = sup.rotate_batch_to_local_frame(kwargs['batch_y'], kwargs['batch_x'])
			batch_x = batch_x[:, :, 2:]
		else:
			batch_x = kwargs['batch_vel']
			batch_y = kwargs['batch_y']
		n_other_agents = np.zeros([self.args.batch_size])
		for i,other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: batch_x,
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.step: kwargs['step'],
		        self.seq_length: n_other_agents,
		        self.diversity_placeholder: kwargs['batch_div'],
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: np.random.normal(self.cell_state_current_lstm_grid, 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.hidden_state_current_lstm_grid, 0),
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }


	def feed_test_dic(self, **kwargs):
		step = kwargs["step"]
		n_other_agents = np.zeros([self.args.batch_size])
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: np.expand_dims(kwargs["batch_vel"][:, step, :], axis=1),
		        self.input_ped_grid_placeholder: np.expand_dims(kwargs["batch_ped_grid"][:, step, :], axis=1),
		        self.input_grid_placeholder: np.expand_dims(kwargs["batch_grid"][:, step, :, :], axis=1),
		        self.step: 0,
		        self.seq_length: n_other_agents,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), 0),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), 0),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), 0),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), 0),
		        self.cell_state_lstm_ped: self.test_cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.test_hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def feed_val_dic(self, **kwargs):
		n_other_agents = np.zeros([self.args.batch_size])
		for i, other in enumerate(kwargs['other_agents_pos']):
			n_other_agents[i] = len(other)
		return {self.input_state_placeholder: kwargs["batch_vel"],
		        self.input_grid_placeholder: kwargs["batch_grid"],
		        self.input_ped_grid_placeholder: kwargs["batch_ped_grid"],
		        self.diversity_placeholder: kwargs['batch_div'],
		        self.output_placeholder: kwargs['batch_y'],
		        self.seq_length: n_other_agents,
		        self.step: 0,
		        self.cell_state: np.random.normal(self.test_cell_state_current.copy(), kwargs["state_noise"]),
		        self.hidden_state: np.random.normal(self.test_hidden_state_current.copy(), kwargs["state_noise"]),
		        self.cell_state_lstm_grid: np.random.normal(self.test_cell_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
		        self.hidden_state_lstm_grid: np.random.normal(self.test_hidden_state_current_lstm_grid.copy(), kwargs["grid_noise"]),
		        self.cell_state_lstm_ped: np.random.normal(self.test_cell_state_current_lstm_ped.copy(), kwargs["ped_noise"]),
		        self.hidden_state_lstm_ped: np.random.normal(self.test_hidden_state_current_lstm_ped.copy(), kwargs["ped_noise"]),
		        self.cell_state_lstm_concat: self.test_cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.test_hidden_state_current_lstm_concat,
		        }

	def reset_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])


	def reset_test_cells(self, sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.test_cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.test_cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.test_cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.test_cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.test_hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def update_test_hidden_state(self):
		self.test_cell_state_current = self.cell_state_current.copy()
		self.test_hidden_state_current = self.hidden_state_current.copy()
		self.test_cell_state_current_lstm_grid = self.cell_state_current_lstm_grid.copy()
		self.test_hidden_state_current_lstm_grid = self.hidden_state_current_lstm_grid.copy()
		self.test_cell_state_current_lstm_ped = self.cell_state_current_lstm_ped.copy()
		self.test_hidden_state_current_lstm_ped = self.hidden_state_current_lstm_ped.copy()
		self.test_cell_state_current_lstm_concat = self.cell_state_current_lstm_concat.copy()
		self.test_hidden_state_current_lstm_concat = self.hidden_state_current_lstm_concat.copy()

	def validation_step(self, sess, feed_dict_validation,update = True):
		batch_loss, _model_prediction,_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, summary = sess.run([self.reconstruction_loss,
		                                                self.prediction,
		                                                                    self.current_state,
		                                                                    self.current_state_lstm_grid,
		                                                                    self.current_state_lstm_ped,
		                                                                    self.current_state_lstm_concat,
		                                                                    self.loss_summary],
		                                                                   feed_dict=feed_dict_validation)

		if update:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, summary, _model_prediction

	def train_step(self, sess, feed_dict_train, step=0):

		# Diversity Training
		if step > 2000 and self.args.diversity_update:
			dict = {"batch_vel": feed_dict_train[self.input_state_placeholder],
			        "batch_grid": feed_dict_train[self.input_grid_placeholder],
			        "batch_ped_grid": feed_dict_train[self.input_ped_grid_placeholder],
			        "batch_y": feed_dict_train[self.output_placeholder],
			        "batch_div": feed_dict_train[self.diversity_placeholder],
			        "step": step,
			        "state_noise": 0.5,
			        "grid_noise": 0.0,
			        "ped_noise": 0.5
			        }
			self.update_test_hidden_state()
			feed_test_dic = self.feed_val_dic(**dict)
			loss, _, y_model_pred = self.validation_step(sess, feed_test_dic, False)
			batch_y_varible = np.zeros((self.args.batch_size, self.args.truncated_backprop_length,
			                            self.args.prediction_horizon * self.args.output_dim))
			# Grount truth
			batch_y = feed_dict_train[self.output_placeholder]

			diversity_distance = 0
			diff_traj_index = 0
			# Find most different trajectory index
			for batch_id in range(self.args.batch_size):
				for t_id in range(self.args.truncated_backprop_length):
					for mix_id in range(self.args.n_mixtures):
						for pred_id in range(self.args.prediction_horizon):
							new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures+ mix_id]
							new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + mix_id]
							batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
								[new_batch_x, new_batch_y])
						if diversity_distance < np.linalg.norm(batch_y_varible[batch_id, t_id]-batch_y[batch_id, t_id]):
							diversity_distance = np.linalg.norm(batch_y_varible[batch_id, t_id] - batch_y[batch_id, t_id])
							diff_traj_index = mix_id
					# Saving most different traj
					for pred_id in range(self.args.prediction_horizon):
						new_batch_x = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + diff_traj_index]
						new_batch_y = y_model_pred[batch_id,t_id,pred_id * self.args.output_pred_state_dim * self.args.n_mixtures + self.args.n_mixtures + diff_traj_index]
						batch_y_varible[batch_id, t_id, 2 * pred_id:2 * pred_id + self.args.output_dim] = np.array(
							[new_batch_x, new_batch_y])

			# Assemble feed dict for training
			feed_dict_train[self.diversity_placeholder] = batch_y_varible

		_, batch_loss, kl_loss, recons_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta, output_decoder, div_loss, autoencoder_loss = sess.run([self.update,
		                                                                                                  self.total_loss,
		                                                                                                  self.kl_loss,
		                                                                                                  self.reconstruction_loss,
		                                                                                                  self.current_state,
		                                                                                                  self.current_state_lstm_grid,
		                                                                                                  self.current_state_lstm_ped,
		                                                                                                  self.current_state_lstm_concat,
		                                                                                                  self.prediction,
		                                                                                                  self.summary,
		                                                                                                  self.learning_rate,
		                                                                                                  self.beta,
		                                                                                                  self.deconv1,
		                                                                                                  self.div_loss,
		                                                                                                  self.autoencoder_loss],
		                                                                                                 feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		# Train autoencoder if loss larger than threshold
		""""""
		if np.mean(autoencoder_loss) > 0.05:
			self.run_autoencoder(sess, feed_dict_train)
			self.convnet_saver.save(sess, self.args.pretrained_convnet_path + '/final-model.ckpt')

		out = {"batch_loss": recons_loss,
		       "model_predictions": _model_prediction,
		       "summary": _summary_str,
		       "learning_rate": lr,
		       "autoencoder_loss": autoencoder_loss,
		       "log loss": batch_loss,
		       "kl_loss": kl_loss,
		       "output_decoder": output_decoder
		       }

		return out


	def run_autoencoder(self, sess, feed_dict_train):
		_ = sess.run([self.autoencoder_update],
		             feed_dict=feed_dict_train)


	def predict(self, sess, feed_dict_train, update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, output_decoder, _likelihood = sess.run([self.current_state,
		                                              self.current_state_lstm_grid,
		                                              self.current_state_lstm_ped,
		                                              self.current_state_lstm_concat,
		                                              self.prediction,
		                                              self.deconv1,
                                                  self.likelihood],
		                                             feed_dict=feed_dict_train)
		if update_state:
			self.test_cell_state_current, self.test_hidden_state_current = _current_state
			self.test_cell_state_current_lstm_grid, self.test_hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.test_cell_state_current_lstm_ped, self.test_hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.test_cell_state_current_lstm_concat, self.test_hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, _likelihood


	def warmstart_model(self, args, sess):
		# Restore whole model
		try:
			print('Loading session from "{}"'.format(args.model_path))
			ckpt = tf.train.get_checkpoint_state(args.model_path)
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.full_saver.restore(sess, ckpt.model_checkpoint_path)
		except:
			print(Fore.RED + "Model not found..."+Style.RESET_ALL)
			sys.exit()


	def warmstart_model_encoder(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.pretrained_encoder_path))
		ckpt = tf.train.get_checkpoint_state(args.pretrained_encoder_path)
		if ckpt == None:
			print(Fore.RED +"Error. Encoder not found...")
			sys.exit()
		else:
			print('Restoring model {}'.format(ckpt.model_checkpoint_path))
			self.encoder_saver.restore(sess, ckpt.model_checkpoint_path)

		return True


	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not found...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)


	def get_diversity_loss(self, mux, muy, sigmax, sigmay):
		# It only works for n_mixtures =3
		div_loss = 0
		for mix_id in range(self.args.n_mixtures):
			mu1 = tf.concat([tf.expand_dims(mux[:, mix_id], axis=1), tf.expand_dims(muy[:, mix_id], axis=1)], 1)
			sigma1 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
			gaussian1 = tfd.MultivariateNormalDiag(
				loc=mu1,
				scale_diag=sigma1)
			if mix_id + 1 < self.args.n_mixtures:
				mu2 = tf.concat([tf.expand_dims(mux[:, mix_id + 1], axis=1), tf.expand_dims(muy[:, mix_id + 1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id + 1], axis=1), tf.expand_dims(sigmay[:, mix_id + 1], axis=1)],
				                   1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)
			else:
				mu2 = tf.concat([tf.expand_dims(mux[:, -1], axis=1), tf.expand_dims(muy[:, -1], axis=1)], 1)
				sigma2 = tf.concat([tf.expand_dims(sigmax[:, mix_id], axis=1), tf.expand_dims(sigmay[:, mix_id], axis=1)], 1)
				gaussian2 = tfd.MultivariateNormalDiag(
					loc=mu2,
					scale_diag=sigma2)

			div_loss += tfd.kl_divergence(gaussian1, gaussian2)

			return div_loss


	def tf_2d_normal(self, x, y, mux, muy, sx, sy):
		'''
		Function that implements the PDF of a 2D normal distribution
		params:
		x : input x points
		y : input y points
		mux : mean of the distribution in x
		muy : mean of the distribution in y
		sx : std dev of the distribution in x
		sy : std dev of the distribution in y
		rho : Correlation factor of the distribution
		'''
		# eq 3 in the paper
		# and eq 24 & 25 in Graves (2013)
		# Calculate (x - mux) and (y-muy)
		rho = tf.constant(0.0)  # hack because we dont want correlation now
		normx = tf.subtract(x, mux)
		normy = tf.subtract(y, muy)
		# Calculate sx*sy
		sxsy = tf.multiply(sx, sy)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, sx)) + tf.square(tf.div(normy, sy)) - 2 * tf.div(
			tf.multiply(rho, tf.multiply(normx, normy)),
			sxsy)
		negRho = 1 - tf.square(rho)
		# Numerator
		result = tf.exp(tf.div(-z, 2 * negRho))
		# Normalization constant
		denom = 2 * np.pi * tf.multiply(sxsy, tf.sqrt(negRho))
		# Final PDF calculation
		result = tf.div(result, denom)
		return result


	def tf_kl_gaussgauss(self, mu_1, sigma_1, mu_2, sigma_2):
		with tf.variable_scope("kl_gaussgauss"):
			# return tf.reduce_sum(0.5 * (
			#    2 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 2 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_1 - mu_2)) / tf.maximum(1e-9, (tf.square(sigma_2))) - 1
			# ), 1)
			# return tf.reduce_sum(0.5 * (
			#    1.0 * tf.math.log(tf.maximum(1e-9, sigma_2), name='log_sigma_2')
			#    - 1.0 * tf.math.log(tf.maximum(1e-9, sigma_1), name='log_sigma_1')
			#    + (tf.square(sigma_1) + tf.square(mu_2 - mu_1)) / tf.maximum(1e-9, tf.square(sigma_2)) - 1.0
			# ), 1)
			kl = tfd.MultivariateNormalDiag(loc=mu_1, scale_diag=sigma_1).kl_divergence(
				tfd.MultivariateNormalDiag(loc=mu_2, scale_diag=sigma_2))
			return tf.reduce_mean(kl)

	def get_log_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, x, y, pi):

		# gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		# gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		normx = tf.subtract(x, dec_mux)
		normy = tf.subtract(y, dec_muy)
		# Calculate sx*sy
		sxsy = tf.multiply(dec_sigmax, dec_sigmay)
		# Calculate the exponential factor
		z = tf.square(tf.div(normx, dec_sigmax)) + tf.square(tf.div(normy, dec_sigmay))

		loss = tf.multiply(z + tf.log(sxsy), pi)

		return tf.reduce_mean(loss)


	def get_mixture_lossfunc(self, dec_mux, dec_muy, dec_sigmax, dec_sigmay, dec_pi, x, y):
		gaussian = self.tf_2d_normal(x, y, dec_mux, dec_muy, dec_sigmax, dec_sigmay)
		gaussian = tf.multiply(gaussian, dec_pi)  # does element-wise multiplication
		gaussian = tf.reduce_sum(gaussian, axis=1)  # do inner summation

		likelihood_loss = -tf.math.log(tf.maximum(gaussian, 1e-10))
		return tf.reduce_mean(likelihood_loss)  #


	def process_grid(self, occ_grid_series):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""
		# Convolutional layer 1
		conv1_kernel_size = 5
		conv1_number_filters = 64
		conv1_stride_length = 2

		# Convolutional layer 2
		conv2_kernel_size = 3
		conv2_number_filters = 32
		conv2_stride_length = 2

		# Convolutional layer 3
		conv3_kernel_size = 3
		conv3_number_filters = 8
		conv3_stride_length = 2

		# Flatten convolutional filter vector such that it can be fed to a FC layer
		conv_grid_size = 8
		fc_grid_hidden_dim = 64

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)

			if idx == 0:
				with tf.variable_scope('auto_encoder') as scope:
					conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
					                                 conv1_stride_length, trainable=True)
					conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
					                                 conv2_stride_length, trainable=True)
					conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
					                                 conv3_stride_length, padding='VALID', trainable=False)

					conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])

					fc_grid_weights = self.get_weight_variable(shape=[conv3_flat.get_shape()[1], fc_grid_hidden_dim],
					                                           name="fc_grid_weights",
					                                           trainable=True)
					fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
					                                        trainable=True)

					fc_final = self.fc_layer(conv3_flat, fc_grid_weights, fc_grid_biases)

					tf.summary.histogram("fc_activations", fc_final)
					fc_grid_dec_biases = self.get_bias_variable(shape=[conv3_flat.get_shape()[1]],
					                                            name="fc_grid_decode_biases", trainable=True)
					fc_grid_dec = self.fc_layer(fc_final, tf.transpose(fc_grid_weights), fc_grid_dec_biases)
					print("Dim out fc_grid_dec (conv2_flat_out): {}".format(fc_grid_dec.get_shape()))

					# Reshaping into known convolutional grid size
					conv3_dec = tf.reshape(fc_grid_dec,
					                       shape=[-1, conv_grid_size - 1, conv_grid_size - 1, conv3_number_filters])
					print("Dim out conv3 reshape (conv3_out): {}".format(conv3_dec.get_shape()))

					# From output of conv3 to input of conv3 (output of conv2)
					deconv3 = tf.contrib.layers.conv2d_transpose(conv3_dec, conv2_number_filters,
					                                             [conv3_kernel_size, conv3_kernel_size], conv3_stride_length,
					                                             padding='VALID', trainable=True)
					print("Dim out deconv3 (conv2_out): {}".format(deconv3.get_shape()))

					# From output of conv2 to input of conv2 (output of conv1)
					deconv2 = tf.contrib.layers.conv2d_transpose(deconv3, conv1_number_filters,
					                                             [conv2_kernel_size, conv2_kernel_size], conv2_stride_length,
					                                             padding='SAME', trainable=True)
					print("Dim out deconv2 (conv1_out): {}".format(deconv2.get_shape()))

					# From output of conv1 to input of conv1 (raw grid input)
					self.deconv1 = tf.contrib.layers.conv2d_transpose(deconv2, 1, [conv1_kernel_size, conv1_kernel_size],
					                                                  conv1_stride_length, padding='SAME', trainable=True)
					print("Dim out deconv1 (image input): {}".format(self.deconv1.get_shape()))

					self.autoencoder_var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='auto_encoder')

					self.autoencoder_loss = tf.reduce_mean(
						tf.square(tf.abs(grid_batch - self.deconv1)), name="autoencoder_loss")


			else:
				conv1 = tf.contrib.layers.conv2d(grid_batch, conv1_number_filters, [conv1_kernel_size, conv1_kernel_size],
				                                 conv1_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv')
				conv2 = tf.contrib.layers.conv2d(conv1, conv2_number_filters, [conv2_kernel_size, conv2_kernel_size],
				                                 conv2_stride_length, trainable=True, reuse=True, scope='auto_encoder/Conv_1')
				conv3 = tf.contrib.layers.conv2d(conv2, conv3_number_filters, [conv3_kernel_size, conv3_kernel_size],
				                                 conv3_stride_length, padding='VALID', trainable=True, reuse=True,
				                                 scope='auto_encoder/Conv_2')

				conv3_flat = tf.reshape(conv3, [-1, (conv_grid_size - 1) * (conv_grid_size - 1) * conv3_number_filters])
				fc_final = self.fc_layer(input=conv3_flat, weights=fc_grid_weights, biases=fc_grid_biases)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors


	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			# fc_feature_vector.append(self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid,activation = tf.nn.relu, name="fc_ped"))
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector


	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		#                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1), trainable=trainable)
		# var = tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(),
		#                      trainable=trainable)
		var = tf.get_variable(name=name, shape=shape,
		                      initializer=tf.contrib.layers.variance_scaling_initializer(factor=2.0, mode='FAN_IN',
		                                                                                 uniform=False, seed=None,
		                                                                                 dtype=tf.float32),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, mean=0.1, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var


	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations


	def deconv_layer(self, input, weights, biases, out_shape, conv_stride_length=1, padding="SAME"):
		#
		input_shape = input.get_shape()
		deconv = tf.nn.conv2d_transpose(input, weights, out_shape, [1, conv_stride_length, conv_stride_length, 1],
		                                padding=padding)
		activation = tf.nn.relu(deconv + biases)
		return activation


	def fc_layer(self, input, weights, biases, activation=None, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if activation is None:
			activations = affine_result
		else:
			activations = activation(affine_result)
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations

src/models/VRNNwLikelihood.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
import numpy as np
import os
import sys
sys.path.append('./autoencoder_grid')
tfk = tf.keras
tfkl = tf.keras.layers

class EncoderState(tfk.Model):
  def __init__(self, rnn_state_size,truncated_backprop_length, input_state_dim,batch_size):
    super(EncoderState, self).__init__()

    self.rnn_state_size = rnn_state_size
    self.truncated_backprop_length = truncated_backprop_length
    self.input_state_dim = input_state_dim
    self.batch_size = batch_size
    self.lambda_ = 0.01

    self.agent_state_lstm = tfkl.GRU(self.rnn_state_size, input_shape=[self.truncated_backprop_length, self.input_state_dim],
                                 name='ped_lstm',return_sequences=True,
                                 return_state=True,
                                 kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))

  def call(self, x, hidden):
    output, state = self.agent_state_lstm(x, initial_state = hidden)
    return output, state

  def initialize_hidden_state(self):
    return tf.zeros((self.batch_size, self.rnn_state_size))

class EncoderPeds(tfk.Model):
  def __init__(self, rnn_state_size_lstm_ped, truncated_backprop_length, input_state_dim, batch_size):
    super(EncoderPeds, self).__init__()

    self.rnn_state_size_lstm_ped = rnn_state_size_lstm_ped
    self.truncated_backprop_length = truncated_backprop_length
    self.input_state_dim = input_state_dim
    self.lambda_ = 0.01
    self.batch_size = batch_size
    self.ped_embed = tfkl.Dense(self.rnn_state_size_lstm_ped, activation='relu')

    self.agent_other_peds_lstm = tfkl.GRU(self.rnn_state_size_lstm_ped,
                                     input_shape=[self.truncated_backprop_length, self.input_state_dim],
                                     name='other_peds_lstm', return_sequences=True,
                                     return_state=True,
                                     kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))

  def call(self, x, hidden):
    x = self.ped_embed(x)
    output, state = self.agent_other_peds_lstm(x, initial_state=hidden)
    return output, state

  def initialize_hidden_state(self):
    return tf.zeros((self.batch_size, self.rnn_state_size_lstm_ped))

class BahdanauAttention(tfkl.Layer):
  def __init__(self, units):
    super(BahdanauAttention, self).__init__()
    self.W1 = tf.keras.layers.Dense(units)
    self.W2 = tf.keras.layers.Dense(units)
    self.V = tf.keras.layers.Dense(1)

  def call(self, query, values):
    # hidden shape == (batch_size, hidden size)
    # hidden_with_time_axis shape == (batch_size, 1, hidden size)
    # we are doing this to perform addition to calculate the score
    hidden_with_time_axis = tf.expand_dims(query, 1)

    # score shape == (batch_size, max_length, 1)
    # we get 1 at the last axis because we are applying score to self.V
    # the shape of the tensor before applying self.V is (batch_size, max_length, units)
    score = self.V(tf.nn.tanh(
        self.W1(values) + self.W2(hidden_with_time_axis)))

    # attention_weights shape == (batch_size, max_length, 1)
    attention_weights = tf.nn.softmax(score, axis=1)

    # context_vector shape after sum == (batch_size, hidden_size)
    context_vector = attention_weights * values
    context_vector = tf.reduce_sum(context_vector, axis=1)

    return context_vector, attention_weights

def scaled_dot_product_attention(q, k, v, mask):
  """Calculate the attention weights.
  q, k, v must have matching leading dimensions.
  k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
  The mask has different shapes depending on its type(padding or look ahead)
  but it must be broadcastable for addition.

  Args:
    q: query shape == (..., seq_len_q, depth)
    k: key shape == (..., seq_len_k, depth)
    v: value shape == (..., seq_len_v, depth_v)
    mask: Float tensor with shape broadcastable
          to (..., seq_len_q, seq_len_k). Defaults to None.

  Returns:
    output, attention_weights
  """

  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

  # scale matmul_qk
  dk = tf.cast(tf.shape(k)[-1], tf.float32)
  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

  # add the mask to the scaled tensor.
  if mask is not None:
    scaled_attention_logits += (mask * -1e9)

    # softmax is normalized on the last axis (seq_len_k) so that the scores
  # add up to 1.
  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

  return output, attention_weights

class MultiHeadAttention(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads):
    super(MultiHeadAttention, self).__init__()
    self.num_heads = num_heads
    self.d_model = d_model

    assert d_model % self.num_heads == 0

    self.depth = d_model // self.num_heads

    self.wq = tf.keras.layers.Dense(d_model)
    self.wk = tf.keras.layers.Dense(d_model)
    self.wv = tf.keras.layers.Dense(d_model)

    self.dense = tf.keras.layers.Dense(d_model)

  def split_heads(self, x, batch_size):
    """Split the last dimension into (num_heads, depth).
    Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
    """
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
    return tf.transpose(x, perm=[0, 2, 1, 3])

  def call(self, v, k, q, mask):
    batch_size = tf.shape(q)[0]

    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)

    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
    scaled_attention, attention_weights = scaled_dot_product_attention(
      q, k, v, mask)

    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention,
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

    return output, attention_weights

class Decoder(tfk.Model):
  def __init__(self, args):
    super(Decoder, self).__init__()
    self.batch_size = args.batch_size
    self.truncated_backprop_length = args.truncated_backprop_length
    self.prediction_horizon = args.prediction_horizon
    self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
    self.fc_hidden_unit_size = args.fc_hidden_unit_size
    self.output_dim = args.output_dim
    self.lambda_ = 0.01

    self.gru = tf.keras.layers.GRU(self.rnn_state_size_lstm_concat,
                                   return_sequences=True,
                                   return_state=True,
                                   recurrent_initializer='glorot_uniform')
    self.fc1 = tfkl.Dense(self.fc_hidden_unit_size, activation='relu', name='concat_fc',
                                              kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))
    self.fco = tfkl.Dense(self.output_dim*self.prediction_horizon, activation='linear', name='out_fc',
                                              kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))

    # used for attention
    self.attention = BahdanauAttention(self.rnn_state_size_lstm_concat)

  def call(self, x, hidden,ped_info):

    # x shape after concatenation == (batch_size, 1, embedding_dim + hidden_size)
    concat_x = tf.concat([ped_info, x], axis=-1)

    context_vector, attention_weights = self.attention(hidden, concat_x)

    context_vector = tf.expand_dims(context_vector, 1)

    # passing the concatenated vector to the GRU
    output, state = self.gru(concat_x,initial_state=hidden)

    # two FCL
    output = self.fc1(output)
    output = self.fco(output)

    return output, state, attention_weights

  def initialize_hidden_state(self):
    return tf.zeros((self.batch_size, self.rnn_state_size_lstm_concat))

class RNN(tf.keras.Model):

  def __init__(self,args):
    super(RNN, self).__init__()
    self.args = args

    # Network Elements

    self.agent_state_encoder = EncoderState(self.args.rnn_state_size, self.args.truncated_backprop_length, self.args.input_state_dim, self.args.batch_size)
    self.ped_info_encoder = EncoderPeds(self.args.rnn_state_size_lstm_ped, self.args.truncated_backprop_length, self.args.input_state_dim, self.args.batch_size)
    self.decoder = Decoder(self.args)
    # Initialize States
    self.state_hidden = self.agent_state_encoder.initialize_hidden_state()
    self.ped_info_hidden = self.ped_info_encoder.initialize_hidden_state()
    self.dec_hidden = self.decoder.initialize_hidden_state()

    # Optimizer and Loss
    self.optimizer = tf.keras.optimizers.Adam()
    self.loss_object = tf.keras.losses.Huber()

  def call(self, x, ped_info):
    enc_state_output, self.state_hidden = self.agent_state_encoder(x, self.state_hidden)

    enc_ped_output, self.ped_info_hidden = self.ped_info_encoder(ped_info, self.ped_info_hidden)

    return self.decoder(enc_state_output, self.dec_hidden,enc_ped_output)

  def loss_function(self,real, pred):
    loss_ = self.loss_object(real, pred)

    return tf.reduce_mean(loss_)

  @tf.function
  def train_step(self,x, ped, targ):
    loss = 0

    with tf.GradientTape() as tape:
      enc_state_output, enc_state_hidden = self.agent_state_encoder(x, self.state_hidden)

      enc_ped_output, enc_ped_info_hidden = self.ped_info_encoder(ped, self.ped_info_hidden)

      predictions, dec_hidden, attention_weights = self.decoder(enc_state_output, self.dec_hidden,
                                                           enc_ped_output)

      # Teacher forcing - feeding the target as the next input
      for t in range(1, targ.shape[1]):
        loss += self.loss_function(targ[:, t], predictions[:, t])

    batch_loss = (loss / int(targ.shape[1]))

    gradients = tape.gradient(loss, self.trainable_variables)

    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

    return batch_loss, enc_state_hidden, enc_ped_info_hidden, dec_hidden, predictions

  @tf.function
  def predict(self,batch_x, other_agents_info):
    enc_state_output, enc_state_hidden = self.agent_state_encoder(batch_x, self.state_hidden)

    enc_ped_output, enc_ped_info_hidden = self.ped_info_encoder(other_agents_info, self.ped_info_hidden)

    predictions, dec_hidden, attention_weights = self.decoder(enc_state_output, self.dec_hidden,
                                                              enc_ped_output)

    return predictions
src/models/modelKerasRNN.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
tfk = tf.keras
tfkm = tf.keras.models
tfkl = tf.keras.layers

class StateEncoder(tfkl.Layer):
    def __init__(self, args):
        super().__init__()
        
        self.lambda_ = 0.01
        self.rnn_state_size = args.rnn_state_size
        
        self.lstm_ped = tfkl.LSTM(self.rnn_state_size, name = 'lstm_state', return_sequences = False, kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))
        
    def call(self, x):
        x = self.lstm_ped(x)
        return x

class OtherPedsEncoder(tfkl.Layer):
    def __init__(self, args):
        super().__init__()

        self.lambda_ = 0.01
        self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
        self.rnn_state_size_bilstm = args.rnn_state_size_lstm_ped # TODO: Possibly add another argument to modify this state size

        self.lstm_other_peds = tfkl.LSTM(self.rnn_state_size_lstm_ped, name = 'lstm_other_peds', return_sequences = False, kernel_regularizer=tfk.regularizers.l2(l=self.lambda_))
        self.bilstm = tfkl.Bidirectional(tfkl.LSTM(self.rnn_state_size_bilstm,  name = 'bilstm_other_peds', return_sequences = False, return_state = False, kernel_regularizer=tfk.regularizers.l2(l=self.lambda_)), merge_mode = 'ave') # Average of forward and backward LSTMs
        
    def call(self, x):
        other_peds = []
        for x_ped in x:
            other_peds.append(self.lstm_other_peds(x_ped))
        stacked_other_peds = tf.stack(other_peds, axis = 1)
        out = self.bilstm(stacked_other_peds)
        return out

class Decoder(tfkl.Layer):
    def __init__(self, args):
        super().__init__()
        
        self.output_dim = args.output_dim
    
        self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
        self.fc_hidden_unit_size = args.fc_hidden_unit_size

        self.lstm_concat = tfkl.LSTM(self.rnn_state_size_lstm_concat, name = 'lstm_decoder', return_sequences = True, return_state = False)
        self.fc1 = tfkl.TimeDistributed(tfkl.Dense(self.fc_hidden_unit_size, name = 'fc_decoder', activation = 'relu'))
        self.fco = tfkl.TimeDistributed(tfkl.Dense(self.output_dim, name = 'fc_out'))
                
    def call(self, x):        
        x = self.lstm_concat(x)
        x = self.fc1(x)
        x = self.fco(x)
        return x

class NetworkModel(tfk.Model):
    def __init__(self, args):
        super().__init__()
        
        future_steps = args.prediction_horizon
        
        # Define optimizer and loss
        self.loss_object = tfk.losses.MeanSquaredError()
        self.optimizer = tfk.optimizers.Adam()
        self.train_loss = tfk.metrics.MeanSquaredError(name='train_loss')
        self.val_loss = tfk.metrics.MeanSquaredError(name='val_loss')
        
        # Define architecture
        self.state_encoder = StateEncoder(args)
        self.other_peds_encoder = OtherPedsEncoder(args)
        self.concat = tfkl.Concatenate()
        self.repeat = tfkl.RepeatVector(future_steps)
        self.decoder = Decoder(args)
                
    def call(self, x):
        x1 = self.state_encoder(x[0])
        
        x2 = self.other_peds_encoder(x[1:])
        concat = self.concat([x1, x2])
        repeated = self.repeat(concat)
        out = self.decoder(repeated)
        return out
    
    @tf.function
    def train_step(self, X, Y):
        with tf.GradientTape() as tape:
            predictions = self(X)
            loss = self.loss_object(Y, predictions)
        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        self.train_loss(Y, predictions)

    @tf.function
    def val_step(self, X, Y):
        predictions = self(X)
        v_loss = self.loss_object(Y, predictions)

        self.val_loss(Y, predictions)



src/models/modelKerasRNN_arbitraryAgents.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
# import tensorflow_probability as tfp
import numpy as np
import os


class NetworkModel():

	def __init__(self, args, batch_size=None, is_training=True):

		# Input- / Output dimensions
		self.input_dim = args.input_dim  # [x, y, vx, vy]
		self.input_state_dim = args.input_state_dim  # [vx, vy]
		self.output_state_dim = args.output_dim
		self.output_pred_state_dim = args.output_pred_state_dim
		self.truncated_backprop_length = args.truncated_backprop_length
		self.n_mixtures = args.n_mixtures
		self.output_sequence_length = args.prediction_horizon
		self.output_placeholder_dim = self.output_sequence_length * self.output_state_dim
		self.output_dim = self.output_sequence_length * self.output_pred_state_dim
		self.pedestrian_vector_dim = args.pedestrian_vector_dim

		# Network attributes
		self.rnn_state_size = args.rnn_state_size
		self.rnn_state_size_lstm_grid = args.rnn_state_size_lstm_grid
		self.rnn_state_size_lstm_ped = args.rnn_state_size_lstm_ped
		self.rnn_state_size_lstm_concat = args.rnn_state_size_lstm_concat
		self.fc_hidden_unit_size = args.fc_hidden_unit_size

		# Training parameters
		self.lambda_ = args.regularization_weight
		self.batch_size = None
		self.is_training = is_training
		self.grads_clip = args.grads_clip
		self.regularization_weight = args.regularization_weight
		self.keep_prob = args.keep_prob
		# Specify placeholders
		self.input_state_placeholder = tf.placeholder(dtype=tf.float32,
		                                              shape=[self.batch_size, self.truncated_backprop_length,
		                                                     self.input_state_dim], name='input_state')
		self.input_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.truncated_backprop_length,
		                                                    int(args.submap_width/args.submap_resolution),
		                                                    int(args.submap_height/args.submap_resolution)], name='input_grid')
		self.input_ped_grid_placeholder = tf.placeholder(dtype=tf.float32,
		                                                 shape=[self.batch_size, self.truncated_backprop_length,
		                                                        self.pedestrian_vector_dim], name='ped_grid')
		self.output_placeholder = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.truncated_backprop_length,
		                                                                  self.output_placeholder_dim], name='output')

		self.step = tf.placeholder(dtype=tf.float32,
		                           shape=[], name='global_step')

		# Initialize hidden states with zeros
		self.cell_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.hidden_state_current = np.zeros([args.batch_size, args.rnn_state_size])
		self.cell_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.hidden_state_current_lstm_grid = np.zeros([args.batch_size, args.rnn_state_size_lstm_grid])
		self.cell_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.hidden_state_current_lstm_ped = np.zeros([args.batch_size, args.rnn_state_size_lstm_ped])
		self.cell_state_current_lstm_concat = np.zeros([args.batch_size, args.rnn_state_size_lstm_concat])
		self.hidden_state_current_lstm_concat = np.zeros([args.batch_size, args.rnn_state_size_lstm_concat])

		# Pedestrian state LSTM
		self.cell_state = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size],
		                                 name='cell_state')  # internal state of the cell (before output gate)
		self.hidden_state = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size],
		                                   name='hidden_state')  # output of the cell (after output gate)
		self.init_state_tuple = tf.contrib.rnn.LSTMStateTuple(self.cell_state, self.hidden_state)

		# Static occupancy grid LSTM
		self.cell_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_grid],
		                                           name='cell_state_lstm_grid')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_grid = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_grid],
		                                             name='hidden_state_lstm_grid')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_grid = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_grid,
		                                                                self.hidden_state_lstm_grid)

		# Pedestrian LSTM
		self.cell_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, self.rnn_state_size_lstm_ped],
		                                          name='cell_state_lstm_ped')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_ped = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_ped],
		                                            name='hidden_state_lstm_ped')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_ped = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_ped, self.hidden_state_lstm_ped)

		# Concatenation LSTM (state, grid, pedestrians)
		self.cell_state_lstm_concat = tf.placeholder(dtype=tf.float32,
		                                             shape=[self.batch_size, self.rnn_state_size_lstm_concat],
		                                             name='cell_state_lstm_concat')  # internal state of the cell (before output gate)
		self.hidden_state_lstm_concat = tf.placeholder(dtype=tf.float32, shape=[None, self.rnn_state_size_lstm_concat],
		                                               name='hidden_state_lstm_concat')  # output of the cell (after output gate)
		self.init_state_tuple_lstm_concat = tf.contrib.rnn.LSTMStateTuple(self.cell_state_lstm_concat,
		                                                                  self.hidden_state_lstm_concat)

		inputs_series = tf.unstack(self.input_state_placeholder, axis=1)
		outputs_series = tf.unstack(self.output_placeholder, axis=1)

		# Print network info
		print("Input list length: {}".format(len(inputs_series)))
		print("Output list length: {}".format(len(outputs_series)))
		print("Single input shape: {}".format(inputs_series[0].get_shape()))
		print("Single output shape: {}".format(outputs_series[0].get_shape()))

		# Set up model with variables
		with tf.variable_scope('model') as scope:

			# CNN weights for static occupancy grid processing
			train_grid_encoder_conv = False
			train_grid_encoder_fc = False
			use_summary_convnet = False

			# Convolutional layer 1
			self.conv1_kernel_size = 5
			self.conv1_number_filters = 64
			self.conv1_stride_length = 2
			self.conv1_weights = self.get_weight_variable(name="conv1_weights",
			                                              shape=[self.conv1_kernel_size, self.conv1_kernel_size, 1,
			                                                     self.conv1_number_filters],
			                                              trainable=train_grid_encoder_conv, summary=use_summary_convnet)
			self.conv1_biases = self.get_bias_variable(name="conv1_biases", shape=[self.conv1_number_filters],
			                                           trainable=train_grid_encoder_conv, summary=use_summary_convnet)

			# Convolutional layer 2
			self.conv2_kernel_size = 3
			self.conv2_number_filters = 32
			self.conv2_stride_length = 2
			self.conv2_weights = self.get_weight_variable(name="conv2_weights",
			                                              shape=[self.conv2_kernel_size, self.conv2_kernel_size,
			                                                     self.conv1_number_filters, self.conv2_number_filters],
			                                              trainable=train_grid_encoder_conv, summary=use_summary_convnet)
			self.conv2_biases = self.get_bias_variable(name="conv2_biases", shape=[self.conv2_number_filters],
			                                           trainable=train_grid_encoder_conv, summary=use_summary_convnet)

			# Convolutional layer 3
			self.conv3_kernel_size = 3
			self.conv3_number_filters = 8
			self.conv3_stride_length = 2
			self.conv3_weights = self.get_weight_variable(name="conv3_weights",
			                                              shape=[self.conv3_kernel_size, self.conv3_kernel_size,
			                                                     self.conv2_number_filters, self.conv3_number_filters],
			                                              trainable=train_grid_encoder_conv, summary=use_summary_convnet)
			self.conv3_biases = self.get_bias_variable(name="conv3_biases", shape=[self.conv3_number_filters],
			                                           trainable=train_grid_encoder_conv, summary=use_summary_convnet)

			fc_grid_hidden_dim = 64
			self.fc_grid_weights = self.get_weight_variable(shape=[512, fc_grid_hidden_dim], name="fc_grid_weights",
			                                                trainable=train_grid_encoder_fc)
			self.fc_grid_biases = self.get_bias_variable(shape=[fc_grid_hidden_dim], name="fc_grid_biases",
			                                             trainable=train_grid_encoder_fc)
			with tf.variable_scope('encoder_model') as scope:
				# Pedestrian grid embedding layer
				ped_grid_out_dim = 128
				self.W_pedestrian_grid = self.get_weight_variable(name='w_ped_grid',
				                                                  shape=[self.pedestrian_vector_dim, ped_grid_out_dim],
				                                                  trainable=self.is_training)
				self.b_pedestrian_grid = self.get_bias_variable(name='b_ped_grid', shape=[1, ped_grid_out_dim],
				                                                trainable=self.is_training)

				###################### Assemble Neural Net #####################
				# Process occupancy grid, apply convolutional filters
				occ_grid_series = tf.unstack(self.input_grid_placeholder, axis=1)
				conv_grid_feature_series = self.process_grid(occ_grid_series, trainable=self.is_training)

				# Process pedestrian grid
				ped_grid_series = tf.unstack(self.input_ped_grid_placeholder, axis=1)
				# Embedding layer of pedestrian grid
				ped_grid_feature_series = self.process_pedestrian_grid(ped_grid_series)

				# LSTM cells to process static occupancy grid
				with tf.variable_scope('lstm_grid') as scope:
					cell_grid = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_grid, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_grid, self.current_state_lstm_grid = tf.contrib.rnn.static_rnn(cell_grid,
					                                                                                             conv_grid_feature_series,
					                                                                                             dtype=tf.float32,
					                                                                                             initial_state=self.init_state_tuple_lstm_grid)

				# LSTM cells to process pedestrian grid
				with tf.variable_scope('lstm_ped') as scope:
					cell_ped = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_ped, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_ped, self.current_state_lstm_ped = tf.contrib.rnn.static_rnn(cell_ped,
					                                                                                           ped_grid_feature_series,
					                                                                                           dtype=tf.float32,
					                                                                                           initial_state=self.init_state_tuple_lstm_ped)

				# LSTM cells to process pedestrian state
				with tf.variable_scope('lstm_state') as scope:
					cell = tf.keras.layers.LSTMCell(self.rnn_state_size, name='basic_lstm_cell')
					self.cell_outputs_series_state, self.current_state = tf.contrib.rnn.static_rnn(cell,
					                                                                               inputs_series,
					                                                                               dtype=tf.float32,
					                                                                               initial_state=self.init_state_tuple)

				# Concatenate the outputs of the three previous LSTMs and feed in first concatenated LSTM cell
				concat_lstm = []
				for lstm_out, grid_feature, ped_feature in zip(self.cell_outputs_series_state, self.cell_outputs_series_lstm_grid,
				                                               self.cell_outputs_series_lstm_ped):
					concat_lstm.append(
						tf.concat([lstm_out, grid_feature, ped_feature], axis=1, name="concatenate_lstm_out_with_grid_and_ped"))

			with tf.variable_scope('decoder_model') as scope:
				# Hidden layer (fully connected) after concatenation LSTM
				W_fc = self.get_weight_variable(name='w_fc',
				                                shape=[self.rnn_state_size_lstm_concat, self.fc_hidden_unit_size],
				                                trainable=self.is_training)
				b_fc = self.get_bias_variable(name='b_fc', shape=[1, self.fc_hidden_unit_size],
				                              trainable=self.is_training)

				# Output layer (fully connected)
				W_out = self.get_weight_variable(name='w_out', shape=[self.fc_hidden_unit_size, self.output_dim],
				                                 trainable=self.is_training)
				b_out = self.get_bias_variable(name='b_out', shape=[1, self.output_dim], trainable=self.is_training)

				# Second concatenated LSTM layer
				with tf.variable_scope('lstm_concat') as scope:
					cell_lstm_concat = tf.keras.layers.LSTMCell(self.rnn_state_size_lstm_concat, name='basic_lstm_cell')
					self.cell_outputs_series_lstm_concat, self.current_state_lstm_concat = tf.contrib.rnn.static_rnn(
						cell_lstm_concat, concat_lstm, dtype=tf.float32, initial_state=self.init_state_tuple_lstm_concat)

				# FC layers
				fc_fuse = [tf.nn.relu(tf.add(tf.matmul(lstm_out, W_fc), b_fc)) for lstm_out in
				           self.cell_outputs_series_lstm_concat]
				fc_drop = [tf.layers.dropout(inputs=dense, rate=1-self.keep_prob) for dense in
				           fc_fuse]
				# fc_fuse2 = [tf.nn.elu(tf.add(tf.matmul(fc_out, W_fc2), b_fc2)) for fc_out in fc_fuse]
				# Model prediction
				self.prediction = [tf.add(tf.matmul(fc_out, W_out), b_out) for fc_out in fc_fuse]

			# Loss function
			loss_list = []
			lossfunc = []
			# Iterate through tbp steps
			for pred, out in zip(self.prediction,
			                     outputs_series):  # prediction dimension: output_sequence_length * output_states (e.g. [x_1, y_1, x_2, y_2, x_3, y_3])
				prediction_loss_list = []
				val_prediction_loss_list = []
				for prediction_step in range(self.output_sequence_length):
					# Compute Euclidean error for each prediction step within each tbp step
					idx = prediction_step * self.output_pred_state_dim * self.n_mixtures
					idx_x = prediction_step * self.output_state_dim
					idx_y = idx_x + 1

					# Calculate loss for the current ped
					lossfunc= tf.sqrt(tf.add(tf.square(tf.subtract(pred[:, idx_x], out[:, idx_x])),  # error in x
						               tf.square(tf.subtract(pred[:, idx_y], out[:, idx_y]))))  # error in y

					prediction_loss_list.append(lossfunc)  # sum over number of batches

				loss_list.append(tf.reduce_mean(prediction_loss_list)) # mean over prediction horizon

			# Get all trainable variables
			tvars = tf.trainable_variables()

			# L2 loss
			l2 = self.lambda_ * sum(tf.nn.l2_loss(tvar) for tvar in tvars)

			# Reduce mean in all dimensions
			self.total_loss = tf.reduce_mean(loss_list, axis=0) #+ l2 # mean over tbp steps

			# Parameter optimization

			# Exponential / stepwise decrease of learning rate
			global_step = tf.Variable(0, trainable=False)
			self.learning_rate = tf.train.exponential_decay(args.learning_rate_init, global_step, decay_steps=5000,
			                                                decay_rate=0.975, staircase=True)
			self.beta = (tf.tanh((tf.to_float(self.step) - 3500) / 1000) + 1) / 2
			# Optimizer specification
			#self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
			self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate)
			# Compute gradients
			train_vars = tf.trainable_variables()
			self.gradients = tf.gradients(self.total_loss, train_vars)

			# Clip gradients
			self.grads, _ = tf.clip_by_global_norm(self.gradients, self.grads_clip)

			# Apply gradients to update model
			self.update = self.optimizer.apply_gradients(zip(self.grads, train_vars), global_step=global_step)
			# self.update = self.optimizer.minimize(self.total_loss, global_step=global_step)
			# Tensorboard summary

			## Loss
			tf.summary.scalar('loss', self.total_loss)
			tf.summary.scalar('learning_rate', self.learning_rate)
			self.summary = tf.summary.merge_all()
			self.summary_writer = tf.summary.FileWriter(logdir=os.path.join(args.log_dir, 'train'), graph=tf.Session().graph)

			# Saver

			# Save / restore all variables
			self.full_saver = tf.train.Saver()

			# Only save / restore the parameters responsible for occupancy grid processing
			self.convnet_saver = tf.train.Saver(var_list={'conv1_weights': self.conv1_weights,
			                                              'conv1_biases': self.conv1_biases,
			                                              'conv2_weights': self.conv2_weights,
			                                              'conv2_biases': self.conv2_biases,
			                                              'conv3_weights:': self.conv3_weights,
			                                              'conv3_biases': self.conv3_biases,
			                                              'fc_grid_weights': self.fc_grid_weights,
			                                              'fc_grid_biases': self.fc_grid_biases,
			                                              })
			# Retrieve just the LSTM variables.
			lstm_variables = [v for v in tf.all_variables()
			                  if "encoder" in v.name]
			self.encoder_saver = tf.train.Saver(var_list=lstm_variables)

	def feed_dic(self, batch_x, batch_grid, batch_ped_grid, step, batch_y):
		return {self.input_state_placeholder: batch_x[:, :, 2:],
		        self.input_grid_placeholder: batch_grid,
		        self.input_ped_grid_placeholder: batch_ped_grid,
		        self.step: step,
		        self.output_placeholder: batch_y,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
		        self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
						self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat
		        }

	def feed_test_dic(self, batch_x, batch_grid, batch_ped_grid,batch_target):
		return {self.input_state_placeholder: batch_x,
		        self.input_grid_placeholder: batch_grid,
		        self.input_ped_grid_placeholder: batch_ped_grid,
		        self.output_placeholder: batch_target,
		        self.cell_state: self.cell_state_current,
		        self.hidden_state: self.hidden_state_current,
		        self.cell_state_lstm_grid: self.cell_state_current_lstm_grid,
		        self.hidden_state_lstm_grid: self.hidden_state_current_lstm_grid,
		        self.cell_state_lstm_ped: self.cell_state_current_lstm_ped,
		        self.hidden_state_lstm_ped: self.hidden_state_current_lstm_ped,
		        self.cell_state_lstm_concat: self.cell_state_current_lstm_concat,
		        self.hidden_state_lstm_concat: self.hidden_state_current_lstm_concat,
		        }

	def run(self, sess, feed_dict_train):
		_, batch_loss, _current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, _summary_str, lr, beta = sess.run([self.update,
		                                                      self.total_loss,
		                                                      self.current_state,
		                                                      self.current_state_lstm_grid,
		                                                      self.current_state_lstm_ped,
		                                                      self.current_state_lstm_concat,
		                                                      self.prediction,
		                                                      self.summary,
		                                                      self.learning_rate,
		                                                      self.beta],
		                                                     feed_dict=feed_dict_train)
		self.cell_state_current, self.hidden_state_current = _current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return batch_loss, _model_prediction, _summary_str, lr, beta

	def predict(self, sess, feed_dict_train,update_state=True):
		_current_state, _current_state_lstm_grid, _current_state_lstm_ped, \
		_current_state_lstm_concat, \
		_model_prediction, loss = sess.run([self.current_state,
		                              self.current_state_lstm_grid,
		                              self.current_state_lstm_ped,
		                              self.current_state_lstm_concat,
		                              self.prediction,
		                              self.total_loss],
		                             feed_dict=feed_dict_train)
		if update_state:
			self.cell_state_current, self.hidden_state_current = _current_state
			self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = _current_state_lstm_grid
			self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = _current_state_lstm_ped
			self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = _current_state_lstm_concat

		return _model_prediction, loss

	def reset_cells(self,sequence_reset):
		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		if np.any(sequence_reset):
			for sequence_idx in range(sequence_reset.shape[0]):
				if sequence_reset[sequence_idx] == 1:
					# Hidden state of specific batch entry will be initialized with zeros if there was a jump in sequence
					self.cell_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.hidden_state_current[sequence_idx, :] = np.zeros([1, self.rnn_state_size])
					self.cell_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.hidden_state_current_lstm_grid[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_grid])
					self.cell_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.hidden_state_current_lstm_ped[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_ped])
					self.cell_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])
					self.hidden_state_current_lstm_concat[sequence_idx, :] = np.zeros([1, self.rnn_state_size_lstm_concat])

	def update_states(self, current_state, current_state_lstm_grid, current_state_lstm_ped, current_state_lstm_concat):
		# Store the hidden states to re-initialize the in the next training iteration if necessary
		self.cell_state_current, self.hidden_state_current = current_state
		self.cell_state_current_lstm_grid, self.hidden_state_current_lstm_grid = current_state_lstm_grid
		self.cell_state_current_lstm_ped, self.hidden_state_current_lstm_ped = current_state_lstm_ped
		self.cell_state_current_lstm_concat, self.hidden_state_current_lstm_concat = current_state_lstm_concat

	def warmstart_model(self, args, sess):
		# Restore whole model
		print('Loading session from "{}"'.format(args.model_path))
		ckpt = tf.train.get_checkpoint_state(args.model_path)
		print('Restoring model {}'.format(ckpt.model_checkpoint_path))
		self.full_saver.restore(sess, ckpt.model_checkpoint_path)

	def warmstart_convnet(self, args, sess):
		# Restore convnet parameters from pretrained ones (autoencoder encoding part)
		print('Loading convnet parameters from "{}"'.format(args.pretrained_convnet_path))
		ckpt_conv = tf.train.get_checkpoint_state(args.pretrained_convnet_path)
		if ckpt_conv == None:
			print("Error. Convnet not finded...")
			exit()
		else:
			print('Restoring convnet {}'.format(ckpt_conv.model_checkpoint_path))
			self.convnet_saver.restore(sess, ckpt_conv.model_checkpoint_path)

	def process_grid(self, occ_grid_series, trainable):
		"""
		Process occupancy grid with series of convolutional and fc layers.
		input: occupancy grid series (shape: tbpl * [batch_size, grid_dim_x, grid_dim_y])
		output: convolutional feature vector (shape: list of tbpl elements with [batch_size, feature_vector_size] each)
		"""

		conv_feature_vectors = []
		print("Length occ grid series: {}".format(len(occ_grid_series)))
		for idx, grid_batch in enumerate(occ_grid_series):
			grid_batch = tf.expand_dims(input=grid_batch, axis=3)
			conv1 = self.conv_layer(input=grid_batch, weights=self.conv1_weights, biases=self.conv1_biases,
			                        conv_stride_length=self.conv1_stride_length, name="conv1_grid")
			conv2 = self.conv_layer(input=conv1, weights=self.conv2_weights, biases=self.conv2_biases,
			                        conv_stride_length=self.conv2_stride_length, name="conv2_grid")
			conv3 = self.conv_layer(input=conv2, weights=self.conv3_weights, biases=self.conv3_biases,
			                        conv_stride_length=self.conv3_stride_length, name="conv3_grid")
			conv_grid_size = 8
			conv5_flat = tf.reshape(conv3, [-1, conv_grid_size * conv_grid_size * self.conv3_number_filters])
			fc_final = self.fc_layer(input=conv5_flat, weights=self.fc_grid_weights, biases=self.fc_grid_biases,
			                         use_activation=True, name="fc_grid")
			if idx == 0:
				tf.summary.histogram("fc_activations", fc_final)

			# Flatten to obtain feature vector
			conv_features = tf.contrib.layers.flatten(fc_final)
			conv_feature_vectors.append(conv_features)

		return conv_feature_vectors

	def process_pedestrian_grid(self, ped_grid_series):
		"""
		Process pedestrian grid with a FC layer
		"""
		fc_feature_vector = []
		for idx, grid_batch in enumerate(ped_grid_series):
			grid_batch = tf.contrib.layers.flatten(grid_batch)
			fc_feature_vector.append(
				self.fc_layer(grid_batch, weights=self.W_pedestrian_grid, biases=self.b_pedestrian_grid, use_activation=False,
				              name="fc_ped"))
			if idx == 0:
				tf.summary.histogram("fc_ped_activations", fc_feature_vector)
		#       fc_feature_vector.append(grid_batch)
		return fc_feature_vector

	def get_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		                      regularizer=tf.contrib.layers.l2_regularizer(0.01), trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.get_variable(name=name, shape=shape, initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1),
		                      trainable=trainable)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_weight_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get weight variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def get_bayesian_bias_variable(self, shape, name, trainable=True, regularizer=None, summary=True):
		"""
		Get bias variable with specific initializer.
		"""
		var = tf.Variable(tf.truncated_normal(shape, stddev=0.1), dtype=tf.float32, name=name)
		if summary:
			tf.summary.histogram(name, var)
		return var

	def conv_layer(self, input, weights, biases, conv_stride_length=1, padding="SAME", name="conv", summary=False):
		"""
		Convolutional layer including a ReLU activation but excluding pooling.
		"""
		conv = tf.nn.conv2d(input, filter=weights, strides=[1, conv_stride_length, conv_stride_length, 1], padding=padding,
		                    name=name)
		activations = tf.nn.relu(conv + biases)
		if summary:
			tf.summary.histogram(name, activations)
		return activations

	def fc_layer(self, input, weights, biases, use_activation=False, name="fc", summary=False):
		"""
		Fully connected layer with given weights and biases.
		Activation and summary can be activated with the arguments.
		"""
		affine_result = tf.matmul(input, weights) + biases
		if use_activation:
			activations = tf.nn.sigmoid(affine_result)
		else:
			activations = affine_result
		if summary:
			tf.summary.histogram(name + "_activations", activations)
		return activations




src/models/modelRNN_v2.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
# import tensorflow_probability as tfp
import numpy as np
import os
import sys

tfk = tf.keras
tfkl = tf.keras.layers
# tfpl = tfp.layers
# tfd = tfp.distributions

class SeqToSeqModel2():#tfk.Model

	def __init__(self, args):
		#used for subclassing
		#super(SeqToSeqModel2, self).__init__()
		#Network parameters
		if args.input_dim is not None:
			self.input_dim = args.input_dim
		if args.latent_dim is not None:
			self.latent_dim = args.latent_dim
		if args.output_dim is not None:
			self.output_dim = args.output_dim
		if args.batch_size is not None:
			self.batch_size = args.batch_size
		self.args = args
		# Define an input sequence and process it.
		self.encoder_inputs = tfkl.Input(shape=(None, self.input_dim))
		#During training we should use the previous hidden_states right ?
		self.encoder_lstm = tfkl.Bidirectional(tfkl.LSTM(self.latent_dim, return_state=True,stateful=args.stateful))

		self.encoder_outputs, self.state_h_f, self.state_c_f, self.state_h_b, self.state_c_b = self.encoder_lstm(self.encoder_inputs)

		# We discard `encoder_outputs` and only keep the states.
		#self.enc_states_f = [self.state_h_f, self.state_c_f]
		#self.enc_states_b = [self.state_h_b, self.state_c_b]

		self.enc_state_h = tfkl.Concatenate()([self.state_h_f, self.state_h_b])
		self.enc_state_c = tfkl.Concatenate()([self.state_c_f, self.state_c_b])

		self.enc_states = [self.enc_state_h, self.enc_state_c]

		# Set up the decoder, using `encoder_states` as initial state.
		self.dec_inputs = tfkl.Input(shape=(None, self.output_dim))

		self.decoder_lstm = tfkl.LSTM(self.latent_dim * 2, return_sequences=True, return_state=True)

		self.decoder_o, _, _ = self.decoder_lstm(self.dec_inputs,initial_state=self.enc_states)

		self.decoder_fcl = tfkl.Dense(self.latent_dim,activation=tfk.activations.relu)

		self.decoder_linear = tfkl.Dense(self.output_dim)

		self.dec_outputs = self.decoder_fcl(self.decoder_o)

		self.dec_outputs = self.decoder_linear(self.dec_outputs)

		# Define the model that will turn
		# `encoder_input_data` & `decoder_input_data` into `decoder_target_data`
		self.model = tfk.Model([self.encoder_inputs, self.dec_inputs], self.dec_outputs)

		self.encoder_state_input_h = tfkl.Input(shape=(self.latent_dim,))
		self.encoder_state_input_c = tfkl.Input(shape=(self.latent_dim,))
		self.encoder_states_inputs = [self.encoder_state_input_h, self.encoder_state_input_c]

		self.encoder_model = tfk.Model([self.encoder_inputs], self.enc_states)

		self.decoder_state_input_h = tfkl.Input(shape=(self.latent_dim * 2,))
		self.decoder_state_input_c = tfkl.Input(shape=(self.latent_dim * 2,))
		self.decoder_states_inputs = [self.decoder_state_input_h, self.decoder_state_input_c]

		self.decoder_o, self.dec_state_h, self.dec_state_c = self.decoder_lstm(self.dec_inputs,
		                                                                       initial_state=self.decoder_states_inputs)

		self.dec_states = [self.dec_state_h, self.dec_state_c]

		self.dec_out = self.decoder_fcl(self.decoder_o)

		self.dec_out = self.decoder_linear(self.dec_out)

		self.decoder_model = tfk.Model([self.dec_inputs] + self.decoder_states_inputs,[self.dec_out]+ self.dec_states)

	def get_encoder_zero_initial_state(self, inputs):
		return [tf.zeros((self.batch_size, self.latent_dim)), tf.zeros((self.batch_size, self.latent_dim))]

	def get_initial_state(self, inputs):
		return self.initial_state

	"""Used for subclassing
	def __call__(self, inputs, training,states=None):
		if states is None:
			self.encoder_lstm.get_initial_state = self.get_encoder_zero_initial_state
		else:
			self.initial_state = states
			self.encoder_lstm.get_initial_state = self.get_initial_state

		self.encoder_outputs, self.state_h, self.state_c = self.encoder_lstm(inputs[0])

		self.enc_states = [self.state_h, self.state_c]

		self.decoder_outputs, _, _ = self.decoder_lstm(inputs[1], initial_state=self.enc_states)

		self.decoder_dense_o = self.decoder_dense(self.decoder_outputs)

		return self.decoder_dense_o
	"""
	def decode_sequence(self,input_seq):
		# Populate the first character of target sequence with the start character.
		target_seq = np.zeros((1, 1, self.args.output_dim))
		predictions = []
		#states_value = [np.zeros((1, self.args.latent_dim )), np.zeros((1, self.args.latent_dim ))]
		for i in range(len(input_seq[0])):
			# Encode the input as state vectors.
			in_enc = [input_seq[0][i]]
			in_dec = [input_seq[1][i]]
			target_seq[0, 0, :] = in_enc[0][-1, 6:8]
			states_value = self.encoder_model.predict([in_enc],batch_size=1)

			decoded_sentence = np.zeros((self.args.prediction_horizon, self.args.output_dim))
			step = 0
			stop_condition = False
			while not stop_condition:
				# Exit condition: either hit max length
				# or find stop character.
				if step >= self.args.prediction_horizon - 1:
					stop_condition = True

				output_tokens, h, c = self.decoder_model.predict(
					[target_seq] + states_value)

				# Sample a token
				sampled_output = output_tokens[0, -1, :]

				decoded_sentence[step, :] = sampled_output

				# Update the target sequence (of length 1).
				target_seq = output_tokens

				# Update states
				states_value = [h, c]

				step += 1
			predictions.append(decoded_sentence.copy())

		return predictions
src/models/seqBiModel.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
# import tensorflow_probability as tfp
import numpy as np
import os
import sys

tfk = tf.keras
tfkl = tf.keras.layers
# tfpl = tfp.layers
# tfd = tfp.distributions

class SeqToSeqModel2():#tfk.Model

	def __init__(self, args):
		#used for subclassing
		#super(SeqToSeqModel2, self).__init__()
		#Network parameters
		if args.input_dim is not None:
			self.input_dim = args.input_dim
		if args.latent_dim is not None:
			self.latent_dim = args.latent_dim
		if args.output_dim is not None:
			self.output_dim = args.output_dim
		if args.batch_size is not None:
			self.batch_size = args.batch_size
		self.args = args
		# Define an input sequence and process it.
		self.encoder_inputs = tfkl.Input(shape=(None, self.input_dim))
		#During training we should use the previous hidden_states right ?
		self.encoder_lstm = tfkl.LSTM(self.latent_dim, return_state=True,stateful=args.stateful)

		self.encoder_outputs, self.state_h, self.state_c = self.encoder_lstm(self.encoder_inputs)

		# We discard `encoder_outputs` and only keep the states.
		self.enc_states = [self.state_h, self.state_c]

		# Set up the decoder, using `encoder_states` as initial state.
		self.dec_inputs = tfkl.Input(shape=(None, self.output_dim))

		self.decoder_lstm = tfkl.LSTM(self.latent_dim, return_sequences=True, return_state=True)

		self.decoder_o, _, _ = self.decoder_lstm(self.dec_inputs,initial_state=self.enc_states)

		self.decoder_fcl = tfkl.Dense(self.latent_dim,activation=tfk.activations.relu)

		self.decoder_linear = tfkl.Dense(self.output_dim)

		self.dec_outputs = self.decoder_fcl(self.decoder_o)

		self.dec_outputs = self.decoder_linear(self.dec_outputs)

		# Define the model that will turn
		# `encoder_input_data` & `decoder_input_data` into `decoder_target_data`
		self.model = tfk.Model([self.encoder_inputs, self.dec_inputs], self.dec_outputs)

		self.encoder_state_input_h = tfkl.Input(shape=(self.latent_dim,))
		self.encoder_state_input_c = tfkl.Input(shape=(self.latent_dim,))
		self.encoder_states_inputs = [self.encoder_state_input_h, self.encoder_state_input_c]

		self.encoder_model = tfk.Model([self.encoder_inputs]+self.encoder_states_inputs, self.enc_states)

		# This model is used only for inference

		self.decoder_state_input_h = tfkl.Input(shape=(self.latent_dim,))
		self.decoder_state_input_c = tfkl.Input(shape=(self.latent_dim,))
		self.decoder_states_inputs = [self.decoder_state_input_h, self.decoder_state_input_c]

		self.decoder_o, self.dec_state_h, self.dec_state_c = self.decoder_lstm(self.dec_inputs,
		                                                                       initial_state=self.decoder_states_inputs)

		self.dec_states = [self.dec_state_h, self.dec_state_c]

		self.dec_out = self.decoder_fcl(self.decoder_o)

		self.dec_out = self.decoder_linear(self.dec_out)

		self.decoder_model = tfk.Model([self.dec_inputs] + self.decoder_states_inputs,[self.dec_out]+ self.dec_states)

	def get_encoder_zero_initial_state(self, inputs):
		return [tf.zeros((self.batch_size, self.latent_dim)), tf.zeros((self.batch_size, self.latent_dim))]

	def get_initial_state(self, inputs):
		return self.initial_state

	"""Used for subclassing
	def __call__(self, inputs, training,states=None):
		if states is None:
			self.encoder_lstm.get_initial_state = self.get_encoder_zero_initial_state
		else:
			self.initial_state = states
			self.encoder_lstm.get_initial_state = self.get_initial_state

		self.encoder_outputs, self.state_h, self.state_c = self.encoder_lstm(inputs[0])

		self.enc_states = [self.state_h, self.state_c]

		self.decoder_outputs, _, _ = self.decoder_lstm(inputs[1], initial_state=self.enc_states)

		self.decoder_dense_o = self.decoder_dense(self.decoder_outputs)

		return self.decoder_dense_o
	"""
	def decode_sequence(self,input_seq):
		# Populate the first character of target sequence with the start character.
		target_seq = np.zeros((1, 1, self.args.output_dim))
		predictions = []
		states_value = [np.zeros((1, self.args.latent_dim)), np.zeros((1, self.args.latent_dim))]
		for i in range(len(input_seq[0])):
			# Encode the input as state vectors.
			in_enc = [input_seq[0][i]]
			in_dec = [input_seq[1][i]]
			target_seq[0, 0, :] = in_enc[0][-1, 6:8]
			states_value = self.encoder_model.predict([in_enc]+states_value,batch_size=1)

			decoded_sentence = np.zeros((self.args.prediction_horizon, self.args.output_dim))
			step = 0
			stop_condition = False
			while not stop_condition:
				# Exit condition: either hit max length
				# or find stop character.
				if step >= self.args.prediction_horizon - 1:
					stop_condition = True

				output_tokens, h, c = self.decoder_model.predict(
					[target_seq] + states_value)

				# Sample a token
				sampled_output = output_tokens[0, -1, :]

				decoded_sentence[step, :] = sampled_output

				# Update the target sequence (of length 1).
				target_seq = output_tokens

				# Update states
				states_value = [h, c]

				step += 1
			predictions.append(decoded_sentence.copy())

		return predictions
src/models/seqModel.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf
#import tensorflow_probability as tfp
import numpy as np
import os
import sys

sys.path.append('./autoencoder_grid')
tfk = tf.keras
tfkl = tf.keras.layers
#tfpl = tfp.layers
#tfd = tfp.distributions


class RNNModel():

	def __init__(self, args, batch_size=None):
		# Input- / Output dimensions
		if args.input_dim is not None:
			self.input_dim = args.input_dim
		if args.latent_dim is not None:
			self.latent_dim = args.latent_dim
		if args.output_dim is not None:
			self.output_dim = args.output_dim
		if args.batch_size is not None:
			self.batch_size = args.batch_size
		self.args = args

		# Define an input sequence and process it.
		self.encoder_inputs = tfkl.Input(shape=(None,self.input_dim*self.args.prev_horizon))

		# Encoder state LSTM
		self.enc_state_lstm = tfkl.LSTM(self.latent_dim,
		                                  input_shape=[None, self.input_dim],
		                                  dropout=self.args.dropout, name='enc_lstm', return_sequences=True,
		                                  return_state=True, stateful=args.stateful,
		                                  kernel_regularizer=tfk.regularizers.l2(l=self.args.regularization_weight))

		self.encoder_outputs, self.state_h, self.state_c = self.enc_state_lstm(self.encoder_inputs)

		# We discard `encoder_outputs` and only keep the states.
		self.enc_states = [self.state_h, self.state_c]

		# Set up the decoder, using `encoder_states` as initial state.
		self.dec_inputs = tfkl.Input(shape=(None, self.args.output_dim))

		self.decoder_lstm = tfkl.LSTM(self.latent_dim,
		                             input_shape=[None, self.latent_dim],
		                             dropout=self.args.dropout, name='dec_lstm',
		                             return_sequences=True, return_state=True, stateful=args.stateful,
		                             kernel_regularizer=tfk.regularizers.l2(l=self.args.regularization_weight))

		self.decoder_o, _, _ = self.decoder_lstm(self.dec_inputs, initial_state=self.enc_states)

		self.fcl = tfkl.TimeDistributed(tfkl.Dense(self.latent_dim, activation='relu', name='dec_fc',
		                                           kernel_regularizer=tfk.regularizers.l2(l=self.args.regularization_weight)))
		self.fco = tfkl.TimeDistributed(tfkl.Dense(self.output_dim*self.args.prediction_horizon, activation='linear', name='out_fc',
		                                           kernel_regularizer=tfk.regularizers.l2(l=self.args.regularization_weight)))

		self.dec_outputs = self.fcl(self.decoder_o)

		self.dec_outputs = self.fco(self.dec_outputs)

		# Define the model that will turn
		# `encoder_input_data` & `decoder_input_data` into `decoder_target_data`
		self.model = tfk.Model([self.encoder_inputs, self.dec_inputs], self.dec_outputs)

		self.encoder_state_input_h = tfkl.Input(shape=(self.latent_dim,))
		self.encoder_state_input_c = tfkl.Input(shape=(self.latent_dim,))
		self.encoder_states_inputs = [self.encoder_state_input_h, self.encoder_state_input_c]

		self.encoder_model = tfk.Model([self.encoder_inputs] + self.encoder_states_inputs, self.enc_states)

		self.decoder_state_input_h = tfkl.Input(shape=(self.latent_dim,))
		self.decoder_state_input_c = tfkl.Input(shape=(self.latent_dim,))
		self.decoder_states_inputs = [self.decoder_state_input_h, self.decoder_state_input_c]

		self.decoder_o, self.dec_state_h, self.dec_state_c = self.decoder_lstm(self.dec_inputs,
		                                                                       initial_state=self.decoder_states_inputs)

		self.dec_states = [self.dec_state_h, self.dec_state_c]

		self.dec_out = self.fcl(self.decoder_o)

		self.dec_out = self.fco(self.dec_out)

		self.decoder_model = tfk.Model([self.dec_inputs] + self.decoder_states_inputs, [self.dec_out] + self.dec_states)

	def decode_sequence(self,input_seq):
		# Populate the first character of target sequence with the start character.
		target_seq = np.zeros((1, 1, self.args.output_dim))
		predictions = []
		states_value = [np.zeros((1, self.args.latent_dim)), np.zeros((1, self.args.latent_dim))]
		for i in range(len(input_seq[0])):
			# Encode the input as state vectors.
			in_enc = [input_seq[0][i]]

			target_seq[0, 0, :] = input_seq[1][i]
			states_value = self.encoder_model.predict([in_enc]+states_value,batch_size=1)

			decoded_sentence = np.zeros((1, self.args.output_dim*self.args.prediction_horizon))


			output_tokens, h, c = self.decoder_model.predict(
					[target_seq] + states_value)

			# Sample a token
			sampled_output = output_tokens[0, -1, :]

			decoded_sentence[0, :] = sampled_output

			# Update states
			states_value = [h, c]

			predictions.append(decoded_sentence.copy())

		return predictions
src/models/seqRNNModel.py
eliatrevisan/IA-MPPI-LBM
import tensorflow as tf

def linear(input_, output_size,id=0, scope=None, stddev=0.02, bias_start=0.0, with_w=False):
    shape = input_.get_shape().as_list()

    with tf.variable_scope(scope or "Linear",reuse=tf.AUTO_REUSE) as scope:
        matrix = tf.get_variable("Matrix"+str(id), [shape[1], output_size], tf.float32,
                                 tf.random_normal_initializer(stddev=stddev))
        bias = tf.get_variable("bias"+str(id), [output_size],
            initializer=tf.constant_initializer(bias_start))
        if with_w:
            return tf.matmul(input_, matrix) + bias, matrix, bias
        else:
            return tf.matmul(input_, matrix) + bias
src/models/tf_utils.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
import json
import importlib
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *

import pickle as pkl
import time
from copy import deepcopy
from multiprocessing.pool import ThreadPool
from progressbar import *

home = os.path.expanduser('~')

class Map:
    size = [1715, 881]
    data = pl.imread(home + '/I-LSTM/data/simulation/roboat/map.png')
    origin = [-78, -40, 0.0]
    resolution = 0.081

def plot_map(map, ax):
    ax.imshow(map.data,
               extent = (map.origin[0], map.origin[0] + map.size[0] * map.resolution,
                         map.origin[1], map.origin[1] + map.size[1] * map.resolution),
               cmap='gray_r')
    ax.set_facecolor('#262626')
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    ax.yaxis.set_label_coords(-0.08, .5)
    ax.xaxis.set_label_coords(0.5, -0.09)

#os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

pretrained_convnet_path = "../trained_models/autoencoder_with_ped"

#data_path = '../data/2_agents_swap/trajs/'
#scenario = 'GA3C-CADRL-10-py27'
#data_path = '../data/cyberzoo_experiments/'
#scenario = 'all_trajectories'
exp_num = 2
data_path = '../data/'
scenario = 'simulation/roboat'

# Hyperparameters
n_epochs = 2

batch_size = 1
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 10
prediction_horizon = 10
prev_horizon = 0

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 32
rnn_state_ped_size = 16
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 256
x_dim = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 1
tensorboard_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 3

print_freq = 2000
save_freq = 2000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = True
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
submap_resolution = 0.081
submap_width = 4.86
submap_height = 4.86
diversity_update = False
predict_positions = False
warm_start_convnet = True

train_set = 1.0

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default= "RNN")
	parser.add_argument('--model_path',
	                    help='Path to directory to save the model (default=""../trained_models/"+model_name").',
	                    type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
	                    help='Path to directory that comprises the pre-trained convnet model (default=" ").',
	                    type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
	                    help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
	                    type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
	                    type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
	                    default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
	                    default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
	                    default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
	                    default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
	                    default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
	                    default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
	                    default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
	                    default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
	                    default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
	                    default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
	                    default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
	                    default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
	                    help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
	                    type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
	                    type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
	                    type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
	                    help='Number of fully connected layer units after LSTM layer (default=64).',
	                    type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
	                    default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
	                    default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
	                    default=output_dim)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
	                    default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
	                    default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
	                    default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
	                    default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
	                    default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
	                    default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
	                    default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the netwprk.', type=int,
	                    default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=int, default=0)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
	                    default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
	                    type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
	                    default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
	                    default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
	                    default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
	                    default=0)
	parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=submap_width)
	parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=submap_height)
	parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
	                    default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--relative_info', help='Use relative info for other agents.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
	                    default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
	                    default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
	                    default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
	                    default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	args = parser.parse_args()

	return args


args = parse_args()

if args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf
tfk = tf.keras

args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)
args.log_dir = args.model_path + '/log'
if not os.path.exists(args.log_dir):
	os.makedirs(args.log_dir)
args.dataset = '/' + args.scenario + '.pkl'
model_parameters = {"args": args}
print(args)
# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()
with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

data_prep = dhlstm.DataHandlerLSTM(args)
# Only used to create a map from png
map_args = {"file_name": 'map.png',
	            "resolution": 0.1,
	            "map_size": np.array([30., 30.]), }

data_prep.processData(**map_args)

fig_animate = pl.figure('Dataset Trajectories')
fig_animate.tight_layout()
fig_width = 12  # width in inches
fig_height = 10  # height in inches
fig_size = [fig_width, fig_height]
fontsize = 26

params = {'backend': 'ps',
          'axes.labelsize': fontsize,
          'font.size': fontsize,
          'xtick.labelsize': fontsize,
          'ytick.labelsize': fontsize,
          'figure.figsize': fig_size}

plt_colors = [0.8500, 0.3250, 0.0980]

colors = ['r', 'g', 'b','y','m','c','k']
ax_pos = pl.subplot()

ax_pos.set_xlabel('x [m]')
ax_pos.set_ylabel('y [m]')
ax_pos.axis('on')
ax_pos.set_aspect('equal')

new_epoch = False

while not new_epoch:
	plt_colors = np.random.uniform(0,1,[3])

	batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch = data_prep.getBatch()
	mij = np.min(batch_ped_grid,axis=2)
	#print(np.min(mij,axis=1))
	if np.min(mij, axis=1) < 1.0:
		#ax_pos.plot(batch_pos_target[0,0,0::2], batch_pos_target[0,0,1::2], color=plt_colors)
		ax_pos.plot(batch_pos[0,:,0], batch_pos[0,:,1], color=plt_colors)
map = Map()
plot_map(map, ax_pos)
pl.show(block=True)

print(batch_pos[0,:,0])
print(batch_pos[0,:,1])


fig_animate.savefig(args.data_path + args.scenario + "/trajectories.jpg")

src/plot_dataset.py
eliatrevisan/IA-MPPI-LBM
# -*- coding: utf-8 -*-

from math import floor, exp, sqrt
import numpy as np
from src.simulators import social_forces_utils as utils
import sys
from os.path import dirname, abspath
import pylab as pl
import random
import csv
import time

def load_map(data_path):
    try:
        file_path = data_path + '/map.npy'
        if sys.version_info[0] < 3:
            grid_map = np.load(file_path,allow_pickle=True)[()]
        else:
            grid_map = np.load(file_path,encoding='latin1',allow_pickle=True)[()]
        print("Map loaded")
        return grid_map
    except IOError:
        print("ERROR: Cannot load map.")


class ForceModel(object):
    """
    This class is used to simulate and predict social force model.
    """

    def __init__(self, prediction_steps, occ_grid,ped_num,dt=0.1):
        self.desired_speed = 1.0
        self.rel_time = 0.5
        self.delta_t = 0.5
        self.tol = 0.3
        self.ped_num = ped_num
        self.prediction_steps = prediction_steps
        self.occ_grid = occ_grid
        self.dt = dt  # Simulation step in seconds
        try:
            self.map_center = self.occ_grid['Size']/2.0
        except:
            self.map_center = self.occ_grid.map_size
        self.corridor = True

    def goal_force(self, pose, current_speed, goal):
        """
        This function calculates the attractive force towards the goal
        """
        direction = utils.calc_direction(current_pose=pose, desired_pose=goal)
        speed_mod = utils.calc_mod(current_speed)

        force = np.zeros((2))
        force = (direction*self.desired_speed - current_speed)/self.rel_time
        return force

    def repulsive_obj_forces(self, pose, speed, goal):
        """
        This function uses the occupancy grid to calculate the repulsive forces
        of obstacles in a given pose.
        It is not done only once and stored so we can avoid the quantization error
        given by the discretization of the map.
        """
        A = 4.3
        B = 1.07
        l = 0.7

        # Rescale position
        position = (np.array(pose) + self.map_center)/self.occ_grid['Resolution']

        # Get indexes for map
        position_idx = np.ndarray.astype(position, 'int')

        closest_obj = np.array((self.occ_grid['Closest X'][tuple(position_idx)],
                                self.occ_grid['Closest Y'][tuple(position_idx)]))

        rel_distance = pose - closest_obj
        step = - np.array(speed) * self.delta_t
        r_s = rel_distance - step

        mod_rd = utils.calc_mod(rel_distance)
        mod_step = utils.calc_mod(step)
        mod_rs = utils.calc_mod(r_s)

        b = sqrt( pow(mod_rd + mod_rs, 2) - pow(mod_step, 2) )/2

        force = np.zeros((2))
        # if mod_rd != 0 and mod_rs != 0 and mod_step != 0 and b != 0:
        g = A * exp(-b/B) * (mod_rd + mod_rs)/(2*b)
        g = g*(rel_distance/mod_rd + r_s/mod_rs)/2

        force = g


        return force

    def repulsive_ped_force(self, pose, speed, goal, other_ped):
        """
        This function calculates the repulsive forces exerted by other pedestrian.
        """
        A = 4.3
        B = 1.07
        l = 0.7
        # other_ped is in the format [[pose_x, pose_y], [vel_x, vel_y], [goal_x, goal_y]]
        r_ij = np.array(pose) - np.array(other_ped[0])
        s_ij = (np.array(other_ped[1]) - np.array(speed)) * self.delta_t
        r_s_ij = r_ij - s_ij

        mod_r_ij = utils.calc_mod(r_ij)
        mod_r_s_ij = utils.calc_mod(r_s_ij)
        mod_s_ij = utils.calc_mod(s_ij)

        b = sqrt(pow(mod_r_ij + mod_r_s_ij, 2) - pow(mod_s_ij, 2))/2
        force = np.zeros((2))
        if (mod_r_ij != 0) and (mod_r_s_ij != 0) and (mod_s_ij != 0) and (b != 0):
            g = A * exp(-b/B) * (mod_r_ij + mod_r_s_ij) / (2*b)
            g = g * (r_ij/mod_r_ij + r_s_ij/mod_r_s_ij)/2

            e_t_i = utils.calc_direction(current_pose=pose,
                                         desired_pose=goal)

            cos_phi = (e_t_i[0]*r_ij[0]/mod_r_ij) + (e_t_i[1]*r_ij[1]/mod_r_ij)
            w = l + (1 - l)*(1 + cos_phi)/2

            force = g * w

        return force

    def predict(self, ped_id, initial_pose, current_speed, goal, dt):
        """
        This function predicts the next positions and speed of the pedestrian
        using the social force model.
        """
        # data must be in the shape [current + predicted_time_step, ped_index, pose_speed, XY_dims]
        current_pose = initial_pose
        pred_velocities = np.zeros((self.prediction_steps, 2))
        pred_poses = np.zeros((self.prediction_steps, 2))

        for i in range(self.prediction_steps):
            force = self.goal_force(current_pose, current_speed, goal)
            force = force + self.repulsive_obj_forces(pose=current_pose,
                                                      speed=current_speed,
                                                      goal=goal)

            for pp in range(self.ped_num):
                if pp != ped_id:
                    ped_data = [self.current_poses[pp], self.current_speed[pp], self.goals[pp]]
                    force = force + self.repulsive_ped_force(pose=initial_pose,
                                                                         speed=current_speed,
                                                                         goal=goal,
                                                                         other_ped=ped_data)

            current_speed = current_speed + dt * force # Integrate speed
            mod_speed = utils.calc_mod(current_speed)
            max_speed = 1.3 * self.desired_speed
            if mod_speed >= max_speed:
                current_speed = current_speed * max_speed/mod_speed
            current_pose = current_pose + dt * current_speed # Integrate pose

            pred_velocities[i] = current_speed
            pred_poses[i] = current_pose
        return pred_poses, pred_velocities

    def set_initial_pose(self,pose,speed,goals):
        self.initial_pose = pose
        self.current_speed = speed
        self.goals = goals
        self.current_poses = pose

    def get_initial_poses(self, ped_num):
        """
        This functions calculates the random initial positions for the
        pedestrians, making sure they do not spawn into obstacles
        """
        index = 0
        poses = np.zeros((ped_num, 2))
        while index < ped_num:
            if (index % 2):
                pose = np.array((random.uniform(-9.9, -9),
                                 random.uniform(-0.5, 0.5)))
            else:
                pose = np.array((random.uniform(9, 9.9),
                                 random.uniform(-0.5, 0.5)))

            position = (pose + self.map_center)/self.occ_grid['Resolution']
            obstacle = self.occ_grid['Map'][int(position[0]), int(position[1])]
            occupied = utils.check_if_occupied(pose, poses[:index+1])
            if obstacle != 1 and not occupied:
                poses[index] = pose
                index = index + 1
        return poses

    def get_goal(self, pose=None, corridor=False):
        done = False

        if corridor:
            if pose[0] >= 0:
                x_range = np.array([-9.9, -9])
                y_range = np.array([-0.5, 0.5])
            elif pose[0] <= 0:
                x_range = np.array([9, 9.9])
                y_range = np.array([-0.5, 0.5])
        else:
            x_range = np.array([-self.occ_grid['Size'][0]/4,
                                self.occ_grid['Size'][0]/4])
            y_range = np.array([-self.occ_grid['Size'][1]/4,
                                self.occ_grid['Size'][1]/4])

        while not done:
            goal = np.array((random.uniform(x_range[0], x_range[1]),
                             random.uniform(y_range[0], y_range[1])))

            # Check that the goal is not into any obstacle
            position = (goal + self.map_center)/self.occ_grid['Resolution']
            obstacle = self.occ_grid['Map'][int(position[0]), int(position[1])]
            position_idx = np.ndarray.astype(position, 'int')
            closest_obj = np.array((self.occ_grid['Closest X'][tuple(position_idx)],
                                    self.occ_grid['Closest Y'][tuple(position_idx)]))
            if obstacle != 1 or utils.calc_mod(closest_obj - goal) >= 0.5:
                done = True
        return goal

    def simulate(self, ped_num, duration, data_path):
        """
        This function is used to simulate SF model in the given map
        for the given time.
        Input:
            ped_num: Number of pedestrians to simulate
            time: [hours, minutes, seconds]
        """
        self.ped_num = ped_num
        self.min_time = duration[0] * 60 + duration[1]
        sec_time = self.min_time * 60 + duration[2]
        #print("Simulation length: {}h {}m {}s".format(duration[0],
        #                                              duration[1],
        #                                              duration[2]))
        dt = 0.1 # Simulation step in seconds
        sim_elapsed_mins = 0

        self.csv_files = []
        self.writers = []
        self.fieldnames = ['Pedestrian', 'Time s', 'Time ns',
                           'Position X', 'Position Y',
                           'Velocity X', 'Velocity Y',
                           'Goal X', 'Goal Y']

        random.seed()
        counts = np.zeros((self.ped_num))
        poses = self.get_initial_poses(self.ped_num)
        speed = np.zeros((self.ped_num, 2))
        goals = np.zeros((self.ped_num, 2))
        forces = np.zeros((self.ped_num, 2))

        for i in range(self.ped_num):
            goals[i] = self.get_goal(poses[i], corridor=self.corridor)
            # Create logger files
            try:
                self.csv_files.append(open('{}/log_ped{}.csv'.format(data_path, i), 'wb'))
            except:
                print("Cannot open log file number {}. Quitting.".format(i))
                return
            self.writers.append(csv.DictWriter(self.csv_files[i], fieldnames=self.fieldnames))
            self.writers[i].writeheader()
        # Create total log file
        try:
            self.csv_files.append(open('{}/total_log.csv'.format(data_path), 'wb'))
        except:
            print("Cannot open total log file. Quitting.".format(i))
            return
        self.total_log_writer = csv.DictWriter(self.csv_files[-1], fieldnames=self.fieldnames)
        self.total_log_writer.writeheader()

        # Now we are ready to start simulate
        t_start = time.time()
        time_steps = int(sec_time/dt)
        for ts in range(time_steps):
            # NB this could be made for all the peds at once but the force functions
            # for now are designed to work with just one ped at a time
            for ped in range(self.ped_num):
                forces[ped] = self.goal_force(pose=poses[ped],
                                              current_speed=speed[ped],
                                              goal=goals[ped])

                forces[ped] = forces[ped] + self.repulsive_obj_forces(pose=poses[ped],
                                                                      speed=speed[ped],
                                                                      goal=goals[ped])
                for pp in range(self.ped_num):
                    if pp != ped:
                        ped_data = [poses[pp], speed[pp], goals[pp]]
                        forces[ped] = forces[ped] + self.repulsive_ped_force(pose=poses[ped],
                                                                             speed=speed[ped],
                                                                             goal=goals[ped],
                                                                             other_ped=ped_data)
                # Add some noise
                forces[ped] = forces[ped]
            # Once we have all the forces, we update velocity for each ped

            speed = speed + dt * (forces + np.random.normal(0, 0.3, (self.ped_num, 2)))

            for ped in range(self.ped_num):

                mod_speed = utils.calc_mod(speed[ped])
                max_speed = 1.3 * self.desired_speed
                if mod_speed >= max_speed:
                    speed[ped] = speed[ped] * max_speed/mod_speed
                new_pose = poses[ped] + dt * speed[ped]

                # Check if we end up into an obstacle
                # Rescale position
                position = (np.array(new_pose) + self.map_center)/self.occ_grid['Resolution']
                # Get indexes for map
                position = np.ndarray.astype(position, 'int')
                # If we would end up into an obst we do not move
                if self.occ_grid['Map'][tuple(position)] == 1:
                    print("Inside obstacle")
                    new_pose = poses[ped]
                    speed[ped] = np.zeros((2))

                # Update poses
                poses[ped] = new_pose
                # Save data
                data_dict = {}
                data_dict['Pedestrian'] = ped
                data_dict['Time s'] = int(ts*dt)
                data_dict['Time ns'] = int((ts*dt * 10**9)%10**9)
                data_dict['Position X'] = new_pose[0]
                data_dict['Position Y'] = new_pose[1]
                data_dict['Velocity X'] = speed[ped][0]
                data_dict['Velocity Y'] = speed[ped][1]
                data_dict['Goal X'] = goals[ped][0]
                data_dict['Goal Y'] = goals[ped][1]
                self.writers[ped].writerow(data_dict)
                self.total_log_writer.writerow(data_dict)
                # Update goal
                if utils.calc_mod(goals[ped] - new_pose) <= self.tol:
                    goals[ped] = self.get_goal(new_pose, corridor=self.corridor)
                    counts[ped] = 0
                # Check if stuck
                if utils.calc_mod(speed[ped]) <= 0.01:
                    counts[ped] = counts[ped] + 1
                    if counts[ped] == 5:
                        goals[ped] = self.get_goal(new_pose, corridor=self.corridor)
                        counts[ped] = 0
                else: # This way we just count when they are consecutive
                    counts[ped] = 0

            # Log passing time
            rt_time = time.time() - t_start
            if (ts*dt/60 - sim_elapsed_mins) >= 1:
                sim_elapsed_mins = int(ts*dt/60)
                fact = ts*dt / rt_time
                rt_time = time.gmtime(rt_time)
                print("Elapsed simulation time since start: {} mins".format(sim_elapsed_mins))
                print("Elapsed real time since start: {}h {}m {}s".format(rt_time.tm_hour,
                                                                          rt_time.tm_min,
                                                                          rt_time.tm_sec))
                print("Sim to RT factor: {}".format(fact))
                print("")

        #
        for i in range(self.ped_num + 1):
            self.csv_files[i].close()

        print("Simulation finished.")
        rt_time = time.gmtime(time.time() - t_start)
        print("Total real time: {}h {}m {}s".format(rt_time.tm_hour,
                                                    rt_time.tm_min,
                                                    rt_time.tm_sec))
    def simulate_ped(self, ped_id,ini_pose,goal,duration,other_agents_pos):
        """
        This function is used to simulate SF model in the given map
        for the given time.
        Input:
            ped_num: Number of pedestrians to simulate
            time: [hours, minutes, seconds]
        """

        self.min_time = duration[0] * 60 + duration[1]
        sec_time = self.min_time * 60 + duration[2]
        print("Simulation length: {}h {}m {}s".format(duration[0],
                                                      duration[1],
                                                      duration[2]))
        dt = 0.3 # Simulation step in seconds
        sim_elapsed_mins = 0

        random.seed()
        counts = np.zeros((self.ped_num))
        poses = np.expand_dims(ini_pose[0:2].copy(),0)
        speed = np.expand_dims(ini_pose[2:].copy(),0)
        goals = np.zeros((self.ped_num, 2))
        forces = np.zeros((self.ped_num, 2))

        for i in range(self.ped_num):
            goals[i] = goal

        # Now we are ready to start simulate
        t_start = time.time()
        time_steps = int(sec_time/dt)
        traj = np.zeros((time_steps,2))
        velocities = np.zeros((time_steps,2))
        velocities[0,:] = speed
        traj[0,:] = ini_pose[0:2]
        for ts in range(time_steps-1):
            # NB this could be made for all the peds at once but the force functions
            # for now are designed to work with just one ped at a time
            ped=0
            forces[ped] = self.goal_force(pose=poses[ped],
                                         current_speed=speed[ped],
                                         goal=goals[ped])

            forces[ped] = forces[ped] + self.repulsive_obj_forces(pose=poses[ped],
                                                                  speed=speed[ped],
                                                                  goal=goals[ped])
            for pp in range(len(other_agents_pos)):
                agent_pos = np.reshape(other_agents_pos[pp],2)
                ped_data = [agent_pos, np.zeros((2)), agent_pos]
                forces[ped] = forces[ped] + self.repulsive_ped_force(pose=poses[ped],
                                                                     speed=speed[ped],
                                                                     goal=goals[ped],
                                                                     other_ped=ped_data)
            # Add some noise
            forces[ped] = forces[ped]
            # Once we have all the forces, we update velocity for each ped

            speed = speed + dt * (forces + np.random.normal(0, 0.01, (self.ped_num, 2)))

            mod_speed = utils.calc_mod(speed[ped])
            max_speed = 1.3 * self.desired_speed
            if mod_speed >= max_speed:
                speed[ped] = speed[ped] * max_speed/mod_speed
            new_pose = poses[ped] + dt * speed[ped]

            # Check if we end up into an obstacle
            # Rescale position
            position = (np.array(new_pose) + self.map_center)/self.occ_grid['Resolution']
            # Get indexes for map
            position = np.ndarray.astype(position, 'int')
            # If we would end up into an obst we do not move
            if self.occ_grid['Map'][tuple(position)] == 1:
                new_pose = poses[ped]
                speed[ped] = np.zeros((2))

            # Update poses
            poses[ped] = new_pose

            # Update goal
            if utils.calc_mod(goals[ped] - new_pose) <= self.tol:
                goals[ped] = self.get_goal(new_pose, corridor=self.corridor)
                counts[ped] = 0
            # Check if stuck
            if utils.calc_mod(speed[ped]) <= 0.01:
                counts[ped] = counts[ped] + 1
                if counts[ped] == 5:
                    goals[ped] = self.get_goal(new_pose, corridor=self.corridor)
                    counts[ped] = 0
            else: # This way we just count when they are consecutive
                counts[ped] = 0

            traj[ts+1,:]=poses[ped]
            velocities[ts+1,:]=speed[ped]
        return traj, velocities
    
    def simulation_step(self,sec_time):
        """
        This function is used to simulate one step of SF model in the given map
        for the given time.
        Input:
            ped_num: Number of pedestrians to simulate
            time: [hours, minutes, seconds]
        """

        forces = np.zeros((self.ped_num, 2))

        # Now we are ready to start simulate
        time_steps = int(sec_time/self.dt)
        for ts in range(time_steps):
            # NB this could be made for all the peds at once but the force functions
            # for now are designed to work with just one ped at a time
            for ped in range(self.ped_num):
                forces[ped] = self.goal_force(pose=self.current_poses[ped],
                                              current_speed=self.current_speed[ped],
                                              goal=self.goals[ped])

                forces[ped] = forces[ped] + self.repulsive_obj_forces(pose=self.current_poses[ped],
                                                                      speed=self.current_speed[ped],
                                                                      goal=self.goals[ped])
                for pp in range(self.ped_num):
                    if pp != ped:
                        ped_data = [self.current_poses[pp], self.current_speed[pp], self.goals[pp]]
                        forces[ped] = forces[ped] + self.repulsive_ped_force(pose=self.current_poses[ped],
                                                                             speed=self.current_speed[ped],
                                                                             goal=self.goals[ped],
                                                                             other_ped=ped_data)
                # Add some noise
                forces[ped] = forces[ped]
            # Once we have all the forces, we update velocity for each ped

            self.current_speed = self.current_speed + self.dt * (forces + np.random.normal(0, 0.1, (self.ped_num, 2)))

            for ped in range(self.ped_num):

                mod_speed = utils.calc_mod(self.current_speed[ped])
                max_speed = 1.3 * self.desired_speed
                if mod_speed >= max_speed:
                    self.current_speed[ped] = self.current_speed[ped] * max_speed/mod_speed
                new_pose = self.current_poses[ped] + self.dt * self.current_speed[ped]

                # Check if we end up into an obstacle
                # Rescale position
                position = (np.array(new_pose) + self.map_center)/self.occ_grid['Resolution']
                # Get indexes for map
                position = np.ndarray.astype(position, 'int')
                # If we would end up into an obst we do not move
                if self.occ_grid['Map'][tuple(position)] == 1:
                    print("Inside obstacle")
                    new_pose = self.current_poses[ped]
                    self.current_speed[ped] = np.zeros((2))

                # Update poses
                self.current_poses[ped] = new_pose

                # Update goal
                if utils.calc_mod(self.goals[ped] - new_pose) <= self.tol:
                    print("Goal reached")
                # Check if stuck
                if utils.calc_mod(self.current_speed[ped]) <= 0.01:
                    print("Stuck")

        return self.current_poses , self.current_speed

# All the stuff here is for test
if __name__ == "__main__":

    data_path ="../../data/multipath"

    grid = load_map(data_path)
    model = ForceModel(predicted=10, occ_grid=grid,ped_num=2)

    model.simulate(2, [5,0,0], data_path)
src/simulators/social_force_model.py
eliatrevisan/IA-MPPI-LBM
# -*- coding: utf-8 -*-

from math import ceil
import numpy as np
import os
import cv2 as cv
import csv
import pylab as pl
from PIL import ImageDraw
from PIL import Image
from numpngw import write_png

# TODO Merge this functions into Support.py
# ------------------------------------------------------------------------------
def create_map_from_obstacles(map_shape=[30, 30],
                              map_resolution=0.1,
                              save_data_path='../data'):
  grid_map = {}
  grid_map['Resolution'] = map_resolution
  grid_map['Size'] = np.array(map_shape)
  grid_map['Map'] = np.zeros((int(grid_map['Size'][0] / grid_map['Resolution']),
                              int(grid_map['Size'][1] / grid_map['Resolution'])))
  map_center = grid_map['Size'] /2 # other datasets /2

  xx, yy = np.mgrid[:map_shape[0] / map_resolution, :map_shape[1] / map_resolution]
  # First read and draw cylinders
  with open("../objects/cylinders.csv", "rb") as f:
    reader = csv.reader(f, delimiter=",")
    for i, line in enumerate(reader):
      if i == 0:
        continue
      center = np.array((float(line[0]) + map_center[0],
                         float(line[1]) + map_center[1])) / map_resolution
      circle = (xx - int(center[0]))**2 + (yy - int(center[1]))**2
      grid_map['Map'][circle <= int(float(line[2]) / map_resolution)**2] = 1

  # Now go with the rectangles
  with open("../objects/rectangles.csv", "rb") as f:
    reader = csv.reader(f, delimiter=",")
    img = Image.fromarray(grid_map['Map'])
    draw = ImageDraw.Draw(img)
    for i, line in enumerate(reader):
      if i == 0:
        continue
      center = np.array((float(line[0]), float(line[1]))) / map_resolution
      shape = np.array((float(line[2]), float(line[3]))) / map_resolution
      theta = (np.pi / 180.0) * float(line[4])  # In radians
      rect = np.array(((-shape[1] / 2., -shape[0] / 2.),
                       (-shape[1] / 2., shape[0] / 2.),
                       (shape[1] / 2., shape[0] / 2.),
                       (shape[1] / 2., -shape[0] / 2.)))

      R = np.array([[np.cos(theta), -np.sin(theta)],
                    [np.sin(theta), np.cos(theta)]])

      rect = np.dot(rect, R) + center[::-1] + map_center / map_resolution
      draw.polygon([tuple(p) for p in rect], fill=1)
    grid_map['Map'] = np.asarray(img)
  print("Occupancy grid done.")

  # Calculate closest obstacles
  grid_map['Closest X'] = np.zeros_like(grid_map['Map']) + 10000
  grid_map['Closest Y'] = np.zeros_like(grid_map['Map']) + 10000

  # Get idx in map of places where there are obstacles
  obst_idx = np.array(np.where(grid_map['Map'] == 1))

  for xx in range(int(grid_map['Size'][0] / grid_map['Resolution'])):
    for yy in range(int(grid_map['Size'][1] / grid_map['Resolution'])):
      distances = np.sqrt(np.sum(np.square(obst_idx - np.expand_dims([xx, yy], axis=1)), axis=0))
      closest_obj = np.argmin(distances)
      # Get index in map of the closest_obj
      grid_map['Closest X'][xx, yy] = obst_idx[0, closest_obj]
      grid_map['Closest Y'][xx, yy] = obst_idx[1, closest_obj]

  grid_map['Closest X'] = grid_map['Closest X'] * grid_map['Resolution'] - map_center[0]
  grid_map['Closest Y'] = grid_map['Closest Y'] * grid_map['Resolution'] - map_center[1]

  print("Closest obstacles positions calculated.")
  print("Saving map...")
  np.save(os.path.join(save_data_path, 'map'), grid_map)
  print("Map saved.")
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def create_map_from_png(load_data_path='../data/datasets/ewap_dataset/seq_eth',
                        save_data_path='../data'):
  # Create grid for SF model
  print("create_map_from_png")
  grid_map = {}
  grid_map['Resolution'] = 0.1
  grid_map['Size'] = np.array([20., 7.])  # map size in [m]
  grid_map['Map'] = np.zeros((int(grid_map['Size'][0] / grid_map['Resolution']),
                              int(grid_map['Size'][1] / grid_map['Resolution'])))

  map_center = grid_map['Size'] / 2. # hack for my dataset
  H = np.genfromtxt(os.path.join(load_data_path, 'H.txt'),
                    delimiter='  ',
                    unpack=True).transpose()

  # Extract static obstacles
  obst_threshold = 200
  static_obst_img = cv.imread(os.path.join(load_data_path, 'map.png'), 0)
  #static_obst_img = cv.imread(os.path.join(load_data_path, 'map.png'), 0)
  obstacles = np.zeros([0, 3])
  #grid = (static_obst_img/255*-1)+1
  #static_obst_img = np.transpose(static_obst_img)
  for xx in range(static_obst_img.shape[0]):
    for yy in range(static_obst_img.shape[1]):
      if static_obst_img[xx, yy] > obst_threshold:
        obstacles = np.append(obstacles,
                              np.dot(H, np.array([[xx], [yy], [1]])).transpose(),
                              axis=0)
  # Compute obstacles in 2D
  obstacles_2d = np.zeros([obstacles.shape[0], 2])
  obstacles_2d[:, 0] = obstacles[:, 0] / obstacles[:, 2]
  obstacles_2d[:, 1] = obstacles[:, 1] / obstacles[:, 2]

  # Get obstacle idx on map
  obst_idx = []
  for obst_ii in range(obstacles_2d.shape[0]):
    obst_idx.append(idx_from_pos(obstacles_2d[obst_ii, 0],
                    obstacles_2d[obst_ii, 1],
                    map_center,grid_map['Resolution']))
    grid_map['Map'][obst_idx[-1]] = 1

  grid_map['Closest X'] = np.zeros_like(grid_map['Map']) + 10000
  grid_map['Closest Y'] = np.zeros_like(grid_map['Map']) + 10000
  # Calculate closest obstacle
  obst_idx = np.array(obst_idx)
  for xx in range(int(grid_map['Size'][0] / grid_map['Resolution'])):
    for yy in range(int(grid_map['Size'][1] / grid_map['Resolution'])):
      delta_idx = obst_idx - np.array([xx, yy])
      distances = np.sqrt(np.sum(np.square(delta_idx), axis=1))
      closest_obj = np.argmin(distances)
      grid_map['Closest X'][xx, yy] = obst_idx[closest_obj][0]
      grid_map['Closest Y'][xx, yy] = obst_idx[closest_obj][1]

  grid_map['Closest X'] = grid_map['Closest X'] * grid_map['Resolution'] - map_center[0]
  grid_map['Closest Y'] = grid_map['Closest Y'] * grid_map['Resolution'] - map_center[1]

  np.save(os.path.join(save_data_path, 'map'), grid_map)
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def rotateGridAroundCenter(grid, angle):
  # Rotate grid into direction of initial heading
  grid = grid.copy()
  rows, cols = grid.shape
  M = cv.getRotationMatrix2D(center=(rows/2, cols/2), angle=angle, scale=1)
  grid = np.array(cv.warpAffine(grid, M, (rows, cols)))

  grid[grid > 0] = 1 # This way we remove artifacts around borders
  grid[grid < 1] = 0 # This way we remove any negative value

  return grid
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def idx_from_pos(x, y, center, res=0.1):
  idx_x = round((x + float(center[0])) / res)
  idx_y = round((y + float(center[1])) / res)

  # Projecting index on map if out of bounds
  idx_x = max(0, min(idx_x, -1 + center[0] * 2. / res))
  idx_y = max(0, min(idx_y, -1 + center[1] * 2. / res))
  return int(idx_x), int(idx_y)
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_mod(vect):
  vect = np.array(vect)
  return np.sqrt(pow(vect[0], 2) + pow(vect[1], 2))
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_angle(x, y):
  angles = np.array(np.arctan2(y, x) * 180 / np.pi)
  angles[angles == 180.] = -180. # This way output is in [-180,180) and not (-180,180]
  return angles
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_rotation_matrix(angle):
  """
  This function calculates the active rotation matrix.
  Thus if you want to rotate the ref system and not the point you shoud use
  its transpose.
  Thus R = [c, -s; s, c]
  """
  theta = np.array(np.radians(angle))
  c, s = np.cos(theta), np.sin(theta)
  R = np.zeros((len(theta), 2, 2))
  R[:, 0, 0] = c
  R[:, 0, 1] = -s
  R[:, 1, 0] = s
  R[:, 1, 1] = c
  return R
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_direction(current_pose, desired_pose):
  """
  This function calculates the normalized vector going from the current pose
  towards the desired pose.
  """
  tmp = np.array(desired_pose) - np.array(current_pose)
  mod_e_t = calc_mod(tmp)

  e_t = np.zeros((2))

  if mod_e_t is not 0:
    e_t = tmp / mod_e_t

  return e_t
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def check_if_occupied(pose, other_poses):
  """
  This function is used in the SF model to check if a position
  is already occupied by another pedestrian.
  """
  ped_radius = 0.2
  for index, ped_pose in enumerate(other_poses):
    distance = calc_mod(pose - ped_pose)
    if distance <= 2 * ped_radius:
      return True

  return False
# ------------------------------------------------------------------------------

def write_map_png(file, grid):
  grid = rotateGridAroundCenter(grid,90)
  #grid = (grid-1)*-255
  img = grid.astype(np.uint8)
  write_png(file, img, bitdepth=1)

# ------------------------------------------------------------------------------
def load_map(data_path):
  try:
    file_path = data_path + '/map.npy'
    grid_map = np.load(file_path)[()]
    print("Map loaded")
    return grid_map
  except IOError:
    print("ERROR: Cannot load map.")
    return False
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def show_map(grid_map):
  i = 1
  for k in grid_map.keys():
    if k in ["Size", "Resolution"]:
      continue
    fig = pl.figure(i)
    pl.imshow(grid_map[k], cmap=pl.cm.ocean)
    pl.colorbar()
    pl.title(k)
    i = i + 1
  pl.show()
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
def calc_direction(current_pose, desired_pose):
    """
    This function calculates the normalized vector going from the current pose
    towards the desired pose.
    """
    tmp = np.array(desired_pose) - np.array(current_pose)
    mod_e_t = calc_mod(tmp)

    e_t = np.zeros((2))

    if mod_e_t != 0:
        e_t = tmp/mod_e_t

    return e_t
# ------------------------------------------------------------------------------
if __name__ == "__main__":
  create_map_from_png(
    load_data_path='/media/bdebrito/7697ec91-468a-4763-b1c3-135caa7f5aed/home/code/lstm_pedestrian_prediction/data/multipath/corridor_3peds',
  save_data_path = '/media/bdebrito/7697ec91-468a-4763-b1c3-135caa7f5aed/home/code/lstm_pedestrian_prediction/data/multipath/corridor_3peds')
  grid_map = load_map("/media/bdebrito/7697ec91-468a-4763-b1c3-135caa7f5aed/home/code/lstm_pedestrian_prediction/data/multipath/corridor_3peds")
  show_map(grid_map)
  print("DOne!")
src/simulators/social_forces_utils.py
eliatrevisan/IA-MPPI-LBM
import sys
import os

sys.path.append('../')

# With debugger:
#os.chdir(os.getcwd() + "/src/")

import numpy as np
import argparse
import pylab as pl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pickle as pkl
import importlib
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import random
import progressbar
import scipy.io as sio
import json
from copy import deepcopy
import colorama
from colorama import Fore, Style

if sys.version_info[0] < 3:
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from Performance import *
	from utils import *
	import Recorder as rec
else:
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.Performance import *
	from src.data_utils.utils import *
	from src.data_utils.Recorder import Recorder as rec
	from src.models.CV_model import ConstantVelocity

"""
exp_num = 33
model_name = "VGDNN" 
num_test_sequences = 100 
scenario = "simulation/roboat" 
record = False 
n_samples = 1 
unit_testing = False 
freeze_other_agents = False
"""

from random import seed
seed(1)

# Model directories
def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
											help='Path to directory that comprises the model (default="model_name").',
											type=str, default="VGDNN")
	parser.add_argument('--num_test_sequences', help='Number of test sequences', type=int, default=10)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=33)
	parser.add_argument('--n_samples', help='Number of samples', type=int, default=10)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
											type=str, default="datasets/ewap_dataset/seq_eth")
	parser.add_argument('--record', help='Is grid rotated? (default=True).', type=sup.str2bool,
											default=True)
	parser.add_argument('--save_figs', help='Save figures?', type=sup.str2bool,
											default=True)
	parser.add_argument('--noise_cell_state', help='Adding noise to cell state of the agent', type=float,
											default=0)
	parser.add_argument('--noise_cell_grid', help='Adding noise to cell state of the grid', type=float,
											default=5)
	parser.add_argument('--real_world_data', help='real_world_data', type=sup.str2bool,
											default=False)
	parser.add_argument('--update_state', help='update_state', type=sup.str2bool,
											default=False)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
											default=False)
	parser.add_argument('--freeze_other_agents', help='FReeze other agents.', type=sup.str2bool,
											default=False)
	parser.add_argument('--unit_testing', help='Run  Unit Tests.', type=sup.str2bool,
											default=False)
	parser.add_argument('--constant_velocity', help='Run CV comparison', type=sup.str2bool,
											default=False)
	args = parser.parse_args()

	return args


test_args = parse_args()

if test_args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

cwd = os.getcwd()

model_path = os.path.normpath(cwd + '/../') + '/trained_models/' + test_args.model_name + "/" + str(test_args.exp_num)

print("Loading data from: '{}'".format(model_path))
file = open(model_path + '/model_parameters.pkl', 'rb')
if sys.version_info[0] < 3:
	model_parameters = pkl.load(file)  # ,encoding='latin1')
else:
	model_parameters = pkl.load(file, encoding='latin1')
file.close()
args = model_parameters["args"]

with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

print("Training steps: ", args.total_training_steps)

# change some args because we are doing inference
truncated_backprop_length = args.truncated_backprop_length
args.truncated_backprop_length = 1
args.batch_size = 1
args.keep_prob = 1.0


training_scenario = args.scenario
args.scenario = test_args.scenario
args.real_world_data = test_args.real_world_data
args.dataset = '/' + args.scenario + '.pkl'
data_prep = dhlstm.DataHandlerLSTM(args)

# Only used to create a map from png
# Load Map Parameters
map_params = os.path.join(args.data_path+args.scenario, 'map.json')
with open(map_params) as json_file:
	data = json.load(json_file)
map_args = {"file_name": data["file_name"],
						"resolution": data["resolution"],
						"map_size": np.array(data["map_size"]),
						"map_center": np.array(data["map_center"])}

data_prep.processData(**map_args)
if args.normalize_data:
	data_prep.compute_min_max_values()

# Import model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

model = NetworkModel(args)

# Lists for logging of the input / output data of the model
input_list = []
grid_list = []
goal_list = []
ped_grid_list = []
y_ground_truth_list = []
y_pred_list = []  # uses ground truth as input at every step
other_agents_list = []
all_predictions = []
all_traj_likelihood = []
trajectories = []
batch_y = []
batch_loss = []

config = tf.ConfigProto(
	device_count={'GPU': 0}
)


def clear_plot():
	axarr[0].clear()
	axarr[0].set_title('Training Sample')

if test_args.unit_testing:
	data_handler = dhlstm.DataHandlerLSTM(args)
	data_handler.unit_test_data_(map_args)

with tf.Session(config=config) as sess:
	model.warmstart_model(args, sess)
	try:
			model.warmstart_convnet(args, sess)
	except:
		print("")

	for exp_id in range(np.minimum(test_args.num_test_sequences,len(data_prep.trajectory_set)-1)):
		predictions = []
		traj_likelihood = []

		# sample a trajectory id for testing
		traj_id = random.randint(0, len(data_prep.trajectory_set) - 1)
		

		batch_x, batch_vel, batch_pos,batch_goal, batch_grid, other_agents_info, batch_target,batch_end_pos, other_agents_pos, traj = data_prep.getTrajectoryAsBatch(
			traj_id,freeze = test_args.freeze_other_agents)  
		
		trajectories.append(traj)
		x_input_series = np.zeros([0, (args.prev_horizon + 1) * args.input_dim])
		goal_input_series = np.zeros([0, 2])
		grid_input_series = np.zeros(
			[0, int(args.submap_size), int(args.submap_size)])
		if args.others_info == "relative":
			ped_grid_series = np.zeros([0, args.n_other_agents,args.pedestrian_vector_dim])
		elif "sequence" in args.others_info:
			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim*args.prediction_horizon])
		elif args.others_info == "prev_sequence":
			ped_grid_series = np.zeros([0, args.n_other_agents, args.pedestrian_vector_dim *(args.prev_horizon+1)])
		elif args.others_info == "sequence2":
			ped_grid_series = np.zeros([0, args.n_other_agents, args.prediction_horizon,args.pedestrian_vector_dim])
		else:
			ped_grid_series = np.zeros([0, args.pedestrian_vector_dim])
		y_ground_truth_series = np.zeros([0, args.prediction_horizon * 2])
		y_pred_series = np.zeros([0, args.n_mixtures * args.prediction_horizon * args.output_pred_state_dim])

		batch_y.append(batch_target)
		model.reset_test_cells(np.ones((args.batch_size)))
		cell_state_list= []
		cell_ped_list = []
		cell_concat_list = []
		if "grid" in args.model_name:
			batch_ped_grid_backup = np.zeros_like(batch_grid)
			data_prep.add_other_agents_to_grid(batch_ped_grid_backup, batch_x, [other_agents_pos])

		for step in range(batch_x.shape[1]):
			samples = []
			# Assemble feed dict for training
			if "future" in args.others_info:
				if step == 0:
					batch_y_pred = deepcopy(batch_vel)
					for pred_step in range(1, args.prediction_horizon):
						batch_y_pred[:, :, pred_step * 2:pred_step * 2 + 2] = batch_vel[:, :, :2]
				else:
					batch_y_pred = deepcopy(batch_vel)
					batch_y_pred[:, :, 2:] = y_model_pred[:, :, 2:]
			"""
			dict = {"batch_x": batch_x,
							"batch_vel": batch_vel,
							"batch_pos": batch_pos,
							"batch_grid": batch_grid,
							"batch_ped_grid": other_agents_info,
							"step": step,
							"batch_goal": batch_goal,
							"state_noise": 0.0,
							"grid_noise": 0.0,
							"other_agents_pos": [other_agents_pos]
			}
		
			feed_dict_ = model.feed_test_dic(**dict)
			"""

			dict = {
							"batch_vel": batch_vel,
							"batch_grid": batch_grid,
							"batch_ped_grid": other_agents_info,
							"step": step,
							"other_agents_pos": [other_agents_pos]
			}

			feed_dict_ = model.feed_test_dic(
                **dict
            )

			# Append to logging series
			x_input_series = np.append(x_input_series, batch_x[:, step, :], axis=0)
			grid_input_series = np.append(grid_input_series, batch_grid[:, step, :, :], axis=0)
			goal_input_series = np.append(goal_input_series, batch_goal[:, step, :], axis=0)

			y_ground_truth_series = np.append(y_ground_truth_series, batch_target[:, step, :], axis=0)

			y_model_pred, likelihood = model.predict(sess, feed_dict_, True)

			# Backup cell states for later analysis
			cell_state_list.append(model.test_cell_state_current[0,:])
			#cell_ped_list.append(model.test_cell_state_current_lstm_ped[0, :])
			cell_concat_list.append(model.test_cell_state_current_lstm_concat[0, :])

			# Rotate predictions to global frame
			if args.rotated_grid:
				heading = math.atan2(batch_vel[0, step, 1], batch_vel[0, step, 0])
				rot_mat = np.array([[np.cos(heading), -np.sin(heading)], [np.sin(heading), np.cos(heading)]])
				for pred_step in range(args.prediction_horizon):
					y_model_pred[0][0, 2 * pred_step:2 * pred_step + 2] = np.dot(rot_mat, y_model_pred[0][0,2 * pred_step:2 * pred_step + 2])
				samples.append(y_model_pred[0])
			else:
				samples.append(y_model_pred[:,0,:])

			# If sample more than one trajectory from the model
			for sample_id in range(test_args.n_samples - 1):
				dict = {"batch_x": batch_x,
								"batch_vel": batch_vel,
								"batch_pos": batch_pos,
								"batch_grid": batch_grid,
								"batch_ped_grid": other_agents_info,
								"step": step,
								"batch_goal": batch_goal,
								"state_noise": 0.0,
								"grid_noise": 0.0
								}
				feed_dict_ = model.feed_test_dic(**dict)
				y_model_pred = model.predict(sess, feed_dict_, test_args.update_state)
				samples.append(y_model_pred[:,0,:])

			traj_likelihood.append(likelihood)
			predictions.append(samples)

		latent_analysis = False

		# Plot hidden states evolution over a trajectory
		if latent_analysis:

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			pl.show(block=False)
			x_state = np.arange(0, args.rnn_state_size, 1)
			x_ped = np.arange(0, args.rnn_state_size_lstm_ped * 2, 1)
			x_concat = np.arange(0, args.rnn_state_size_lstm_concat, 1)

			time = np.arange(0, len(cell_state_list), 1)
			R= np.array(cell_state_list)
			print(R)
			ax.clear()
			X, Y = np.meshgrid(x_state,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			save_img_to_file = model_path + "/figs"
			if not os.path.exists(save_img_to_file):
				os.makedirs(save_img_to_file)
			pl.savefig(save_img_to_file+"/cell_state" + str(exp_id) + '.png')
			pl.show()
			pl.close()

			# fig = pl.figure()
			# ax = pl.axes(projection="3d")
			# time = np.arange(0, len(cell_ped_list), 1)
			# R= np.array(cell_ped_list)
			# print(R)
			# ax.clear()
			# X, Y = np.meshgrid(x_ped,time)
			# surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			# ax.view_init(elev=90,azim=90)
			# ax.set_xlabel("Time [s]")
			# ax.set_xlabel("Features")
			# fig.colorbar(surf)
			# pl.savefig(save_img_to_file+"/cell_ped" + str(exp_id) + '.png')
			# pl.close()

			fig = pl.figure()
			ax = pl.axes(projection="3d")
			time = np.arange(0, len(cell_concat_list), 1)
			R= np.array(cell_concat_list)
			ax.clear()
			X, Y = np.meshgrid(x_concat,time)
			surf = ax.plot_surface(X, Y, R,cmap=cm.coolwarm)
			ax.view_init(elev=90,azim=90)
			ax.set_xlabel("Time [s]")
			ax.set_xlabel("Features")
			fig.colorbar(surf)
			pl.savefig(save_img_to_file+"/cell_concat" + str(exp_id) + '.png')
			pl.show()
			pl.close()

		all_predictions.append(predictions)
		all_traj_likelihood.append(traj_likelihood)
		input_list.append(x_input_series)
		goal_list.append(goal_input_series)
		grid_list.append(grid_input_series)
		y_ground_truth_list.append(y_ground_truth_series)
		other_agents_list.append(other_agents_pos)
		# update progress bar

sess.close()

# Save data in Matlab format for post-processing
if not os.path.exists(args.model_path + '/../results/'):
	os.makedirs(args.model_path + '/../results/')
results= {
	"trajectories": trajectories,
	"predictions": all_predictions
}
scenario = args.scenario.split('/')[-1]
results_file = args.model_path + '/'+ scenario + "_results.mat"
sio.savemat(results_file,results)

if test_args.constant_velocity:
	# Make constant velocity predictions
	CVModel = ConstantVelocity(args)
	cv_predictions = CVModel.predict(input_list)
else:
	cv_predictions = None


if test_args.record:
		recorder = rec(args, data_prep.agent_container.occupancy_grid)
		if ( "real_world" in test_args.scenario) and not test_args.unit_testing:
			print("Real data!!")
			recorder.plot_on_video(input_list, grid_list, all_predictions, y_ground_truth_list,
														 other_agents_list,
														 trajectories,all_traj_likelihood, test_args)
		else:
			#recorder.plot_on_image(input_list, grid_list, all_predictions, y_ground_truth_list, other_agents_list,
			#	                       trajectories,test_args)
			recorder.animate_local(input_list, grid_list, ped_grid_list, all_predictions, y_ground_truth_list, other_agents_list,
										 trajectories,test_args)

			recorder.animate_global(input_list, grid_list, all_predictions, y_ground_truth_list,
														 other_agents_list,
														 trajectories, all_traj_likelihood, cv_predictions, test_args)
			#cv_fde_error = compute_fde(args, trajectories, cv_predictions)
			#cv_ade_error = compute_ade(args, trajectories, cv_predictions)

			#print(" ADE: ", cv_ade_error, " FDE: ", cv_fde_error)

			# cv_fde = compute_rolling_fde(args, trajectories, cv_predictions)
	
			# fde_roll = []
			# for horizon in range(0,args.prediction_horizon):
			# 	pred_fde, pred_error_summary_lstm_fde = compute_rolling_trajectory_fde(args, horizon, trajectories, all_predictions)
			# 	fde_roll.append(pred_fde)
			# 	print("Step: ", horizon, " FDE: ", pred_fde)

			# print("CV: ", cv_fde)
			# print("SVRNN: ", fde_roll)

			print("Recorder is done!")
else:
	print("Performance tests")
	
	cv_fde, cv_std = compute_rolling_fde(args, trajectories, cv_predictions)
	
	fde_roll = []
	std_roll = []
	for horizon in range(0,args.prediction_horizon):
		pred_fde, pred_error_summary_lstm_fde, std = compute_rolling_trajectory_fde(args, horizon, trajectories, all_predictions)
		fde_roll.append(np.mean(pred_error_summary_lstm_fde))
		std_roll.append(std)

	# print("CV: ", cv_fde)
	# print("SVRNN: ", fde_roll)

	std_roll = np.array(std_roll)*0.3
	cv_std = cv_std * 0.3

	# fig, ax = plt.subplots(figsize=(4,3))
	# x = [1,2,3,4,5,6,7,8,9,10,11,12]
	# ax.set_title("Herengracht")
	# ax.plot(x, cv_fde, label="CV", color="b")
	# ax.plot(x, fde_roll, label="SVRNN", color="r")
	# ax.fill_between(x, cv_fde+cv_std, cv_fde-cv_std, color="b", alpha=0.2)
	# ax.fill_between(x, fde_roll+std_roll, fde_roll-std_roll, color="r", alpha=0.2)
	# ax.set_xlabel("Timestep")
	# ax.set_ylabel("Displacement Error [m]")
	# plt.tight_layout()
	# plt.legend()
	# plt.show()
	
	pred_error, pred_error_summary_lstm = compute_trajectory_prediction_mse(args, trajectories, all_predictions)
	pred_fde, pred_error_summary_lstm_fde = compute_trajectory_fde(args, trajectories, all_predictions)
	if test_args.constant_velocity: 
		#cv_pred_error = compute_ade_cv(args, trajectories, cv_predictions)
		#cv_fde_error = compute_fde_cv(args, trajectories, cv_predictions)
		cv_fde_error = compute_fde(args, trajectories, cv_predictions)
		cv_pred_error = compute_ade(args, trajectories, cv_predictions)
		print(
		Fore.LIGHTBLUE_EX + "\nConstant Velocity: MSE: {:01.2f}, FDE: {:01.2f}".format(np.mean(cv_pred_error), np.mean(cv_fde_error))+Style.RESET_ALL)
	diversity, diversity_summary = compute_2_wasserstein(args, all_predictions)
	args.scenario = training_scenario
	args.truncated_backprop_length = truncated_backprop_length
	write_results_summary(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde), np.mean(diversity_summary), args, test_args, cv_pred_error, cv_fde_error)
	print(
		Fore.LIGHTBLUE_EX + "\nMSE: {:01.2f}, FDE: {:01.2f}, DIVERSITY: {:01.2f}".format(np.mean(pred_error_summary_lstm), np.mean(pred_error_summary_lstm_fde),np.mean(diversity_summary))+Style.RESET_ALL)
	

src/test_VGDNN.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import tensorflow as tf
import pickle as pkl
# import tensorflow_probability as tfp
import matplotlib.pyplot as plt
if sys.version_info[0] < 3:
  sys.path.append('../src/data_utils')
  sys.path.append('../src/models')
  import DataHandler as dh
  from seqModel import *
  from plot_utils import *
else:
  from src.data_utils import DataHandler as dh
  from src.models.seqBiModel import *
  from src.data_utils.plot_utils import *

tfk = tf.keras

model_name = "seq_to_seq"
data_path = '../data/2_agents/trajs'
input_dim = 8
latent_dim = 256
output_dim = 2
cmd_vector_dim = 2
batch_size = 16
learning_rate_init = 0.01
grads_clip = 10.0
n_epochs = 100
stateful = False
prediction_horizon = 10
prev_horizon = 10
truncated_backprop_length = 10
regularization_weight = 0.0001
dropout = 0.2

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default=model_name)
	parser.add_argument('--data_path', help='Data directory', type=str,
	                    default=data_path)
	parser.add_argument('--save_path', help='Path to directory that saves pickle data (default=" ").', type=str,
	                    default=data_path+'/GA3C-CADRL-10-py27.pkl')
	parser.add_argument('--batch_size', help='Batch size.', type=int,
	                    default=batch_size)
	parser.add_argument('--learning_rate_init', help='Learning rate.', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--grads_clip', help='Grad clip.', type=float,
	                    default=grads_clip)
	parser.add_argument('--n_epochs', help='Number of epochs.', type=float,
	                    default=n_epochs)
	parser.add_argument('--stateful', help='Stateful.', type=bool,
	                    default=stateful)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--prediction_horizon', help='Prediction Horizon.', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--dropout', help='Dropout probability (default=0.005).', type=float,
	                    default=dropout)
	parser.add_argument('--output_dim', help='Output dimension.', type=int,
	                    default=output_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
	                    default=cmd_vector_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=int,
	                    default=input_dim)
	parser.add_argument('--latent_dim', help='Latent dimension.', type=int,
	                    default=latent_dim)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--distance_threshold', help='Distance threshold to start new trajectory', type=float, default=1)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
	                    default='../config/topics.yaml')
	parser.add_argument('--grid_width', help='width of occupancy grid', type=int, default=120)
	parser.add_argument('--grid_height', help='height of occupancy grid', type=int, default=120)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--file_name', help='PNG map file name', type=str,
	                    default='simple_multipath_rotated.png')
	parser.add_argument('--resolution', help='Distance threshold to start new trajectory', type=float,
	                    default=0.1)
	parser.add_argument('--map_size_x', help='Distance threshold to start new trajectory', type=float,
	                    default=20.0)
	parser.add_argument('--map_size_y', help='Distance threshold to start new trajectory', type=float,
	                    default=7.0)

	args = parser.parse_args()

	return args


args = parse_args()

print("Loading data from: '{}'".format(args.save_path))
file = open(args.save_path, 'rb')
if sys.version_info[0] < 3:
	tmp_self = pkl.load(file)
else:
	tmp_self = pkl.load(file, encoding='latin1')
	dataset = tmp_self

data_prep = dh.DataHandler(args)
# Collect data online
if os.path.isfile(args.save_path):
	data_prep.load_data()
else:
	data_prep.start_node()

batch_enc, batch_dec, batch_dec_target, new_epoch = data_prep.getBatch()

model = SeqToSeqModel2(args)

model.model.compile(optimizer=tfk.optimizers.Adam(lr=args.learning_rate_init,clipnorm=args.grads_clip),
                             loss='mean_squared_error')


history = model.model.fit([batch_enc,batch_dec],batch_dec_target,
                epochs=n_epochs,
                batch_size=16,
								#validation_split=0.2,
                shuffle=True)

# Plot training & validation loss values
plt.plot(history.history['loss'])
plt.title('Modtrain.pyel loss during training')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train'], loc='upper left')
plt.show()

batch_enc_seq, batch_dec_seq, traj = data_prep.getTrajectoryAsBatch()

predictions = model.decode_sequence([batch_enc_seq,batch_dec_seq])

#plot_scenario(batch_enc_seq,batch_dec_seq,traj,predictions,data_prep.occupancy_grid,args)
plot_scenario_vel(batch_enc_seq,batch_dec_seq,traj,predictions,args)
print("test")


src/train.py
eliatrevisan/IA-MPPI-LBM
import sys
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl

if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
	from autoencoder_grid import ae_model as ae
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
	from src.autoencoder_grid import ae_model as ae

import tensorflow as tf
import pickle as pkl
import time
from copy import deepcopy

tfk = tf.keras

model_name = "autoencoder_with_ped2"
model_path = '../trained_models/' + model_name
pretrained_convnet_path = "../trained_models/autoencoder_with_ped"
log_dir = model_path + '/log'
data_path = '../data/2_agents_swap/trajs/'
scenario = 'GA3C-CADRL-10-py27'
data_path = '../data/cyberzoo_experiments/'
scenario = 'all_trajectories'
exp_num = 6
data_path = '../data/'
scenario = '20_ped_with_obstacles/short_few_obstacles'

# Hyperparameters
n_epochs = 1200

batch_size = 16
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 10
prediction_horizon = 10
prev_horizon = 0

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 128
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 0.8
dropout = False
reg = 1e-4
n_mixtures = 3  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
dropout = False
tensorboard_logging = False

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 5  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 10

print_freq = 2000
save_freq = 500
dt = 0.4

warmstart_model = False
pretrained_convnet = True
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = True
rotated_grid = False
centered_grid = True
noise = False
normalize_data = True
real_world_data = False

# Map parameters
submap_resolution = 0.1
submap_width = 6
submap_height = 6


def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
	                    help='Path to directory that comprises the model (default="model_name").',
	                    type=str, default=model_name)
	parser.add_argument('--model_path',
	                    help='Path to directory to save the model (default=""../trained_models/"+model_name").',
	                    type=str, default=model_path)
	parser.add_argument('--pretrained_convnet_path',
	                    help='Path to directory that comprises the pre-trained convnet model (default=" ").',
	                    type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
	                    help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
	                    type=str, default=log_dir)
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
	                    type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
	                    default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
	                    default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
	                    default='/' + scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
	                    default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
	                    default=warmstart_model)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
	                    default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
	                    default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
	                    default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
	                    default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
	                    default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
	                    default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
	                    default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
	                    type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
	                    default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
	                    default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
	                    default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_size_lstm_grid',
	                    help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
	                    help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
	                    type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
	                    type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
	                    type=int, default=latent_space_size)
	parser.add_argument('--fc_hidden_unit_size',
	                    help='Number of fully connected layer units after LSTM layer (default=64).',
	                    type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=float,
	                    default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
	                    default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
	                    default=output_dim)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
	                    default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
	                    default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
	                    default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
	                    default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
	                    default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
	                    default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
	                    default=pedestrian_radius)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=int, default=0)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
	                    default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
	                    type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
	                    default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
	                    default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
	                    default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
	                    default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
	                    default=0)
	parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=submap_width)
	parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=submap_height)
	parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
	                    default=False)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
	                    default='../config/topics.yaml')
	args = parser.parse_args()

	return args


args = parse_args()

args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)

model_parameters = {"args": args}
print(args)
# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()

if "GA3C" in args.scenario:
	data_prep = dh.DataHandler(args)
	# Collect data online
	if os.path.isfile(args.data_path + args.dataset):
		data_prep.load_data()
	else:
		data_prep.start_node()
else:
	data_prep = dhlstm.DataHandlerLSTM(args)
	# Only used to create a map from png
	map_args = {"file_name": 'map.png',
	            "resolution": 0.1,
	            "map_size": np.array([30., 30.]), }
	data_pickle = args.data_path + args.scenario + "/data" + str(prediction_horizon) + ".pickle"
	data_prep.processData(data_pickle, **map_args)

model = ae.ConvAEModel(learning_rate_init=learning_rate_init, grid_width=submap_width*10, grid_height=submap_height*10,
                         log_dir=log_dir, is_training=True, latent_space_dim=64)

config = tf.ConfigProto(
	device_count={'GPU': 0}
)

fig, axarr = pl.subplots(1, 2)
pl.show(block=False)


def clear_plot():
	axarr[0].clear()
	axarr[1].clear()
	axarr[0].set_title('Training Sample')
	axarr[1].set_title('Training Output')


with tf.Session(config=config) as sess:
	# Timing
	avg_time_comp = 0.0
	avg_time_total = 0.0

	if args.warmstart_model:
		model.warmstart_model(args, sess)
	else:
		# Initialize all TF variables
		sess.run(tf.global_variables_initializer())
	model.warmstart_convnet(args, sess)

	# if the trainning was interrupted
	try:
		step = int(open(args.model_path + "/tf_log", 'r').read().split('\n')[-2]) + 1
	except:
		step = 1
	epoch = 0
	training_loss = []
	diversity_loss = []
	while epoch <= n_epochs:
		start_time_loop = time.time()
		# if res == None:
		batch_x, batch_vel, batch_grid, batch_ped_grid, batch_y, other_agents_pos, new_epoch = data_prep.getBatch()

		if new_epoch:
			epoch += 1

		data_prep.add_other_agents_to_grid(batch_grid, batch_x, other_agents_pos, args)

		feed_dict_train = {model.input_placeholder: batch_grid[:, 0, :, :],
		                   model.output_placeholder: batch_grid[:, 0, :, :]}
		time_start = time.time()
		_, total_loss, deconv1, summary_str = sess.run([model.update, model.total_loss, model.deconv1, model.summary],
		                                               feed_dict=feed_dict_train)

		model.summary_writer.add_summary(summary_str,   step)
		time_avg += time.time() - time_start
		if step % print_freq == 0:
			clear_plot()
			sup.plotGrid(batch_grid[0, 0, :, :], axarr[0, 0], color='k', alpha=0.5)
			sup.plotGrid(deconv1[0, :, :], axarr[0, 1], color='k', alpha=0.5)
			fig.canvas.draw()
			time_avg = 0.0

		# Save model every now and then
		if step % save_freq == 0:
			save_path = model_path + '/model_ckpt'
			model.saver.save(sess, save_path, global_step=  step)
			print('Epoch {} Step {}: Saving model under {}'.format(epoch,   step, save_path))
			full_path = model_path + '/final-model.ckpt'
			model.saver.save(sess, full_path)
			print('Saved final model under "{}"'.format(full_path))

		step +=1

sess.close()

print("test")

src/train_autoencoder.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
import json
import importlib
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils import DataHandlerRoboat as dhroboat
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
from datetime import datetime
import pickle as pkl
import time
from copy import deepcopy
from multiprocessing.pool import ThreadPool
import colorama
from colorama import Fore, Style

#os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

pretrained_convnet_path = "../trained_models/autoencoder"
exp_num = 6
data_path = '../data/roboat_dt4/'
scenario = 'herengracht'
no_stopping = True

# Hyperparameters
n_epochs = 0

batch_size = 16 #128
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 8
prediction_horizon = 12
prev_horizon = 8

# Model params
# rnn_state_size = 32
# rnn_state_size_lstm_grid = 256
# rnn_state_size_lstm_ped = 128
# rnn_state_ped_size = 16
# rnn_state_size_lstm_concat = 512
# prior_size = 512
# latent_space_size = 256
# x_dim = 512
# fc_hidden_unit_size = 256

rnn_state_size = 32
rnn_state_size_lstm_grid = 64
rnn_state_size_lstm_ped = 64 
rnn_state_size_lstm_concat = 128
latent_space_size = 0

# Not used
rnn_state_ped_size = 1
prior_size = 1
x_dim = 1
fc_hidden_unit_size = 1

# Regularization and training params
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 18
tensorboard_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 5

print_freq = 200
save_freq = 500
total_training_steps = 50000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = True
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = False
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
#submap_resolution = 0.081
#submap_width = 4.86
#submap_height = 4.86
submap_size = 60 # In pixels
submap_span_real = 10 # In Map coordinates
relative_covariance = 10.0
diversity_update = False
predict_positions = False
warm_start_convnet = True
debug_plotting = False

# Dataset division
train_set = 0.8

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
						help='Path to directory that comprises the model (default="model_name").',
						type=str, default= "RNN")
	parser.add_argument('--model_path',
						help='Path to directory to save the model (default=""../trained_models/"+model_name").',
						type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
						help='Path to directory that comprises the pre-trained convnet model (default=" ").',
						type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
						help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
						type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
						type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
						default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
						default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
						default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
						default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
						default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
						default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
						default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--total_training_steps', help='Number of training steps (default=20000).', type=int, default=total_training_steps)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
						default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
						default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
						default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
						default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
						default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
						default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
						type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
						default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
						default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
						default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
						help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
						type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
						type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
						type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
						help='Number of fully connected layer units after LSTM layer (default=64).',
						type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
						default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
						default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
						default=output_dim)
	parser.add_argument('--goal_size', help='Goal dimension (default=).', type=int,
						default=2)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
						default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
						default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
						default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
						default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
						default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
						default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
						default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the network.', type=int,
						default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=sup.str2bool, default=debug_plotting)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
						default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
						type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
						default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
						default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
						default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
						default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
						default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
						default=0)
	#parser.add_argument('--submap_width', help='width of occupancy grid (For initialization of arrays)', type=int, default=submap_width)
	#parser.add_argument('--submap_height', help='height of occupancy grid (for initiatization of arrays)', type=int, default=submap_height)
	parser.add_argument('--submap_size', help='size of submap in pixels (60)', type=int, default=submap_size)
	parser.add_argument('--submap_span_real', help='span of submap occupancy grid', type=int, default=submap_span_real)
	#parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	#                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
						default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
						default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
						default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
						default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
						default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
						default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
						default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
						default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	parser.add_argument('--kl_weight', help='Experimenting with additional weight for KL loss', type=float, default=1)
	parser.add_argument('--relative_covariance', help='Changing covariance of relative positions', type=float, default=relative_covariance)
	args = parser.parse_args()

	return args

args = parse_args()


# Cyclical KL Annealing weight
def frange_cycle_linear(n_iter, start=0.0, stop=1.0,  n_cycle=4, ratio=0.5):
	L = np.ones(n_iter) * stop
	period = n_iter/n_cycle
	step = (stop-start)/(period*ratio) # linear schedule

	for c in range(n_cycle):
		v, i = start, 0
		while v <= stop and (int(i+c*period) < n_iter):
			L[int(i+c*period)] = v
			v += step
			i += 1
	return L

# Enable / Disable GPU
if args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

# Create Log and Model Directory to save training model
args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)
args.log_dir = args.model_path + '/log'
if not os.path.exists(args.log_dir):
	os.makedirs(args.log_dir)
args.dataset = '/' + args.scenario + '.pkl'
model_parameters = {"args": args}

# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)

# Save Model Parameters
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()
with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

# Specify the number of datasets / scenarios 
n_scenarios = 4

# Create Datahandler class for Herengracht experiment
original_scenario = args.scenario
data_prep = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.081,
				"map_size": np.array([1715,881]),
				"map_center": np.array([-78,-40])}
data_prep.processData(**map_args)

# Create Datahandler class for Prinsengracht scenario
args.scenario = "prinsengracht"
data_prep_2 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.091,
				"map_size": np.array([816,816]),
				"map_center": np.array([-37.1, -37.1])} 
data_prep_2.processData(**map_args)

args.scenario = "bloemgracht"
data_prep_3 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.22,
				"map_size": np.array([1469,928]),
				"map_center": np.array([-161.59,-102.08])} 
data_prep_3.processData(**map_args)

args.scenario = "open_crossing"
data_prep_4 = dhlstm.DataHandlerLSTM(args)
map_args = {"file_name": 'map.png',
				"resolution": 0.6,
				"map_size": np.array([100,100]),
				"map_center": np.array([-30,-30])} 
data_prep_4.processData(**map_args)

args.scenario = original_scenario

# Initialize main data handler
dpreps = [data_prep, data_prep_2, data_prep_3, data_prep_4]
data_roboat = dhroboat.DataHandlerRoboat(args)

# Combine trajectory sets
data_roboat.CombineTrajectorySets(dpreps)

min = 100
print(len(data_roboat.trajectory_set))
for i in range(len(data_roboat.trajectory_set)):
	if len(data_roboat.trajectory_set[i][1].pose_vec) < min:
		min = len(data_roboat.trajectory_set[i][1].pose_vec)
print(min)
# Remove end of trajectories if at goal, and trajectories that are too short
if no_stopping:
	data_roboat.RemoveStops()

min = 100
print(len(data_roboat.trajectory_set))
for i in range(len(data_roboat.trajectory_set)):
	if len(data_roboat.trajectory_set[i][1].pose_vec) < min:
		min = len(data_roboat.trajectory_set[i][1].pose_vec)
print(min)
# Import Deep Learning model
module = importlib.import_module("src.models."+args.model_name)
globals().update(module.__dict__)

# Create Model Graph
model = NetworkModel(args)

config = tf.ConfigProto()

beta_list = frange_cycle_linear(args.total_training_steps)

# Start Training Session
with tf.Session(config=config) as sess:

	# Load a pre-trained model
	if args.warmstart_model:
		model.warmstart_model(args, sess)
	else:
		# Initialize all TF variables
		sess.run(tf.global_variables_initializer())

	# Load Convnet Model
	try:
		if args.warm_start_convnet:
			model.warmstart_convnet(args, sess)
	except:
		print("Failed to initialized Convnet or Convnet does not exist")

	# if the training was interrupted load last training step index
	try:
		initial_step = int(open(args.model_path + "/tf_log", 'r').read().split('\n')[-2]) + 1
	except:
		initial_step = 1

	epoch = 0
	training_loss = []
	diversity_loss = []
	training_loss.append(0)

	# Set up multithreadign for data handler
	pool = ThreadPool(1)
	res = None
	_model_prediction = []
	start_time = time.time()
	best_loss = float('inf')
	avg_training_loss = np.ones(100)


	for step in range(initial_step,args.total_training_steps):
		start_time_loop = time.time()

		# Get Next Batch of Data
		#batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch, train_log = data_roboat.getBatch(dpreps)
		batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch, train_log = data_roboat.getRoboatBatch(dpreps)

		# Create dictionary to feed into the model
		dict = {"batch_x": batch_x,
				"batch_vel": batch_vel,
				"batch_pos": batch_pos,
				"batch_goal": batch_goal,
				  "batch_grid": batch_grid,
				  "batch_ped_grid": batch_ped_grid,
				  "step": step,
				  "beta": beta_list[step],
				  "batch_y": batch_y,
				  "batch_pos_target": batch_pos_target,
				"batch_div": batch_y,
				"other_agents_pos": other_agents_pos
				  }

		feed_dict_train = model.feed_dic(**dict)

		# Here was new epoch before....

		# Initialize the new sequences with a hidden state of zeros, the continuing sequences get assigned the previous hidden state
		model.reset_cells(data_roboat.sequence_reset)

		start_time_training = time.time()

		model_output = model.train_step(sess, feed_dict_train, step)

		avg_training_time = time.time() - start_time_training
		avg_loop_time = time.time() - start_time_loop

		training_loss.append(model_output["batch_loss"])

		# Print training info
		if step % print_freq == 0:
			
			# Get test batch for model validation
			#validation_dict, val_log = data_roboat.getTestBatch(dpreps)
			validation_dict, val_log = data_roboat.getRoboatValidationBatch(dpreps)
			model.reset_test_cells(data_roboat.val_sequence_reset)

			feed_dict_validation = model.feed_val_dic(**validation_dict)
			validation_loss, validation_summary, validation_predictions = model.validation_step(sess, feed_dict_validation)
			#validation_loss, validation_summary, validation_predictions = model.validation_step(sess, feed_dict_train)

			autoencoder_loss = model_output["autoencoder_loss"]

			ellapsed_time = time.time() - start_time

			print(Fore.BLUE + "\n\nEpoch {:d}, Steps: {:d}, Train loss: {:01.2f}, Validation loss: {:01.2f}, AutoEncoder Loss: {:01.2f},  Epoch time: {:01.2f} sec"
				  .format(epoch + 1, step, np.mean(training_loss), validation_loss, autoencoder_loss, ellapsed_time)+Style.RESET_ALL)

			if tensorboard_logging:
				model.summary_writer.add_summary(model_output["summary"], step)
				model.summary_writer.flush()
				model.validation_summary_writer.add_summary(validation_summary, step)
				model.validation_summary_writer.flush()

			with open(args.model_path + "/tf_log", 'a') as f:
				f.write(str(step) + '\n')
				f.write("TrainLog: " + str(train_log) + '\n')
				f.write("ValidationLog: " + str(val_log) + '\n')
			curr_loss = (validation_loss+np.mean(avg_training_loss))/2.0
			if curr_loss < best_loss:
				save_path = args.model_path + '/model_ckpt'
				model.full_saver.save(sess, save_path, global_step=step)
				best_loss = curr_loss
				print(Fore.LIGHTCYAN_EX+'Step {}: Saving model under {}'.format(step, save_path))
		
		# Change this back to the top to test what happens...
		if new_epoch:
			epoch += 1

	write_summary(training_loss[-1], args)
	full_path = args.model_path + '/final-model.ckpt'
	model.full_saver.save(sess, full_path)
	print('Saved final model under "{}"'.format(full_path))

	if tensorboard_logging:
		model.summary_writer.close()

sess.close()



src/train_roboat.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import matplotlib.pyplot as plt

def frange_cycle_linear(n_iter, start=0.0, stop=1.0,  n_cycle=4, ratio=0.5):
	L = np.ones(n_iter) * stop
	period = n_iter/n_cycle
	step = (stop-start)/(period*ratio) # linear schedule

	for c in range(n_cycle):
		v, i = start, 0
		while v <= stop and (int(i+c*period) < n_iter):
			L[int(i+c*period)] = v
			v += step
			i += 1
	return L

def plot_beta(beta_list):

	plt.figure()
	plt.plot(beta_list)
	plt.show()


def main():
	n = 10

	beta_list = []
	"""

	for i in range(1,n):
		beta = 0
		beta = frange_cycle_linear(i)
		print("New")
		print(beta)
		beta_list.append(beta)
	"""
	#plot_beta(beta_list)

	print(frange_cycle_linear(100))

if __name__ == "__main__":
	main()
tests/KLannealing.py
eliatrevisan/IA-MPPI-LBM
import os
import csv
import numpy as np


home = os.path.expanduser("~")
directory = "/compareinputs/combine/"

def main():
    
    names = getFileNames()
    data = []
    agents = 0
    two = 0
    three = 0
    four = 0
    timestamp = 0

    for file in names:

        content = getDataFromFile(file)
        n_agents = getNumofAgents(content)

        #TODO Make a function to discard the first and last 10 time steps (first and last second). Or different amount of time in better 

        if n_agents == 2:
            two += 1
        elif n_agents == 3:
            three += 1
        elif n_agents == 4:
            four += 1

        content = reindex_agents(content, agents,file)
        content = reindex_timestamp(content, timestamp)
        agents = agents + n_agents
        timestamp = last_timestep(content)
        data.append(content)
    
    save_to_file(data)

def getDataFromFile(file):
    """
    Return the data contained in csv file
    """
    data = np.genfromtxt(home + directory + file, delimiter=",")[11:-10,:]
    return data

def getNumofAgents(file):
    """
    Return the number of agents in data file
    """
    return len(np.unique(file[:,0]))

def save_to_file(data):
    """
    Save data to csv file
    """
    fields = ['id', 'timestep_s', 'timestep_ns', 'pos_x', 'pox_y','vel_x', 'vel_y', 'goal_x', 'goal_y' ]

    with open(home + directory + 'total_log.csv', 'w+') as f:
        write = csv.writer(f)
        write.writerow(fields)
        for i in range(len(data)):
            write.writerows(data[i])


def reindex_agents(data, agents,file):
    """
    Reindex the roboat ID's
    """
    for i in range(len(data)):
        if data[i][0] == 0:
            data[i][0] = agents + 1
        elif data[i][0] == 1:
            data[i][0] = agents + 2
        elif data[i][0] == 2:
            data[i][0] = agents + 3
        elif data[i][0] == 3:
            data[i][0] = agents + 4
        else:
            print("Wrong agent ID")
            print("Filename: ", file)
    return data

def reindex_timestamp(data, timestamp):
    """
    Reindex the timestamp to make final file continuous in time
    """
    dt = 0.1
    data[:,1] = data[:,1] + timestamp + dt
    return data

def last_timestep(data):
    last = len(data) - 1
    return data[last][1]

def getFileNames():
    """
    Retrieve file names of all csv files in specified directory
    """
    names = []
    for root, dirs, files in os.walk(home + directory):
        for file in files:
            if file.endswith('.csv'):
                names.append(file)

    return names


if __name__ == "__main__":
    main()
tests/combine_data.py
eliatrevisan/IA-MPPI-LBM
import sys
sys.path.append('../')
import numpy as np
if sys.version_info[0] < 3:
	sys.path.append('../src/data_utils')
	import DataHandler as dh
else:
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
from parser import *
import argparse
import matplotlib.pyplot as plt 

pretrained_convnet_path = "../trained_models/autoencoder_with_ped"
home = os.path.expanduser('~')

exp_num = 2
data_path = '../data/'
scenario = 'simulation/roboat'

# Hyperparameters
n_epochs = 2

batch_size = 16
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 15
prediction_horizon = 15
prev_horizon = 15

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 128
rnn_state_ped_size = 16
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 256
x_dim = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 18
tensorboard_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 4
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 5

print_freq = 200
save_freq = 500
total_training_steps = 20000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = False
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
submap_resolution = 0.081
submap_width = 4.86
submap_height = 4.86
diversity_update = True
predict_positions = False
warm_start_convnet = True
debug_plotting = False

# Dataset division
train_set = 0.8

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
											help='Path to directory that comprises the model (default="model_name").',
											type=str, default= "RNN")
	parser.add_argument('--model_path',
											help='Path to directory to save the model (default=""../trained_models/"+model_name").',
											type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
											help='Path to directory that comprises the pre-trained convnet model (default=" ").',
											type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
											help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
											type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
											type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
											default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
											default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
											default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
											default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
											default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
											default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
											default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--total_training_steps', help='Number of training steps (default=20000).', type=int, default=total_training_steps)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
											default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
											default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
											default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
											default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
											default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
											default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
											type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
											default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
											default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
											default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
											help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
											type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
											help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
											type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
											help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
											type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
											help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
											type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
											type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
											type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
											type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
											help='Number of fully connected layer units after LSTM layer (default=64).',
											type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
											default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
											default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
											default=output_dim)
	parser.add_argument('--goal_size', help='Goal dimension (default=).', type=int,
											default=2)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
											default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
											default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
											default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
											default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
											default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
											default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
											default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the network.', type=int,
											default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=sup.str2bool, default=debug_plotting)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
											default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
											type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
											default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
											default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
											default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
											default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
											default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
											default=0)
	parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=submap_width)
	parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=submap_height)
	parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
											default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
											default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
											default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
											default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
											default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
											default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
											default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
											default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
											default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	args = parser.parse_args()

	return args


args = parse_args()

class Map:
    size = [1715, 881]
    data = pl.imread(home + '/I-LSTM/data/simulation/roboat/map.png')
    origin = [-78, -40, 0.0]
    resolution = 0.081

def plot_map(map, ax):
    ax.imshow(map.data,
               extent = (map.origin[0], map.origin[0] + map.size[0] * map.resolution,
                         map.origin[1], map.origin[1] + map.size[1] * map.resolution),
               cmap='gray_r')
    ax.set_facecolor('#262626')
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    ax.yaxis.set_label_coords(-0.08, .5)
    ax.xaxis.set_label_coords(0.5, -0.09)

data_prep = dhlstm.DataHandlerLSTM(args)
# Only used to create a map from png
# Make sure this parameters are correct otherwise it will fail training and plotting the results
map_args = {"file_name": 'map.png',
							"resolution": 0.081,
							"map_size": np.array([1715,881]),
				"map_center": np.array([-78,-40])} # THIS IS HARDCODED FOR SOME REASON AND NOT EXTRACTED FROM MAP.JSON ?
# Load dataset
data_prep.processData(**map_args)

#print(data_prep.trajectory_set[0]) # Output [ vessel id, trajectory_object ]

print(data_prep.trajectory_set[12][1].vel_vec)

"""
right_start = []
bottom_start = []
left_start = []

for i in range(len(data_prep.trajectory_set)):
	# Check from which direction trajectory comes
	if data_prep.trajectory_set[i][1].pose_vec[0,:][0] > 10:
		right_start.append(i)
	elif data_prep.trajectory_set[i][1].pose_vec[0,:][0] < -20:
		left_start.append(i)
	else:
		bottom_start.append(i)

print(len(right_start), len(left_start), len(bottom_start))


fig, ax = plt.subplots(3)

map = Map()
plot_map(map, ax[0])
plot_map(map, ax[1])
plot_map(map, ax[2])

for i in left_start:
	traj = data_prep.trajectory_set[i][1]
	ax[0].plot(traj.pose_vec[:,0], traj.pose_vec[:,1])
	ax[0].set_title("From left canal")

for i in right_start:
	traj = data_prep.trajectory_set[i][1]
	ax[1].plot(traj.pose_vec[:,0], traj.pose_vec[:,1])
	ax[1].set_title("From right canal")

for i in bottom_start:
	traj = data_prep.trajectory_set[i][1]
	ax[2].plot(traj.pose_vec[:,0], traj.pose_vec[:,1])
	ax[2].set_title("From bottom canal")

plt.tight_layout()
plt.show()

"""
tests/exploredata.py
eliatrevisan/IA-MPPI-LBM
import matplotlib.pyplot as plt
import numpy as np
import os
import sys
sys.path.append('../')
import src.data_utils.Support as sup
import src.data_utils.AgentContainer as ag_cont
import cv2


home = os.path.expanduser("~")
datapath = "/I-LSTM/data/simulation/"
ped = "6_peds_corridor/"
roboat = "roboat/"
mapname = "map.png"


class Map:
    grid = plt.imread(home+datapath+roboat + 'map.png')
    size = [1715, 881]
    origin = [-78, -40]
    resolution = 0.081
    submap_width = 9.72*2
    submap_height = 9.72*2

def main():

    pedmap = plt.imread(home + datapath + ped + mapname)
    roboatmap = plt.imread(home + datapath + roboat + mapname)

    map = Map()
    agcont = ag_cont.AgentContainer()
    agent_pos = [1.0, 1.0]
    agcont.occupancy_grid.resolution = map.resolution
    agcont.occupancy_grid.gridmap = map.grid
    agcont.occupancy_grid.map_size = map.size

    name = home + datapath + ped + mapname

    grid = agcont.occupancy_grid.getSubmapByCoords(agent_pos[0], agent_pos[1], map.submap_width, map.submap_height)

    print(grid.shape[0], grid.shape[1])

    fig, ax = plt.subplots(1,3)
    sup.plot_grid_roboat(ax[0], map.origin, map.grid, map.resolution, map.size)
    ax[0].plot(agent_pos[0], agent_pos[1], marker='.', markersize=10)
    ax[1].imshow(grid)
    ax[2].imshow( cv2.resize(grid, (60,60)) )
    plt.show()

if __name__ == "__main__":
    main()






tests/inspect_mapdata.py
eliatrevisan/IA-MPPI-LBM

import os
import matplotlib.pyplot as plt
import numpy as np


home = os.path.expanduser("~")
#filepath = "/I-LSTM/data/roboat/"
filepath = "/roboat_data/"

def print_results(scenario):
    data = getDataFromFile(scenario)
    print("Scenario: ", scenario)
    # Total number of identical agents
    print(data.shape)
    print(data[0,:])

    print("Unique Agents: ", len(np.unique(data[:,0])))

    print("Unique time steps: ", len(np.unique(data[:,1])))

def getDataFromFile(scenario):
    """
    Return the data contained in csv file
    """
    data = np.genfromtxt(home + filepath + scenario + "/total_log_open_crossing.csv", delimiter=",")[11:-10,:]
    return data

def main():
    #scenario = "herengracht"
    #print_results(scenario)
    #scenario = "prinsengracht"
    #print_results(scenario)
    #scenario = "bloemgracht"
    #print_results(scenario)
    scenario = "all_data"
    print_results(scenario)



if __name__ == "__main__":
    main()
tests/inspect_results.py
eliatrevisan/IA-MPPI-LBM
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal

x = np.linspace(0, 30, 300)
y = multivariate_normal.pdf(x, mean=0.0, cov=5.0)
fig, ax = plt.subplots()
ax.plot(x,y)
#ax.set_title("Probability Density Function")
ax.set_xlabel("Absolute distance [m]")
ax.set_ylabel("Probability density function")
plt.show()

tests/multivariatepdf.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
sys.path.append('../')
import numpy as np
import argparse
import pylab as pl
import json
import importlib
if sys.version_info[0] < 3:
	print("Using Python " + str(sys.version_info[0]))
	sys.path.append('../src/data_utils')
	sys.path.append('../src/models')
	import DataHandler as dh
	import DataHandlerLSTM as dhlstm
	from plot_utils import *
	import Support as sup
	from utils import *
else:
	print("Using Python " + str(sys.version_info[0]))
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils import DataHandlerRoboat as dhroboat
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
from datetime import datetime
import pickle as pkl
import time
from copy import deepcopy
from multiprocessing.pool import ThreadPool
import colorama
import matplotlib.pyplot as plt 

from colorama import Fore, Style

#os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

pretrained_convnet_path = "../trained_models/autoencoder"
exp_num = 6
data_path = '../data/roboat/'
scenario = 'herengracht'
no_stopping = True

# Hyperparameters
n_epochs = 0

batch_size = 16 #128
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 8
prediction_horizon = 12
prev_horizon = 8

# Model params
# rnn_state_size = 32
# rnn_state_size_lstm_grid = 256
# rnn_state_size_lstm_ped = 128
# rnn_state_ped_size = 16
# rnn_state_size_lstm_concat = 512
# prior_size = 512
# latent_space_size = 256
# x_dim = 512
# fc_hidden_unit_size = 256

rnn_state_size = 32
rnn_state_size_lstm_grid = 64
rnn_state_size_lstm_ped = 64 # If change, change hardcoded output value in VGDNN.py to same value!
rnn_state_size_lstm_concat = 128
latent_space_size = 0

# Not used
rnn_state_ped_size = 1
prior_size = 1
x_dim = 1
fc_hidden_unit_size = 1

# Regularization and training params
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 18
tensorboard_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 36
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 5

print_freq = 200
save_freq = 500
total_training_steps = 50000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = True
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = False
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
#submap_resolution = 0.081
#submap_width = 4.86
#submap_height = 4.86
submap_size = 60 # In pixels
submap_span_real = 10 # In Map coordinates
relative_covariance = 10.0
diversity_update = False
predict_positions = False
warm_start_convnet = True
debug_plotting = False

# Dataset division
train_set = 0.8

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
						help='Path to directory that comprises the model (default="model_name").',
						type=str, default= "RNN")
	parser.add_argument('--model_path',
						help='Path to directory to save the model (default=""../trained_models/"+model_name").',
						type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
						help='Path to directory that comprises the pre-trained convnet model (default=" ").',
						type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
						help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
						type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
						type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
						default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
						default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
						default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
						default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
						default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
						default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
						default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--total_training_steps', help='Number of training steps (default=20000).', type=int, default=total_training_steps)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
						default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
						default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
						default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
						default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
						default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
						default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
						type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
						default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
						default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
						default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
						help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
						help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
						type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
						type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
						type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
						type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
						help='Number of fully connected layer units after LSTM layer (default=64).',
						type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
						default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
						default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
						default=output_dim)
	parser.add_argument('--goal_size', help='Goal dimension (default=).', type=int,
						default=2)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
						default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
						default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
						default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
						default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
						default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
						default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
						default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the network.', type=int,
						default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=sup.str2bool, default=debug_plotting)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
						default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
						type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
						default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
						default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
						default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
						default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
						default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
						default=0)
	#parser.add_argument('--submap_width', help='width of occupancy grid (For initialization of arrays)', type=int, default=submap_width)
	#parser.add_argument('--submap_height', help='height of occupancy grid (for initiatization of arrays)', type=int, default=submap_height)
	parser.add_argument('--submap_size', help='size of submap in pixels (60)', type=int, default=submap_size)
	parser.add_argument('--submap_span_real', help='span of submap occupancy grid', type=int, default=submap_span_real)
	#parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
	#                    default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
						default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
						default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
						default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
						default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
						default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
						default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
						default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
						default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	parser.add_argument('--kl_weight', help='Experimenting with additional weight for KL loss', type=float, default=1)
	parser.add_argument('--relative_covariance', help='Changing covariance of relative positions', type=float, default=relative_covariance)
	args = parser.parse_args()

	return args

args = parse_args()


# Cyclical KL Annealing weight
def frange_cycle_linear(n_iter, start=0.0, stop=1.0,  n_cycle=4, ratio=0.5):
	L = np.ones(n_iter) * stop
	period = n_iter/n_cycle
	step = (stop-start)/(period*ratio) # linear schedule

	for c in range(n_cycle):
		v, i = start, 0
		while v <= stop and (int(i+c*period) < n_iter):
			L[int(i+c*period)] = v
			v += step
			i += 1
	return L

# Enable / Disable GPU
if args.gpu:
	import tensorflow as tf
else:
	os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
	import tensorflow as tf

# Suppress TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

# Create Log and Model Directory to save training model
args.model_path = '../trained_models/' + args.model_name + "/" + str(args.exp_num)
args.log_dir = args.model_path + '/log'
if not os.path.exists(args.log_dir):
	os.makedirs(args.log_dir)
args.dataset = '/' + args.scenario + '.pkl'
model_parameters = {"args": args}

# Check whether model folder exists, otherwise make directory
if not os.path.exists(args.model_path):
	os.makedirs(args.model_path)

# Save Model Parameters
param_file = open(args.model_path + '/model_parameters.pkl', 'wb')
pkl.dump(model_parameters, param_file, protocol=2)  # encoding='latin1'
param_file.close()
with open(args.model_path + '/model_parameters.json', 'w') as f:
	json.dump(args.__dict__,f)

# Specify the number of datasets / scenarios 
n_scenarios = 4

# Create Datahandler class for Herengracht experiment
original_scenario = args.scenario
data_prep_heren = dhlstm.DataHandlerLSTM(args)
map_args_heren = {"file_name": 'map.png',
				"resolution": 0.081,
				"map_size": np.array([1715,881]),
				"map_center": np.array([-78,-40])}
data_prep_heren.processData(**map_args_heren)

# Create Datahandler class for Prinsengracht scenario
args.scenario = "prinsengracht"
data_prep_prins = dhlstm.DataHandlerLSTM(args)
map_args_prins = {"file_name": 'map.png',
				"resolution": 0.091,
				"map_size": np.array([816,816]),
				"map_center": np.array([-37.1, -37.1])} 
data_prep_prins.processData(**map_args_prins)

args.scenario = "bloemgracht"
data_prep_bloem = dhlstm.DataHandlerLSTM(args)
map_args_bloem = {"file_name": 'map.png',
				"resolution": 0.22,
				"map_size": np.array([1469,928]),
				"map_center": np.array([-161.59,-102.08])} 
data_prep_bloem.processData(**map_args_bloem)

args.scenario = "open_crossing"
data_prep_open = dhlstm.DataHandlerLSTM(args)
map_args_open = {"file_name": 'map.png',
				"resolution": 0.6,
				"map_size": np.array([100,100]),
				"map_center": np.array([-30,-30])} 
data_prep_open.processData(**map_args_open)

args.scenario = original_scenario

# Initialize main data handler
#dpreps = [data_prep, data_prep_2, data_prep_3, data_prep_4]
#data_roboat = dhroboat.DataHandlerRoboat(args)


import matplotlib.cm as cm

def plot_map(data, map, ax):
    ax.imshow(data[:,:,0],
               extent = (map["map_center"][0], map["map_center"][0] + map["map_size"][0] * map["resolution"],
                         map["map_center"][1], map["map_center"][1] + map["map_size"][1] * map["resolution"]),
               cmap=cm.gray_r)
    ax.set_facecolor('#262626')
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    ax.yaxis.set_label_coords(-0.08, .5)
    ax.xaxis.set_label_coords(0.5, -0.09)

home = os.path.expanduser('~')
heren_map = pl.imread(home + '/I-LSTM/data/roboat/herengracht/map.png')
open_map = pl.imread(home + '/I-LSTM/data/roboat/open_crossing/map.png')
bloem_map = pl.imread(home + '/I-LSTM/data/roboat/bloemgracht/map.png')
prins_map = pl.imread(home + '/I-LSTM/data/roboat/prinsengracht/map.png')

fig, ax = plt.subplots()

plot_map(heren_map, map_args_heren, ax)

traj = data_prep_heren.trajectory_set[27][1]
ax.scatter(traj.pose_vec[:,0], traj.pose_vec[:,1], s=5, color='b')
n_other = len(traj.other_agents_positions[0])

for n in range(n_other):
    for t in range(len(traj.other_agents_positions)):
        ax.scatter(traj.other_agents_positions[t][n][0], traj.other_agents_positions[t][n][1], s=5, color = 'r', )

#ax.set_title("Complex Interactions")
#ax.set_xlim((-20,50))
#ax.set_ylim((-40,30))
ax.set_xlim((-45,15))
ax.set_ylim((-30,30))
plt.show()
tests/plot_trajectories.py
eliatrevisan/IA-MPPI-LBM
import sys
import os
import csv
import numpy as np
import matplotlib.pyplot as plt
sys.path.append('../')


home = os.path.expanduser("~")
directory = "/I-LSTM/src/"
model = "VGDNN"
filename = "_summary.csv"

def main():
    models, labels, results = getResults()

    model_ids = []
    for i in range(1,len(sys.argv)):
        id = sys.argv[i]
        model_ids.append(id)
    
    if len(model_ids) > 0:
        models, results = getResultsfromArgs(models, results, model_ids)

    fig, ax =  plt.subplots(2)
    printResults(ax, models, results,labels)

    #print(results)
    #fig.tight_layout()

    plt.show()

def getResultsfromArgs(models, results, model_ids):
    
    for i in range(len(models)):
        models[i] = models[i].split("_",1)[1]
    idx = []
    print(len(models), len(results))
    for i in range(len(model_ids)):
        idx.append(np.where(models == model_ids[i])[0][0])

    print(len(models), len(results))
    
    return models[idx], results[idx]

def getResults():
    """
    Return the data contained in csv file
    """
    data = np.genfromtxt(home + directory + model + filename, delimiter=",", skip_header=1, usecols=(2,3,4,5,6,7,8,9,10))
    labels = np.genfromtxt(home + directory + model + filename, delimiter=",", dtype=str)[0,:]
    models = np.genfromtxt(home + directory + model + filename, delimiter=",", usecols=0, skip_header=1, dtype=str)

    data[:,2] 

    for i in range(len(data[:,2])):
        data[i,2] = np.round(data[i,2], 2)
        data[i,3] = np.round(data[i,3], 2)
        data[i,4] = np.round(data[i,4], 2)

    return models, labels, data

def printResults(ax, models, results,labels):

    MSE = results[:,2]
    FDE = results[:,3]

    N = len(MSE)
    ind = np.arange(N)  
    barwidth = 0.25

    ax[0].bar(ind + barwidth * 0.5, MSE, align='center', width=barwidth, label='MSE')
    ax[0].bar(ind + barwidth *1.5, FDE, align='center', width=barwidth, label='FDE')

    for index, value in enumerate(MSE):
        ax[0].text(index , value + 0.1,
             str(round(value,2)), fontsize=8)
    for index, value in enumerate(FDE):
        ax[0].text(index + barwidth * 1.0, value + 0.1,
             str(round(value,2)), fontsize=8)


    ax[0].legend()
    ax[0].set_title('Mean Squared Error and Final Displacement Error')
    ax[0].set_xlabel('Models')
    ax[0].set_ylabel('Error')
    ax[0].set_xticks(ind+barwidth)
    ax[0].set_xticklabels( models )
    #print(labels)

    table = list(results[:,[1,4,5,6,7,8]])#,3,4,5,6]])
    table = np.reshape(table, (-1,6)).transpose()
    labels = labels[[3,6,7,8,9,10]]
    
    #print(labels)
    #print(table)

    table = ax[1].table(cellText=table, rowLabels=labels, colLabels=models, loc='center')
    table.auto_set_font_size(False)
    table.set_fontsize(12)
    
    ax[1].axis('off')
    
    return


if __name__ == "__main__":
    main()

tests/quantitative_results.py
eliatrevisan/IA-MPPI-LBM
import sys
sys.path.append('../')
import numpy as np
if sys.version_info[0] < 3:
	sys.path.append('../src/data_utils')
	import DataHandler as dh
else:
	from src.data_utils import DataHandler as dh
	from src.data_utils import DataHandlerLSTM as dhlstm
	from src.data_utils.plot_utils import *
	from src.data_utils import Support as sup
	from src.data_utils.utils import *
from parser import *
import argparse
"""

def parse_args():
	parser = argparse.ArgumentParser(description='data recording')

	parser.add_argument('--data_path', help='Path to directory that comprises the training data (default=" ").',
											type=str, default='../data/')
	parser.add_argument('--save_path', help='Path to directory that saves pickle data (default=" ").',
											type=str,
											default='../data/data.pickle')
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=16)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--truncated_backprop_length', help='truncated back propagation length', type=int, default=10)
	parser.add_argument('--output_sequence_length', help='output sequence length', type=int, default=1)
	parser.add_argument('--cmd_vector_dim', help='robot command vector dimension', type=int, default=2)
	parser.add_argument('--goal_x', help='x coordinate of robot goal state', type=float, default=23.00)
	parser.add_argument('--goal_y', help='y coordinate of robot goal state', type=float, default=0.00)
	parser.add_argument('--prediction_horizon', help='prediction horizon', type=int, default=15)
	parser.add_argument('--bag_file', help='csv file containg control commands (default=" ").',
											type=str, default='data_2019-08-20-12-29-52.bag')
	parser.add_argument('--distance_threshold', help='Distance threshold to start new trajectory', type=float, default=1)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").',
											type=str, default='../config/topics.yaml')



	args = parser.parse_args()

	return args

args = parse_args()

data_prep = dh.DataHandler(args)

data_prep.process_simulation_data_()

for i in range(1000):
	batch_cmd, batch_pred_cmd, batch_robot_state, new_epoch = data_prep.getBatch()
	print("New epoch: " + str(new_epoch))
	if np.any(data_prep.sequence_reset):
		for sequence_idx in range(data_prep.sequence_reset.shape[0]):
			if data_prep.sequence_reset[sequence_idx] == 1:
				print("Reset state: " + str(sequence_idx))

print("successfully obtained training data batch")

print("#########")

data_prep.save_data()

print("pickle data saved")

"""

pretrained_convnet_path = "../trained_models/autoencoder_with_ped"

exp_num = 2
data_path = '../data/'
scenario = 'simulation/roboat'

# Hyperparameters
n_epochs = 2

batch_size = 16
regularization_weight = 0.0001

# Time parameters
truncated_backprop_length = 8
prediction_horizon = 12
prev_horizon = 8

rnn_state_size = 32
rnn_state_size_lstm_grid = 256
rnn_state_size_lstm_ped = 128
rnn_state_ped_size = 16
rnn_state_size_lstm_concat = 512
prior_size = 512
latent_space_size = 256
x_dim = 512
fc_hidden_unit_size = 256
learning_rate_init = 0.001
beta_rate_init = 0.01
keep_prob = 1.0
dropout = False
reg = 1e-4
n_mixtures = 1  # USE ZERO FOR MSE MODEL
grads_clip = 1.0
n_other_agents = 18
tensorboard_logging = True

# Model parameters
input_dim = 4  # [vx, vy]
input_state_dim = 2  # [vx, vy]
output_dim = 2  # data state dimension
output_pred_state_dim = 4  # ux uy simgax sigmay
pedestrian_vector_dim = 4
pedestrian_vector_state_dim = 2
cmd_vector_dim = 2
pedestrian_radius = 0.3
max_range_ped_grid = 5

print_freq = 200
save_freq = 500
total_training_steps = 20000
dt = 0.4

warmstart_model = False
pretrained_convnet = False
pretained_encoder = False
multipath = False
real_world_data = False
end_to_end = True
agents_on_grid = False
rotated_grid = False
centered_grid = True
noise = False
normalize_data = False
real_world_data = False
regulate_log_loss = False
# Map parameters
submap_resolution = 0.081
submap_width = 4.86
submap_height = 4.86
diversity_update = True
predict_positions = False
warm_start_convnet = True
debug_plotting = False

# Dataset division
train_set = 0.8

def parse_args():
	parser = argparse.ArgumentParser(description='LSTM model training')

	parser.add_argument('--model_name',
											help='Path to directory that comprises the model (default="model_name").',
											type=str, default= "RNN")
	parser.add_argument('--model_path',
											help='Path to directory to save the model (default=""../trained_models/"+model_name").',
											type=str, default='../trained_models/')
	parser.add_argument('--pretrained_convnet_path',
											help='Path to directory that comprises the pre-trained convnet model (default=" ").',
											type=str, default=pretrained_convnet_path)
	parser.add_argument('--log_dir',
											help='Path to the log directory of the model (default=""../trained_models/"+model_name").',
											type=str, default="\log")
	parser.add_argument('--scenario', help='Scenario of the dataset (default="").',
											type=str, default=scenario)
	parser.add_argument('--real_world_data', help='Real world dataset (default=True).', type=sup.str2bool,
											default=real_world_data)
	parser.add_argument('--data_path', help='Path to directory that saves pickle data (default=" ").', type=str,
											default=data_path)
	parser.add_argument('--dataset', help='Dataset pkl file', type=str,
											default= scenario + '.pkl')
	parser.add_argument('--data_handler', help='Datahandler class needed to load the data', type=str,
											default='LSTM')
	parser.add_argument('--warmstart_model', help='Restore from pretained model (default=False).', type=bool,
											default=warmstart_model)
	parser.add_argument('--warm_start_convnet', help='Restore from pretained convnet model (default=False).', type=bool,
											default=warm_start_convnet)
	parser.add_argument('--dt', help='Data samplig time (default=0.3).', type=float,
											default=dt)
	parser.add_argument('--n_epochs', help='Number of epochs (default=10000).', type=int, default=n_epochs)
	parser.add_argument('--total_training_steps', help='Number of training steps (default=20000).', type=int, default=total_training_steps)
	parser.add_argument('--batch_size', help='Batch size for training (default=32).', type=int, default=batch_size)
	parser.add_argument('--regularization_weight', help='Weight scaling of regularizer (default=0.01).', type=float,
											default=regularization_weight)
	parser.add_argument('--keep_prob', help='Dropout (default=0.8).', type=float,
											default=keep_prob)
	parser.add_argument('--learning_rate_init', help='Initial learning rate (default=0.005).', type=float,
											default=learning_rate_init)
	parser.add_argument('--beta_rate_init', help='Initial beta rate (default=0.005).', type=float,
											default=beta_rate_init)
	parser.add_argument('--dropout', help='Enable Dropout', type=sup.str2bool,
											default=dropout)
	parser.add_argument('--grads_clip', help='Gridient clipping (default=10.0).', type=float,
											default=grads_clip)
	parser.add_argument('--truncated_backprop_length', help='Backpropagation length during training (default=5).',
											type=int, default=truncated_backprop_length)
	parser.add_argument('--prediction_horizon', help='Length of predicted sequences (default=10).', type=int,
											default=prediction_horizon)
	parser.add_argument('--prev_horizon', help='Previous seq length.', type=int,
											default=prev_horizon)
	parser.add_argument('--rnn_state_size', help='Number of RNN / LSTM units (default=16).', type=int,
											default=rnn_state_size)
	parser.add_argument('--rnn_state_size_lstm_ped',
											help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
											type=int, default=rnn_state_size_lstm_ped)
	parser.add_argument('--rnn_state_ped_size',
											help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
											type=int, default=rnn_state_ped_size)
	parser.add_argument('--rnn_state_size_lstm_grid',
											help='Number of RNN / LSTM units of the grid lstm layer (default=32).',
											type=int, default=rnn_state_size_lstm_grid)
	parser.add_argument('--rnn_state_size_lstm_concat',
											help='Number of RNN / LSTM units of the concatenation lstm layer (default=32).',
											type=int, default=rnn_state_size_lstm_concat)
	parser.add_argument('--prior_size', help='prior_size',
											type=int, default=prior_size)
	parser.add_argument('--latent_space_size', help='latent_space_size',
											type=int, default=latent_space_size)
	parser.add_argument('--x_dim', help='x_dim',
											type=int, default=x_dim)
	parser.add_argument('--fc_hidden_unit_size',
											help='Number of fully connected layer units after LSTM layer (default=64).',
											type=int, default=fc_hidden_unit_size)
	parser.add_argument('--input_state_dim', help='Input state dimension (default=).', type=int,
											default=input_state_dim)
	parser.add_argument('--input_dim', help='Input state dimension (default=).', type=float,
											default=input_dim)
	parser.add_argument('--output_dim', help='Output state dimension (default=).', type=float,
											default=output_dim)
	parser.add_argument('--goal_size', help='Goal dimension (default=).', type=int,
											default=2)
	parser.add_argument('--output_pred_state_dim', help='Output prediction state dimension (default=).', type=int,
											default=output_pred_state_dim)
	parser.add_argument('--cmd_vector_dim', help='Command control dimension.', type=int,
											default=cmd_vector_dim)
	parser.add_argument('--n_mixtures', help='Number of modes (default=).', type=int,
											default=n_mixtures)
	parser.add_argument('--pedestrian_vector_dim', help='Number of angular grid sectors (default=72).', type=int,
											default=pedestrian_vector_dim)
	parser.add_argument('--pedestrian_vector_state_dim', help='Number of angular grid sectors (default=2).', type=int,
											default=pedestrian_vector_state_dim)
	parser.add_argument('--max_range_ped_grid', help='Maximum pedestrian distance (default=2).', type=float,
											default=max_range_ped_grid)
	parser.add_argument('--pedestrian_radius', help='Pedestrian radius (default=0.3).', type=float,
											default=pedestrian_radius)
	parser.add_argument('--n_other_agents', help='Number of other agents incorporated in the network.', type=int,
											default=n_other_agents)
	parser.add_argument('--debug_plotting', help='Plotting for debugging (default=False).', type=sup.str2bool, default=debug_plotting)
	parser.add_argument('--print_freq', help='Print frequency of training info (default=100).', type=int,
											default=print_freq)
	parser.add_argument('--save_freq', help='Save frequency of the temporary model during training. (default=20k).',
											type=int, default=save_freq)
	parser.add_argument('--exp_num', help='Experiment number', type=int, default=exp_num)
	parser.add_argument('--noise', help='Likelihood? (default=True).', type=sup.str2bool,
											default=False)
	parser.add_argument('--agents_on_grid', help='Likelihood? (default=True).', type=sup.str2bool,
											default=agents_on_grid)
	parser.add_argument('--normalize_data', help='Normalize? (default=False).', type=sup.str2bool,
											default=normalize_data)
	parser.add_argument('--rotated_grid', help='Rotate grid? (default=False).', type=sup.str2bool,
											default=rotated_grid)
	parser.add_argument('--centered_grid', help='Center grid? (default=False).', type=sup.str2bool,
											default=centered_grid)
	parser.add_argument('--sigma_bias', help='Percentage of the dataset used for trainning', type=float,
											default=0)
	parser.add_argument('--submap_width', help='width of occupancy grid', type=int, default=submap_width)
	parser.add_argument('--submap_height', help='height of occupancy grid', type=int, default=submap_height)
	parser.add_argument('--submap_resolution', help='Map resolution.', type=float,
											default=submap_resolution)
	parser.add_argument('--min_buffer_size', help='Minimum buffer size (default=1000).', type=int, default=1000)
	parser.add_argument('--max_buffer_size', help='Maximum buffer size (default=100k).', type=int, default=100000)
	parser.add_argument('--max_trajectories', help='maximum number of trajectories to be recorded', type=int, default=30)
	parser.add_argument('--end_to_end', help='End to end trainning.', type=sup.str2bool,
											default=False)
	parser.add_argument('--predict_positions', help='predict_positions.', type=sup.str2bool,
											default=predict_positions)
	parser.add_argument('--gpu', help='Enable GPU training.', type=sup.str2bool,
											default=False)
	parser.add_argument('--sequence_info', help='Use relative info for other agents.', type=sup.str2bool,
											default=False)
	parser.add_argument('--others_info', help='Use relative info for other agents.', type=str,
											default="none")
	parser.add_argument('--regulate_log_loss', help='Enable GPU training.', type=sup.str2bool,
											default=regulate_log_loss)
	parser.add_argument('--diversity_update', help='diversity_update', type=sup.str2bool,
											default=diversity_update)
	parser.add_argument('--topics_config', help='yaml file containg subscription topics (default=" ").', type=str,
											default='../config/topics.yaml')
	parser.add_argument('--min_pos_x', help='min_pos_x', type=float, default=-1)
	parser.add_argument('--min_pos_y', help='min_pos_y', type=float, default=-1)
	parser.add_argument('--max_pos_x', help='max_pos_x', type=float, default=1)
	parser.add_argument('--max_pos_y', help='max_pos_y', type=float, default=1)
	parser.add_argument('--min_vel_x', help='min_vel_x', type=float, default=-1)
	parser.add_argument('--min_vel_y', help='min_vel_y', type=float, default=-1)
	parser.add_argument('--max_vel_x', help='max_vel_x', type=float, default=1)
	parser.add_argument('--max_vel_y', help='max_vel_y', type=float, default=1)
	parser.add_argument('--sx_vel', help='sx_vel', type=float, default=1)
	parser.add_argument('--sy_vel', help='sy_vel', type=float, default=1)
	parser.add_argument('--sx_pos', help='sx_pos', type=float, default=1)
	parser.add_argument('--sy_pos', help='sy_pos', type=float, default=1)
	parser.add_argument('--train_set', help='Percentage of the dataset used for training', type=float, default=train_set)
	args = parser.parse_args()

	return args


args = parse_args()

data_prep = dhlstm.DataHandlerLSTM(args)
# Only used to create a map from png
# Make sure this parameters are correct otherwise it will fail training and plotting the results
map_args = {"file_name": 'map.png',
							"resolution": 0.081,
							"map_size": np.array([1715,881]),
				"map_center": np.array([-78,-40])} # THIS IS HARDCODED FOR SOME REASON AND NOT EXTRACTED FROM MAP.JSON ?
# Load dataset
data_prep.processData(**map_args)

#traj = data_prep.trajectory_set[0][1]

#print(traj.other_agents_positions) # Full traj (pos) of other agents

batch_x, batch_vel, batch_pos,batch_goal,batch_grid, batch_ped_grid, batch_y,batch_pos_target, other_agents_pos, new_epoch = data_prep.getBatch()

print(batch_vel[0].shape)
tests/test_data.py
eliatrevisan/IA-MPPI-LBM