gym_play.py

import os
import select
import sys
import termios

import gym


def play():
    env = gym.make("AdventureNoFrameskip-v4" if len(sys.argv) < 2 else sys.argv[1])
    num_actions = env.action_space.n

    unused_observation = env.reset()

    while True:
        env.render()
        rlist, _, _ = select.select((0,), (), ())
        if rlist:
            c = os.read(0, 1)
            if c == b"q":
                break
            elif c == b" ":
                select.select((0,), (), ())
            elif c[0] in (ord("\n"), ord("\r")):  # Return
                env.reset()
                continue
            else:
                action = int(c[0] - ord("0"))
                if action <= 0 or action >= num_actions:
                    continue

            _, reward, done, info = env.step(action)
            print("reward:", reward, "done:", done, "info:", info)


def main():
    old = termios.tcgetattr(0)
    new = termios.tcgetattr(0)

    # Set to unbuffered, no echo.
    new[3] &= ~(termios.ICANON | termios.ECHO | termios.ECHONL)  # lflags
    termios.tcsetattr(0, termios.TCSANOW, new)

    try:
        play()
    except KeyboardInterrupt:
        pass
    finally:
        termios.tcsetattr(0, termios.TCSANOW, old)


if __name__ == "__main__":
    main()