<reponame>USC-CSCI527-Fall2021/autonomous_car_navigation
import numpy as np
import tensorflow as tf
from tensorflow import keras
import gym
import gym_carla
import carla
import scipy.signal
import time
import cv2

import tensorflow as tf
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.compat.v1.InteractiveSession(config=config)

num_actions = 3
observation_dimensions = (128, 128, 3)

def discounted_cumulative_sums(x, discount):
    # Discounted cumulative sums of vectors for computing rewards-to-go and advantage estimates
    return scipy.signal.lfilter([1], [1, float(-discount)], x[::-1], axis=0)[::-1]


class Buffer:
    # Buffer for storing trajectories
    def __init__(self, observation_dimensions, size, gamma=0.99, lam=0.95):
        # Buffer initialization
        self.observation_buffer = np.zeros(
            (size, *observation_dimensions), dtype=np.float32
        )
        self.action_buffer = np.zeros(size, dtype=np.int32)
        self.advantage_buffer = np.zeros(size, dtype=np.float32)
        self.reward_buffer = np.zeros(size, dtype=np.float32)
        self.return_buffer = np.zeros(size, dtype=np.float32)
        self.value_buffer = np.zeros(size, dtype=np.float32)
        self.logprobability_buffer = np.zeros(size, dtype=np.float32)
        self.gamma, self.lam = gamma, lam
        self.pointer, self.trajectory_start_index = 0, 0

    def store(self, observation, action, reward, value, logprobability):
        # Append one step of agent-environment interaction
        self.observation_buffer[self.pointer] = observation
        self.action_buffer[self.pointer] = action
        self.reward_buffer[self.pointer] = reward
        self.value_buffer[self.pointer] = value
        self.logprobability_buffer[self.pointer] = logprobability
        self.pointer += 1

    def finish_trajectory(self, last_value=0):
        # Finish the trajectory by computing advantage estimates and rewards-to-go
        path_slice = slice(self.trajectory_start_index, self.pointer)
        rewards = np.append(self.reward_buffer[path_slice], last_value)
        values = np.append(self.value_buffer[path_slice], last_value)

        deltas = rewards[:-1] + self.gamma * values[1:] - values[:-1]

        self.advantage_buffer[path_slice] = discounted_cumulative_sums(
            deltas, self.gamma * self.lam
        )
        self.return_buffer[path_slice] = discounted_cumulative_sums(
            rewards, self.gamma
        )[:-1]
        
        self.trajectory_start_index = self.pointer

    def get(self):
        # Get all data of the buffer and normalize the advantages
        self.pointer, self.trajectory_start_index = 0, 0
        advantage_mean, advantage_std = (
            np.mean(self.advantage_buffer),
            np.std(self.advantage_buffer),
        )
        self.advantage_buffer = (self.advantage_buffer - advantage_mean) / advantage_std
        return (
            self.observation_buffer,
            self.action_buffer,
            self.advantage_buffer,
            self.return_buffer,
            self.logprobability_buffer,
        )


def mlp(x, sizes, activation=tf.tanh, output_activation=None):
    # Build a feedforward neural network
    x = keras.layers.Conv2D(filters=32, kernel_size=(3,3), kernel_initializer='he_normal', 
                            padding='same', activation="relu")(x)
    x = keras.layers.AveragePooling2D( (2,2) )(x)
    x = keras.layers.Conv2D(filters=8, kernel_size=(3,3), kernel_initializer='he_normal', 
                            padding='same', activation="relu")(x)
    x = keras.layers.Flatten()(x)
    for size in sizes[:-1]:
        x = keras.layers.Dense(units=size, activation=activation)(x)
    return keras.layers.Dense(units=sizes[-1], activation=output_activation)(x)


def logprobabilities(logits, a):
    # Compute the log-probabilities of taking actions a by using the logits (i.e. the output of the actor)
    logprobabilities_all = tf.nn.log_softmax(logits)
    logprobability = tf.reduce_sum(
        tf.one_hot(a, num_actions) * logprobabilities_all, axis=1
    )
    return logprobability


