FrozenLake-random.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys

import gym


DEBUG = False
#DEBUG = True

ENVS = {
    "4x4" : "FrozenLake-v0",
    "8x8" : "FrozenLake8x8-v0",
}


def error(msg):
    sys.exit(msg)

def usage():
    error("Usage: FrozenLake-random <4x4|8x8> <test_count>")

def main():
    if len(sys.argv) != 3: usage()
    env_name = ENVS[sys.argv[1]]
    test_count = int(sys.argv[2])
    print("# <{}>".format(env_name))

    env = gym.make(env_name)

    reward_total = 0.0
    for episode in range(test_count):
        ob = env.reset()
        if DEBUG: env.render()
        while True:
            ob, reward, done, info = env.step(env.action_space.sample())
            if DEBUG:
                env.render()
                print(ob, reward, done, info)
            if done:
                reward_total += reward
                break

    print("episodes:       {}".format(test_count))
    print("total reward:   {}".format(reward_total))
    print("average reward: {:.3f}".format(reward_total / test_count))

if __name__ == "__main__": main()