FrozenLake-handmade.py

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

# Just a handmade algorithm. NOT AI.

import sys

import gym

### 4x4 MAP ###
# SFFF
# FHFH
# FFFH
# HFFG
###############

### 8x8 MAP ###
# SFFFFFFF
# FFFFFFFF
# FFFHFFFF
# FFFFFHFF
# FFFHFFFF
# FHHFFFHF
# FHFFHFHF
# FFFHFFFG
###############

# 0:L, 1:D, 2:R, 3:U
ACTIONS_MAP = {
    "4x4" : (
         0,  3,  3,  3,
         0, -1,  0, -1,
         3,  1,  0, -1,
        -1,  2,  1, -1,
    ),
    "8x8" : (
         2,  2,  2,  2,  2,  2,  2,  2,
         3,  3,  3,  3,  3,  3,  3,  2,
        -1, -1, -1, -1, -1, -1, -1,  2,
        -1, -1, -1, -1, -1, -1, -1,  2,
        -1, -1, -1, -1, -1, -1, -1,  2,
        -1, -1, -1, -1, -1, -1, -1,  2,
        -1, -1, -1, -1, -1, -1, -1,  2,
        -1, -1, -1, -1, -1, -1, -1, -1,
    ),
}

EPISODE_COUNT = 10000

DEBUG = False
#DEBUG = True

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


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

def get_action(actions, ob):
    act = actions[ob]
    assert 0 <= act <= 3
    return act

def usage():
    error("Usage: FrozenLake-handmade <4x4|8x8> [recdir]")

def main():
    if len(sys.argv) < 2: usage()
    env_name = ENVS[sys.argv[1]]
    actions = ACTIONS_MAP[sys.argv[1]]
    rec_dir = sys.argv[2] if len(sys.argv) >= 3 else None
    print("# <{}>".format(env_name))

    env = gym.make(env_name)
    if rec_dir: env.monitor.start(rec_dir)

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

    if rec_dir: env.monitor.close()

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

if __name__ == "__main__": main()