# Sample action from actor
@tf.function
def sample_action(actor, observation):
    logits = actor(observation)
    action = tf.squeeze(tf.random.categorical(logits, 1), axis=1)
    return logits, action

# Train the policy by maxizing the PPO-Clip objective
@tf.function
def train_policy(
    observation_buffer, action_buffer, logprobability_buffer, advantage_buffer
):
    with tf.GradientTape() as tape:  # Record operations for automatic differentiation.
        ratio = tf.exp(
            logprobabilities(actor(observation_buffer), action_buffer)
            - logprobability_buffer
        )
        min_advantage = tf.where(
            advantage_buffer > 0,
            (1 + clip_ratio) * advantage_buffer,
            (1 - clip_ratio) * advantage_buffer,
        )

        policy_loss = -tf.reduce_mean(
            tf.minimum(ratio * advantage_buffer, min_advantage)
        )
        
    policy_grads = tape.gradient(policy_loss, actor.trainable_variables)
    policy_optimizer.apply_gradients(zip(policy_grads, actor.trainable_variables))
    kl = tf.reduce_mean(
        logprobability_buffer
        - logprobabilities(actor(observation_buffer), action_buffer)
    )
    kl = tf.reduce_sum(kl)
    return kl


# Train the value function by regression on mean-squared error
@tf.function
def train_value_function(observation_buffer, return_buffer):
    with tf.GradientTape() as tape:  # Record operations for automatic differentiation.
        value_loss = tf.reduce_mean((return_buffer - critic(observation_buffer)) ** 2)
    value_grads = tape.gradient(value_loss, critic.trainable_variables)
    value_optimizer.apply_gradients(zip(value_grads, critic.trainable_variables))

# Hyperparameters of the PPO algorithm
steps_per_epoch = 1000
epochs = 500
gamma = 0.99
clip_ratio = 0.2
policy_learning_rate = 3e-4
value_function_learning_rate = 1e-3
train_policy_iterations = 1
train_value_iterations = 1
lam = 0.97
target_kl = 0.01
hidden_sizes = (128, 64)

# True if you want to render the environment
render = False

params = {
    'number_of_vehicles': 0,
    'number_of_walkers': 0,
    'display_size': 250,  # screen size of bird-eye render
    'display_height' : 512,
    'display_main': False,
    'max_past_step': 1,  # the number of past steps to draw
    'dt': 0.1,  # time interval between two frames
    'discrete': False,  # whether to use discrete control space
    'discrete_acc': [1.0, 0.0, 1.0],  # discrete value of accelerations
    'discrete_steer': [-1, 0, 1],  # discrete value of steering angles
    'continuous_accel_range': [-3.0, 3.0],  # continuous acceleration range
    'continuous_steer_range': [-0.2, 0.2],  # continuous steering angle range
    'ego_vehicle_filter': 'vehicle.lincoln*',  # filter for defining ego vehicle
    'address':'192.168.1.173',
    'port': 2000, # connection port
    'town': 'Town02', # which town to simulate
    'task_mode': 'random',  # mode of the task, [random, roundabout (only for Town03)]
    'max_time_episode': 2000,  # maximum timesteps per episode
    'max_waypt': 12,  # maximum number of waypoints
    'obs_range': 32,  # observation range (meter)
    'lidar_bin': 0.125,  # bin size of lidar sensor (meter)
    'd_behind': 12,  # distance behind the ego vehicle (meter)
    'out_lane_thres': 5.0,  # threshold for out of lane
    'desired_speed': 8,  # desired speed (m/s)
    'max_ego_spawn_times': 200,  # maximum times to spawn ego vehicle
    'display_route': True,  # whether to render the desired route
    'pixor_size': 64,  # size of the pixor labels
    'pixor': False,  # whether to output PIXOR observation
  }

  buffer = Buffer(observation_dimensions, steps_per_epoch)

observation_input = keras.Input(shape=observation_dimensions, dtype=tf.float32)
logits = mlp(observation_input, list(hidden_sizes)+[num_actions], tf.tanh, None)
actor = keras.Model(inputs=observation_input, outputs=logits)
value = tf.squeeze(
    mlp(observation_input, list(hidden_sizes)+[1], tf.tanh, None), axis=1
)
critic = keras.Model(inputs=observation_input, outputs=value)

