import multiprocessing as mp
import os
import argparse
import random
import sys
import tempfile

import numpy as np

import chainer
from chainer import optimizers

import fc_tail_q_function
import dqn_head
import q_function
from dqn import DQN
import ale
import random_seed
import async
import rmsprop_ones
import replay_buffer


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('rom', type=str)
    parser.add_argument('--seed', type=int, default=None)
    parser.add_argument('--gpu', type=int, default=0)
    parser.add_argument('--use-sdl', action='store_true')
    parser.add_argument('--final-exploration-frames', type=int, default=1e6)
    parser.set_defaults(use_sdl=False)
    args = parser.parse_args()

    if args.seed is not None:
        random_seed.set_random_seed(args.seed)

    n_actions = ale.ALE(args.rom).number_of_actions

    head = dqn_head.NIPSDQNHead()
    q_func = fc_tail_q_function.FCTailQFunction(
        head, 256, n_actions=n_actions)
    opt = optimizers.RMSpropGraves(
        lr=2.5e-4, alpha=0.95, momentum=0.95, eps=1e-4)
    opt.setup(q_func)
    opt.add_hook(chainer.optimizer.GradientClipping(1.0))
    rbuf = replay_buffer.ReplayBuffer(1e6)
    # TODO: epsilon scheduling
    agent = DQN(q_func, opt, rbuf, gpu=args.gpu, gamma=0.99,
                epsilon=1.0, replay_start_size=1000, target_update_frequency=1000)

    env = ale.ALE(args.rom, use_sdl=args.use_sdl)

    episode_r = 0

    for i in xrange(1000000):
        try:
            episode_r += env.reward

            action = agent.act(env.state, env.reward, env.is_terminal)

            if agent.epsilon >= 0.1:
                agent.epsilon -= 0.9 / args.final_exploration_frames

            if env.is_terminal:
                print 'i:{} epsilon:{} episode_r:{}'.format(i, agent.epsilon, episode_r)
                episode_r = 0
                env.initialize()
            else:
                env.receive_action(action)
        except KeyboardInterrupt:
            tempdir = tempfile.mkdtemp(prefix='drill')
            agent.save_model(tempdir + '/{}_keyboardinterrupt.h5'.format(i))
            print >> sys.stderr, 'Saved the current model to {}'.format(
                tempdir)
            raise

if __name__ == '__main__':
    main()

@@ -642,30 +642,421 @@
 es',
- type=int, default=1e6
+%0A                        type=int, default=int(1e6))%0A    parser.add_argument('--model', type=str, default='')%0A    parser.add_argument('--arch', type=str, default='nature')%0A    parser.add_argument('--steps', type=int, default=int(1e7))%0A    parser.add_argument('--replay-start-size', type=int, default=int(5e4))%0A    parser.add_argument('--target-update-frequency',%0A                        type=int, default=int(1e4)
 )%0A  
@@ -1255,16 +1255,217 @@
 ctions%0A%0A
+    if args.arch == 'nature':%0A        head = dqn_head.NatureDQNHead()%0A        q_func = fc_tail_q_function.FCTailQFunction(%0A            head, 512, n_actions=n_actions)%0A    elif args.arch == 'nips':%0A    
     head
@@ -1490,24 +1490,28 @@
 NHead()%0A    
+  
+  
 q_func = fc_
@@ -1547,24 +1547,28 @@
 on(%0A        
+    
 head, 256, n
@@ -1583,24 +1583,114 @@
 =n_actions)%0A
+    else:%0A        raise RuntimeError('Not supported architecture: %7B%7D'.format(args.arch))%0A%0A
     opt = op
@@ -1790,24 +1790,26 @@
 (q_func)%0A   
+ #
  opt.add_hoo
@@ -1895,42 +1895,11 @@
 r(1e
-6
+5
 )%0A
-    # TODO: epsilon scheduling%0A
     
@@ -2006,42 +2006,166 @@
 ize=
-1000, target_update_frequency=1000
+args.replay_start_size,%0A                target_update_frequency=args.target_update_frequency)%0A%0A    if len(args.model) %3E 0:%0A        agent.load_model(args.model
 )%0A%0A 
@@ -2257,15 +2257,18 @@
 nge(
-1000000
+args.steps
 ):%0A 
@@ -2774,16 +2774,73 @@
 errupt:%0A
+            # Save the current model before being killed%0A
         
@@ -3080,16 +3080,206 @@
  raise%0A%0A
+    # Save the final model%0A    tempdir = tempfile.mkdtemp(prefix='drill')%0A    agent.save_model(tempdir + '/%7B%7D_finish.h5'.format())%0A    print 'Saved the current model to %7B%7D'.format(tempdir)%0A%0A
 if __nam

Add options