Trust_Fed/AV_Data_Tensor at main · gr7-cipher/Trust_Fed · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical

import numpy as np
import gym
import random
import time

from model import QNetwork
from utils import TensorEnv, ExperienceReplayMemory, AgentConfig
import utils

args = utils.arguments()


env = TensorEnv(args.env)

agent = QNetwork(env.observation_space.shape[0], env.action_space.n)
agent.load_state_dict(torch.load('./Data_AVs/{}/model_state_dict'.format(args.env)))

for ep in tqdm(range(10)):

    s = env.reset()
    done = False
    while not done:

        action = agent(s).reshape(-1).argmax().item()

        s, r, done, _ = env.step(action)

        env.render()
        if not args.env.startswith('vehicle'):
           time.sleep(0.02)

env.close()