# Initialize the policy and the value function optimizers
policy_optimizer = keras.optimizers.Adam(learning_rate=policy_learning_rate)
value_optimizer = keras.optimizers.Adam(learning_rate=value_function_learning_rate)

actor.summary()
tf.config.run_functions_eagerly(True)

def transform(img):
    return img[76:204,76:204,:]/255

epoch = 1


env = gym.make('carla-v0', params=params)
observation =  env.reset() # select_env()
episode_return, episode_length = 0, 0
for _ in range(20): observation, _, _, _ = env.step([1,0])
observation = transform( observation['birdeye'] )
    
while epoch != epochs:
    # Initialize the sum of the returns, lengths and number of episodes for each epoch
    sum_return = 0
    sum_length = 0
    num_episodes = 0
        
    # Iterate over the steps of each epoch
    for t in range(steps_per_epoch): #while True
        # Get the logits, action, and take one step in the environment
        logits, action = sample_action(actor, observation.reshape(1, *observation_dimensions))
        observation_new, reward, done, _ = env.step([1, action[0].numpy()-1])
        observation_new = transform(observation_new['birdeye'])
        episode_return += reward
        episode_length += 1

        # Get the value and log-probability of the action
        value_t = critic(observation.reshape(1, *observation_dimensions))
        logprobability_t = logprobabilities(logits, action)

        # Store obs, act, rew, v_t, logp_pi_t
        buffer.store(observation, action, reward, value_t, logprobability_t)

        # Update the observation
        observation = observation_new

        # Finish trajectory if reached to a terminal state
        terminal = done
        if terminal or (t == steps_per_epoch - 1):
            last_value = 0 if done else critic(observation.reshape(1, *observation_dimensions))
            buffer.finish_trajectory(last_value)
            sum_return += episode_return
            sum_length += episode_length
            num_episodes += 1
            observation =  env.reset()#select_env()
            episode_return, episode_length = 0, 0
            for _ in range(20): observation, _, _, _ = env.step([1,0])
            observation = transform(observation['birdeye'])
            
            
    # Print mean return and length for each epoch
    with open("logs_no_vae.txt", "a") as log:
        print(
            f"{epoch},{sum_return / num_episodes},{sum_length / num_episodes},{num_episodes}",
            file = log
        )
    print(
        f" Epoch: {epoch}. Mean Return: {sum_return / num_episodes}. Mean Length: {sum_length / num_episodes}"
    )
    
    # Get values from the buffer
    (
        observation_buffer,
        action_buffer,
        advantage_buffer,
        return_buffer,
        logprobability_buffer,
    ) = buffer.get()

    # Update the policy and implement early stopping using KL divergence
    for _ in range(train_policy_iterations):
        kl = train_policy(
            observation_buffer, action_buffer, logprobability_buffer, advantage_buffer
        )
        if kl > 1.5 * target_kl:
            # Early Stopping
            break

    # Update the value function
    for _ in range(train_value_iterations):
        train_value_function(observation_buffer, return_buffer)

    
    tm = int(time.time())
    actor.save_weights(f"models/novae_actor_{tm}.h5")
    critic.save_weights(f"models/novae_critic_{tm}.h5")
    epoch += 1
import numpy as np
import tensorflow as tf
from tensorflow import keras
import gym
import gym_carla
import carla
import scipy.signal
import time
import cv2

import tensorflow as tf
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.compat.v1.InteractiveSession(config=config)

num_actions = 3
observation_dimensions = (128, 128, 3)

def discounted_cumulative_sums(x, discount):
    # Discounted cumulative sums of vectors for computing rewards-to-go and advantage estimates
    return scipy.signal.lfilter([1], [1, float(-discount)], x[::-1], axis=0)[::-1]


