<reponame>NNpanpan/CS4278-5478-Project-Materials
import ast
import argparse
import logging
import os
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as mp

# Duckietown Specific
import gym
from learning.reinforcement.a3c import a3c_model as a3c
from learning.utils.env import launch_env
from learning.utils.wrappers import NormalizeWrapper, ImgWrapper, \
    DtRewardWrapper, ActionWrapper, ResizeWrapper, DiscreteWrapper_a6

def preprocess_state(obs):
    from PIL import Image
    x = obs.transpose(1, 2, 0).mean(2)
    x = np.uint8(x * 255)
    x = Image.fromarray(x).resize((80, 80))
    x = np.array(x, dtype=np.float32).reshape(1, 80, 80) / 255

    return x

def _enjoy(args):
    env = launch_env()
    # env = ResizeWrapper(env)
    # env = NormalizeWrapper(env)
    env = ImgWrapper(env)  # to make the images from 160x120x3 into 3x160x120
    # env = ActionWrapper(env)
    env = DtRewardWrapper(env)
    env = DiscreteWrapper_a6(env)

    shape_obs_space = env.observation_space.shape  # (3, 120, 160)
    shape_action_space = env.action_space.n  # (2,)

    # Initialize policy

    # Load model
    cwd = os.getcwd()
    path = os.path.join(cwd, args.model_dir, args.model_file)
    print('Loading model from:', path)

    checkpoint = torch.load(path)
    global_net = a3c.Net(channels=1, num_actions=shape_action_space)
    global_net.load_state_dict(checkpoint['model_state_dict'])
    global_net.eval()

    state = torch.tensor(preprocess_state(env.reset()))
    done = True

    while True:
        with torch.no_grad():
            if done:
                hx = torch.zeros(1, 256)
            else:
                hx = hx.detach()

            # Inference
            value, logit, hx = global_net.forward((state.view(-1, 1, 80, 80), hx))
            action_log_probs = F.log_softmax(logit, dim=-1)

            # Take action with highest probability
            action = action_log_probs.max(1, keepdim=True)[1].numpy().squeeze()

            # Perform action
            state, reward, done, _ = env.step(action)
            state = torch.tensor(preprocess_state(state))

            env.render()

            if done:
                state = torch.tensor(preprocess_state(env.reset()))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model-dir', type=str, default='reinforcement/a3c/models')  # Name of the directory where the models are saved
    parser.add_argument('--model-file', type=str,
                        default='a3c-duckie.pth')  # Name of the model file
    _enjoy(parser.parse_args())
