test.py 7 KB
Newer Older
pfrl_rainbow's avatar
pfrl_rainbow committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20

import json
import select
import time
import logging
import os
import threading


from typing import Callable

import aicrowd_helper
import gym
import minerl
import abc
import numpy as np

import coloredlogs
coloredlogs.install(logging.DEBUG)

21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
# our dependencies
import joblib

import sys
sys.path.append(os.path.abspath(os.path.join(__file__, os.pardir, 'mod')))
from dqn_family import get_agent
from env_wrappers import wrap_env

GPU = -1

ARCH = 'distributed_dueling'
NOISY_NET_SIGMA = 0.5
FINAL_EPSILON = 0.01
FINAL_EXPLORATION_FRAMES = 10 ** 6
LR = 0.0000625
ADAM_EPS = 0.00015
PRIORITIZED = True
UPDATE_INTERVAL = 4
REPLAY_CAPACITY = 300000
NUM_STEP_RETURN = 10
AGENT_TYPE = 'CategoricalDoubleDQN'
GAMMA = 0.99
REPLAY_START_SIZE = 5000
TARGET_UPDATE_INTERVAL = 10000
CLIP_DELTA = True
BATCH_ACCUMULATOR = 'mean'
FRAME_SKIP = 4
GRAY_SCALE = False
FRAME_STACK = 4
RANDOMIZE_ACTION = NOISY_NET_SIGMA is None
EVAL_EPSILON = 0.001

maximum_frames = 8000000
STEPS = maximum_frames // FRAME_SKIP


pfrl_rainbow's avatar
pfrl_rainbow committed
57 58 59 60 61 62
# All the evaluations will be evaluated on MineRLObtainDiamondVectorObf-v0 environment
MINERL_GYM_ENV = os.getenv('MINERL_GYM_ENV', 'MineRLObtainDiamondVectorObf-v0')
MINERL_MAX_EVALUATION_EPISODES = int(os.getenv('MINERL_MAX_EVALUATION_EPISODES', 5))

# Parallel testing/inference, **you can override** below value based on compute
# requirements, etc to save OOM in this phase.
pfrl_rainbow's avatar
pfrl_rainbow committed
63
EVALUATION_THREAD_COUNT = int(os.getenv('EPISODES_EVALUATION_THREAD_COUNT', 2))
pfrl_rainbow's avatar
pfrl_rainbow committed
64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95

class EpisodeDone(Exception):
    pass

class Episode(gym.Env):
    """A class for a single episode.
    """
    def __init__(self, env):
        self.env = env
        self.action_space = env.action_space
        self.observation_space = env.observation_space
        self._done = False

    def reset(self):
        if not self._done:
            return self.env.reset()

    def step(self, action):
        s,r,d,i = self.env.step(action)
        if d:
            self._done = True
            raise EpisodeDone()
        else:
            return s,r,d,i



# DO NOT CHANGE THIS CLASS, THIS IS THE BASE CLASS FOR YOUR AGENT.
class MineRLAgentBase(abc.ABC):
    """
    To compete in the competition, you are required to implement a
    SUBCLASS to this class.
pfrl_rainbow's avatar
pfrl_rainbow committed
96

pfrl_rainbow's avatar
pfrl_rainbow committed
97 98
    YOUR SUBMISSION WILL FAIL IF:
        * Rename this class
pfrl_rainbow's avatar
pfrl_rainbow committed
99
        * You do not implement a subclass to this class
pfrl_rainbow's avatar
pfrl_rainbow committed
100

pfrl_rainbow's avatar
pfrl_rainbow committed
101
    This class enables the evaluator to run your agent in parallel,
pfrl_rainbow's avatar
pfrl_rainbow committed
102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121
    so you should load your model only once in the 'load_agent' method.
    """

    @abc.abstractmethod
    def load_agent(self):
        """
        This method is called at the beginning of the evaluation.
        You should load your model and do any preprocessing here.
        THIS METHOD IS ONLY CALLED ONCE AT THE BEGINNING OF THE EVALUATION.
        DO NOT LOAD YOUR MODEL ANYWHERE ELSE.
        """
        raise NotImplementedError()

    @abc.abstractmethod
    def run_agent_on_episode(self, single_episode_env : Episode):
        """This method runs your agent on a SINGLE episode.

        You should just implement the standard environment interaction loop here:
            obs  = env.reset()
            while not done:
pfrl_rainbow's avatar
pfrl_rainbow committed
122
                env.step(self.agent.act(obs))
pfrl_rainbow's avatar
pfrl_rainbow committed
123
                ...
pfrl_rainbow's avatar
pfrl_rainbow committed
124

pfrl_rainbow's avatar
pfrl_rainbow committed
125 126 127 128 129 130 131 132 133 134 135 136 137 138
        NOTE: This method will be called in PARALLEL during evaluation.
            So, only store state in LOCAL variables.
            For example, if using an LSTM, don't store the hidden state in the class
            but as a local variable to the method.

        Args:
            env (gym.Env): The env your agent should interact with.
        """
        raise NotImplementedError()


#######################
# YOUR CODE GOES HERE #
#######################
139 140 141
class MineRLRainbowBaselineAgent(MineRLAgentBase):
    def __init__(self, env):
        self.env = env