class Buffer:
    # Buffer for storing trajectories
    def __init__(self, observation_dimensions, size, gamma=0.99, lam=0.95):
        # Buffer initialization
        self.observation_buffer = np.zeros(
            (size, *observation_dimensions), dtype=np.float32
        )
        self.action_buffer = np.zeros(size, dtype=np.int32)
        self.advantage_buffer = np.zeros(size, dtype=np.float32)
        self.reward_buffer = np.zeros(size, dtype=np.float32)
        self.return_buffer = np.zeros(size, dtype=np.float32)
        self.value_buffer = np.zeros(size, dtype=np.float32)
        self.logprobability_buffer = np.zeros(size, dtype=np.float32)
        self.gamma, self.lam = gamma, lam
        self.pointer, self.trajectory_start_index = 0, 0

    def store(self, observation, action, reward, value, logprobability):
        # Append one step of agent-environment interaction
        self.observation_buffer[self.pointer] = observation
        self.action_buffer[self.pointer] = action
        self.reward_buffer[self.pointer] = reward
        self.value_buffer[self.pointer] = value
        self.logprobability_buffer[self.pointer] = logprobability
        self.pointer += 1

    def finish_trajectory(self, last_value=0):
        # Finish the trajectory by computing advantage estimates and rewards-to-go
        path_slice = slice(self.trajectory_start_index, self.pointer)
        rewards = np.append(self.reward_buffer[path_slice], last_value)
        values = np.append(self.value_buffer[path_slice], last_value)

        deltas = rewards[:-1] + self.gamma * values[1:] - values[:-1]

        self.advantage_buffer[path_slice] = discounted_cumulative_sums(
            deltas, self.gamma * self.lam
        )
        self.return_buffer[path_slice] = discounted_cumulative_sums(
            rewards, self.gamma
        )[:-1]
        
        self.trajectory_start_index = self.pointer

    def get(self):
        # Get all data of the buffer and normalize the advantages
        self.pointer, self.trajectory_start_index = 0, 0
        advantage_mean, advantage_std = (
            np.mean(self.advantage_buffer),
            np.std(self.advantage_buffer),
        )
        self.advantage_buffer = (self.advantage_buffer - advantage_mean) / advantage_std
        return (
            self.observation_buffer,
            self.action_buffer,
            self.advantage_buffer,
            self.return_buffer,
            self.logprobability_buffer,
        )


def mlp(x, sizes, activation=tf.tanh, output_activation=None):
    # Build a feedforward neural network
    x = keras.layers.Conv2D(filters=32, kernel_size=(3,3), kernel_initializer='he_normal', 
                            padding='same', activation="relu")(x)
    x = keras.layers.AveragePooling2D( (2,2) )(x)
    x = keras.layers.Conv2D(filters=8, kernel_size=(3,3), kernel_initializer='he_normal', 
                            padding='same', activation="relu")(x)
    x = keras.layers.Flatten()(x)
    for size in sizes[:-1]:
        x = keras.layers.Dense(units=size, activation=activation)(x)
    return keras.layers.Dense(units=sizes[-1], activation=output_activation)(x)


def logprobabilities(logits, a):
    # Compute the log-probabilities of taking actions a by using the logits (i.e. the output of the actor)
    logprobabilities_all = tf.nn.log_softmax(logits)
    logprobability = tf.reduce_sum(
        tf.one_hot(a, num_actions) * logprobabilities_all, axis=1
    )
    return logprobability


# Sample action from actor
@tf.function
def sample_action(actor, observation):
    logits = actor(observation)
    action = tf.squeeze(tf.random.categorical(logits, 1), axis=1)
    return logits, action

# Train the policy by maxizing the PPO-Clip objective
@tf.function
def train_policy(
    observation_buffer, action_buffer, logprobability_buffer, advantage_buffer
):
    with tf.GradientTape() as tape:  # Record operations for automatic differentiation.
        ratio = tf.exp(
            logprobabilities(actor(observation_buffer), action_buffer)
            - logprobability_buffer
        )
        min_advantage = tf.where(
            advantage_buffer > 0,
            (1 + clip_ratio) * advantage_buffer,
            (1 - clip_ratio) * advantage_buffer,
        )

        policy_loss = -tf.reduce_mean(
            tf.minimum(ratio * advantage_buffer, min_advantage)
        )
        
    policy_grads = tape.gradient(policy_loss, actor.trainable_variables)
    policy_optimizer.apply_gradients(zip(policy_grads, actor.trainable_variables))
    kl = tf.reduce_mean(
        logprobability_buffer
        - logprobabilities(actor(observation_buffer), action_buffer)
    )
    kl = tf.reduce_sum(kl)
    return kl


