gym_eval.py

import os

os.environ["OMP_NUM_THREADS"] = "1"
import argparse
from types import SimpleNamespace
import torch
from environment import atari_env
from utils import read_config, setup_logger
from model import A3Clstm
from player_util import Agent
import gym
import logging
import time
from utils import weights_init

config = dict(
    seed=int(1),
    max_episode_length=int(1e4),
    env='Seaquest-v0',
    env_config='config.json',
    load_model_dir='trained_models',
    log_dir='logs',
    gpu_id=int(0),
    skip_rate=int(4),
    exp_name='exp_wo_reparam',
    num_episodes=int(10),
    render=False,
    render_freq=int(1),
    new_gym_eval=False,
)

args = SimpleNamespace(**config)
setup_json = read_config(args.env_config)
env_conf = setup_json["Default"]
for i in setup_json.keys():
    if i in args.env:
        env_conf = setup_json[i]

gpu_id = args.gpu_id

device = torch.device('cuda:{}'.format(gpu_id) if gpu_id >= 0 else 'cpu')

torch.manual_seed(args.seed)
if gpu_id >= 0:
    torch.cuda.manual_seed(args.seed)

log = {}
setup_logger(
    '{}_mon_log'.format(args.env),
    os.path.join(args.log_dir, '{}-{}_mon_log'.format(args.env,
                                                      args.exp_name)))
log['{}_mon_log'.format(args.env)] = logging.getLogger('{}_mon_log'.format(
    args.env))

d_args = vars(args)
for k in d_args.keys():
    log['{}_mon_log'.format(args.env)].info('{0}: {1}'.format(k, d_args[k]))

env = atari_env("{}".format(args.env), env_conf, args)
num_tests = 0
start_time = time.time()
reward_total_sum = 0
player = Agent(None, env, args, None, gpu_id=gpu_id)
player.model = A3Clstm(player.env.observation_space.shape[0],
                       player.env.action_space)
player.gpu_id = gpu_id

player.model = player.model.to(device)
if args.new_gym_eval:
    player.env = gym.wrappers.Monitor(player.env,
                                      "{}-{}_monitor".format(
                                          args.env, args.exp_name),
                                      force=True)

player.model.load_state_dict(
    torch.load(os.path.join(args.load_model_dir,
                            '{}-{}.dat'.format(args.env, args.exp_name)),
               map_location=device))
player.model.eval()

for i_episode in range(args.num_episodes):
    player.state = player.env.reset()
    player.state = torch.from_numpy(player.state).to(torch.float32)
    player.state = player.state.to(device)

    player.eps_len += 2
    reward_sum = 0
    while True:
        if args.render:
            if i_episode % args.render_freq == 0:
                player.env.render()

        player.action_test()
        reward_sum += player.reward

        if player.done and not player.info:
            state = player.env.reset()
            player.eps_len += 2
            player.state = torch.from_numpy(state).to(torch.float32)
            player.state = player.state.to(device)
        elif player.info:
            num_tests += 1
            reward_total_sum += reward_sum
            reward_mean = reward_total_sum / num_tests
            log['{}_mon_log'.format(args.env)].info(
                "Time {0}, episode reward {1}, episode length {2}, reward mean {3:.4f}"
                .format(
                    time.strftime("%Hh %Mm %Ss",
                                  time.gmtime(time.time() - start_time)),
                    reward_sum, player.eps_len, reward_mean))
            player.eps_len = 0
            break