pfrl_rainbow's avatar
pfrl_rainbow committed
142 143

    def load_agent(self):
144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163
        self.agent = get_agent(
            n_actions=self.env.action_space.n, arch=ARCH, n_input_channels=self.env.observation_space.shape[0],
            noisy_net_sigma=NOISY_NET_SIGMA, final_epsilon=FINAL_EPSILON,
            final_exploration_frames=FINAL_EXPLORATION_FRAMES, explorer_sample_func=self.env.action_space.sample,
            lr=LR, adam_eps=ADAM_EPS,
            prioritized=PRIORITIZED, steps=STEPS, update_interval=UPDATE_INTERVAL,
            replay_capacity=REPLAY_CAPACITY, num_step_return=NUM_STEP_RETURN,
            agent_type=AGENT_TYPE, gpu=GPU, gamma=GAMMA, replay_start_size=REPLAY_START_SIZE,
            target_update_interval=TARGET_UPDATE_INTERVAL, clip_delta=CLIP_DELTA,
            batch_accumulator=BATCH_ACCUMULATOR,
        )

        self.agent.load(os.path.abspath(os.path.join(__file__, os.pardir, 'train')))

    def run_agent_on_episode(self, single_episode_env: Episode):
        with self.agent.eval_mode():
            obs = single_episode_env.reset()
            while True:
                a = self.agent.act(obs)
                obs, r, done, info = single_episode_env.step(a)
pfrl_rainbow's avatar
pfrl_rainbow committed
164

pfrl_rainbow's avatar
pfrl_rainbow committed
165

pfrl_rainbow's avatar
pfrl_rainbow committed
166
#####################################################################
pfrl_rainbow's avatar
pfrl_rainbow committed
167
# IMPORTANT: SET THIS VARIABLE WITH THE AGENT CLASS YOU ARE USING   #
pfrl_rainbow's avatar
pfrl_rainbow committed
168
######################################################################
169
AGENT_TO_TEST = MineRLRainbowBaselineAgent # MineRLMatrixAgent, MineRLRandomAgent, YourAgentHere
pfrl_rainbow's avatar
pfrl_rainbow committed
170 171 172 173 174 175 176



####################
# EVALUATION CODE  #
####################
def main():
177
    assert MINERL_MAX_EVALUATION_EPISODES > 0
pfrl_rainbow's avatar
pfrl_rainbow committed
178
    assert EVALUATION_THREAD_COUNT > 0
179

pfrl_rainbow's avatar
pfrl_rainbow committed
180
    # Create the parallel envs (sequentially to prevent issues!)
181 182
    kmeans = joblib.load(os.path.abspath(os.path.join(__file__, os.pardir, 'train', 'kmeans.joblib')))

pfrl_rainbow's avatar
pfrl_rainbow committed
183 184 185 186 187 188 189
    def wrapper(env):
        return wrap_env(
            env=env, test=True, monitor=False, outdir=None,
            frame_skip=FRAME_SKIP, gray_scale=GRAY_SCALE, frame_stack=FRAME_STACK,
            randomize_action=RANDOMIZE_ACTION, eval_epsilon=EVAL_EPSILON,
            action_choices=kmeans.cluster_centers_,
        )
190

pfrl_rainbow's avatar
pfrl_rainbow committed
191 192 193 194
    envs = [wrapper(gym.make(MINERL_GYM_ENV)) for _ in range(EVALUATION_THREAD_COUNT)]
    # envs = [gym.make(MINERL_GYM_ENV) for _ in range(EVALUATION_THREAD_COUNT)]
    agent = AGENT_TO_TEST(envs[0])
    # agent = AGENT_TO_TEST()
pfrl_rainbow's avatar
pfrl_rainbow committed
195 196 197
    assert isinstance(agent, MineRLAgentBase)
    agent.load_agent()

pfrl_rainbow's avatar
pfrl_rainbow committed
198 199 200 201
    episodes_per_thread = [MINERL_MAX_EVALUATION_EPISODES // EVALUATION_THREAD_COUNT for _ in range(EVALUATION_THREAD_COUNT)]
    episodes_per_thread[-1] += MINERL_MAX_EVALUATION_EPISODES - EVALUATION_THREAD_COUNT *(MINERL_MAX_EVALUATION_EPISODES // EVALUATION_THREAD_COUNT)
    # A simple funciton to evaluate on episodes!
    def evaluate(i, env):
pfrl_rainbow's avatar
pfrl_rainbow committed
202
        print("[{}] Starting evaluator.".format(i))
pfrl_rainbow's avatar
pfrl_rainbow committed
203 204 205 206 207 208 209 210 211 212 213 214 215 216
        for i in range(episodes_per_thread[i]):
            try:
                agent.run_agent_on_episode(Episode(env))
            except EpisodeDone:
                print("[{}] Episode complete".format(i))
                pass

    evaluator_threads = [threading.Thread(target=evaluate, args=(i, envs[i])) for i in range(EVALUATION_THREAD_COUNT)]
    for thread in evaluator_threads:
        thread.start()

    # wait fo the evaluation to finish
    for thread in evaluator_threads:
        thread.join()
217

pfrl_rainbow's avatar
pfrl_rainbow committed
218 219 220

if __name__ == "__main__":
    main()