# Train the value function by regression on mean-squared error
@tf.function
def train_value_function(observation_buffer, return_buffer):
    with tf.GradientTape() as tape:  # Record operations for automatic differentiation.
        value_loss = tf.reduce_mean((return_buffer - critic(observation_buffer)) ** 2)
    value_grads = tape.gradient(value_loss, critic.trainable_variables)
    value_optimizer.apply_gradients(zip(value_grads, critic.trainable_variables))

# Hyperparameters of the PPO algorithm
steps_per_epoch = 1000
epochs = 500
gamma = 0.99
clip_ratio = 0.2
policy_learning_rate = 3e-4
value_function_learning_rate = 1e-3
train_policy_iterations = 1
train_value_iterations = 1
lam = 0.97
target_kl = 0.01
hidden_sizes = (128, 64)

# True if you want to render the environment
render = False

params = {
    'number_of_vehicles': 0,
    'number_of_walkers': 0,
    'display_size': 250,  # screen size of bird-eye render
    'display_height' : 512,
    'display_main': False,
    'max_past_step': 1,  # the number of past steps to draw
    'dt': 0.1,  # time interval between two frames
    'discrete': False,  # whether to use discrete control space
    'discrete_acc': [1.0, 0.0, 1.0],  # discrete value of accelerations
    'discrete_steer': [-1, 0, 1],  # discrete value of steering angles
    'continuous_accel_range': [-3.0, 3.0],  # continuous acceleration range
    'continuous_steer_range': [-0.2, 0.2],  # continuous steering angle range
    'ego_vehicle_filter': 'vehicle.lincoln*',  # filter for defining ego vehicle
    'address':'192.168.1.173',
    'port': 2000, # connection port
    'town': 'Town02', # which town to simulate
    'task_mode': 'random',  # mode of the task, [random, roundabout (only for Town03)]
    'max_time_episode': 2000,  # maximum timesteps per episode
    'max_waypt': 12,  # maximum number of waypoints
    'obs_range': 32,  # observation range (meter)
    'lidar_bin': 0.125,  # bin size of lidar sensor (meter)
    'd_behind': 12,  # distance behind the ego vehicle (meter)
    'out_lane_thres': 5.0,  # threshold for out of lane
    'desired_speed': 8,  # desired speed (m/s)
    'max_ego_spawn_times': 200,  # maximum times to spawn ego vehicle
    'display_route': True,  # whether to render the desired route
    'pixor_size': 64,  # size of the pixor labels
    'pixor': False,  # whether to output PIXOR observation
  }

  buffer = Buffer(observation_dimensions, steps_per_epoch)

observation_input = keras.Input(shape=observation_dimensions, dtype=tf.float32)
logits = mlp(observation_input, list(hidden_sizes)+[num_actions], tf.tanh, None)
actor = keras.Model(inputs=observation_input, outputs=logits)
value = tf.squeeze(
    mlp(observation_input, list(hidden_sizes)+[1], tf.tanh, None), axis=1
)
critic = keras.Model(inputs=observation_input, outputs=value)

# Initialize the policy and the value function optimizers
policy_optimizer = keras.optimizers.Adam(learning_rate=policy_learning_rate)
value_optimizer = keras.optimizers.Adam(learning_rate=value_function_learning_rate)

actor.summary()
tf.config.run_functions_eagerly(True)

def transform(img):
    return img[76:204,76:204,:]/255

epoch = 1


env = gym.make('carla-v0', params=params)
observation =  env.reset() # select_env()
episode_return, episode_length = 0, 0
for _ in range(20): observation, _, _, _ = env.step([1,0])
observation = transform( observation['birdeye'] )
    
