gymz-controller

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

"""
gymz-controller.
Start threads for running an environment in an emulator and expose
input/output/reward buffers via ZeroMQ.

Usage:
    gymz-controller <emulator> <config>
    gymz-controller -h | --help
    gymz-controller --version

Options:
    -h --help         Show this screen.
    --version         Show version.
"""

import docopt
import json
import os
import signal
import threading
import time

import gymz


def run(args):
    if args['--version']:
        print('gymz {version}'.format(version=gymz.__version__))
        exit()

    # Load default configuration file
    config = gymz.misc.read_default_config()

    # Parse user config file and update config
    with open(args['<config>'], 'r') as f:
        gymz.misc.recursively_update_dict(config, json.load(f))

    # Set prefix (defaults to same directory as config)
    if config['All']['prefix'] is None:
        config['All']['prefix'] = os.path.split(args['<config>'])[0]

    # Create wrapper instance
    if args['<emulator>'] == 'gym':
        emu = gymz.GymWrapper(config)
    else:
        raise NotImplementedError('Unknown emulator.')

    # Load an environment
    emu.load_env(config['Env']['env'], monitor_args=config['Env']['monitor_args'])

    # Set random seed if given
    if config['All']['seed'] >= 0:
        emu.seed(config['All']['seed'])

    # Use mutable objects (lists) as containers for buffers to retain
    # correct pointers and avoid globals
    command_buffer = emu.get_command_buffer()
    output_buffer = emu.get_output_buffer()
    reward_buffer = emu.get_reward_buffer()
    done_buffer = emu.get_done_buffer()
    exit_event = threading.Event()

    # Create thread objects
    runner_thread = gymz.EnvRunnerThread(0, 'runner_thread', emu, command_buffer, output_buffer, reward_buffer, config, exit_event)
    zmq_reward_sender_thread = gymz.ZMQRewardSenderThread(1, 'zmq_reward_sender_thread', reward_buffer, done_buffer, config, exit_event)
    zmq_sender_thread = gymz.ZMQObservationSenderThread(2, 'zmq_sender_thread', output_buffer, done_buffer, config, exit_event)
    zmq_receiver_thread = gymz.ZMQCommandReceiverThread(3, 'zmq_receiver_thread', command_buffer, done_buffer, config, exit_event)

    # Set up signal handler for SIGINT to let all threads exit
    # gracefully if user wishes to exit
    signal.signal(signal.SIGINT, gymz.misc.SignalHandler(
        exit_event, [runner_thread, zmq_reward_sender_thread, zmq_sender_thread, zmq_receiver_thread]))

    # Fire up the threads, they will not terminate until SIGINT is received
    runner_thread.start()
    zmq_reward_sender_thread.start()
    zmq_sender_thread.start()
    zmq_receiver_thread.start()

    # Keep main thread alive to be able to receive SIGINT
    while True:
        time.sleep(1000)

if __name__ == '__main__':
    run(docopt.docopt(__doc__))