while epoch != epochs:
    # Initialize the sum of the returns, lengths and number of episodes for each epoch
    sum_return = 0
    sum_length = 0
    num_episodes = 0
        
    # Iterate over the steps of each epoch
    for t in range(steps_per_epoch): #while True
        # Get the logits, action, and take one step in the environment
        logits, action = sample_action(actor, observation.reshape(1, *observation_dimensions))
        observation_new, reward, done, _ = env.step([1, action[0].numpy()-1])
        observation_new = transform(observation_new['birdeye'])
        episode_return += reward
        episode_length += 1

        # Get the value and log-probability of the action
        value_t = critic(observation.reshape(1, *observation_dimensions))
        logprobability_t = logprobabilities(logits, action)

        # Store obs, act, rew, v_t, logp_pi_t
        buffer.store(observation, action, reward, value_t, logprobability_t)

        # Update the observation
        observation = observation_new

        # Finish trajectory if reached to a terminal state
        terminal = done
        if terminal or (t == steps_per_epoch - 1):
            last_value = 0 if done else critic(observation.reshape(1, *observation_dimensions))
            buffer.finish_trajectory(last_value)
            sum_return += episode_return
            sum_length += episode_length
            num_episodes += 1
            observation =  env.reset()#select_env()
            episode_return, episode_length = 0, 0
            for _ in range(20): observation, _, _, _ = env.step([1,0])
            observation = transform(observation['birdeye'])
            
            
    # Print mean return and length for each epoch
    with open("logs_no_vae.txt", "a") as log:
        print(
            f"{epoch},{sum_return / num_episodes},{sum_length / num_episodes},{num_episodes}",
            file = log
        )
    print(
        f" Epoch: {epoch}. Mean Return: {sum_return / num_episodes}. Mean Length: {sum_length / num_episodes}"
    )
    
    # Get values from the buffer
    (
        observation_buffer,
        action_buffer,
        advantage_buffer,
        return_buffer,
        logprobability_buffer,
    ) = buffer.get()

    # Update the policy and implement early stopping using KL divergence
    for _ in range(train_policy_iterations):
        kl = train_policy(
            observation_buffer, action_buffer, logprobability_buffer, advantage_buffer
        )
        if kl > 1.5 * target_kl:
            # Early Stopping
            break

    # Update the value function
    for _ in range(train_value_iterations):
        train_value_function(observation_buffer, return_buffer)

    
    tm = int(time.time())
    actor.save_weights(f"models/novae_actor_{tm}.h5")
    critic.save_weights(f"models/novae_critic_{tm}.h5")
    epoch += 1
# Discounted cumulative sums of vectors for computing rewards-to-go and advantage estimates # Buffer for storing trajectories # Buffer initialization # Append one step of agent-environment interaction # Finish the trajectory by computing advantage estimates and rewards-to-go # Get all data of the buffer and normalize the advantages # Build a feedforward neural network # Compute the log-probabilities of taking actions a by using the logits (i.e. the output of the actor) # Sample action from actor # Train the policy by maxizing the PPO-Clip objective # Record operations for automatic differentiation. # Train the value function by regression on mean-squared error # Record operations for automatic differentiation. # Hyperparameters of the PPO algorithm # True if you want to render the environment # screen size of bird-eye render # the number of past steps to draw # time interval between two frames # whether to use discrete control space # discrete value of accelerations # discrete value of steering angles # continuous acceleration range # continuous steering angle range # filter for defining ego vehicle # connection port # which town to simulate # mode of the task, [random, roundabout (only for Town03)] # maximum timesteps per episode # maximum number of waypoints # observation range (meter) # bin size of lidar sensor (meter) # distance behind the ego vehicle (meter) # threshold for out of lane # desired speed (m/s) # maximum times to spawn ego vehicle # whether to render the desired route # size of the pixor labels # whether to output PIXOR observation # Initialize the policy and the value function optimizers # select_env() # Initialize the sum of the returns, lengths and number of episodes for each epoch # Iterate over the steps of each epoch #while True # Get the logits, action, and take one step in the environment # Get the value and log-probability of the action # Store obs, act, rew, v_t, logp_pi_t # Update the observation # Finish trajectory if reached to a terminal state #select_env() # Print mean return and length for each epoch # Get values from the buffer # Update the policy and implement early stopping using KL divergence # Early Stopping # Update the value function