Update fluid/DeepQNetwork models with Atari environment (#981)

* DQN Atari based on RLLab

* DQN Atari based on RLLab

* DQN Atari based on RLLab

* refactor code without RLLab framework

* add module of saving and loading policy model

* refactor code structure, add DoubleDQN and DuelingDQN modules

* add fluid argmax, flatten utils

* update README.md

* udpdate replay memory

* udpdate replay memory

* update readme, code clean

* clean code and fix codestyle

* fix codestyle

* update README.md

* revisions|history->context, randint->random

* revisions| add comment for max-pooling operation in atari

fluid/DeepQNetwork/DQN.py

-#-*- coding: utf-8 -*-
-#File: DQN.py
-
-from agent import Model
-import gym
-import argparse
-from tqdm import tqdm
-from expreplay import ReplayMemory, Experience
-import numpy as np
-import os
-
-UPDATE_FREQ = 4
-
-MEMORY_WARMUP_SIZE = 1000
-
-
-def run_episode(agent, env, exp, train_or_test):
-    assert train_or_test in ['train', 'test'], train_or_test
-    total_reward = 0
-    state = env.reset()
-    for step in range(200):
-        action = agent.act(state, train_or_test)
-        next_state, reward, isOver, _ = env.step(action)
-        if train_or_test == 'train':
-            exp.append(Experience(state, action, reward, isOver))
-            # train model
-            # start training 
-            if len(exp) > MEMORY_WARMUP_SIZE:
-                batch_idx = np.random.randint(
-                    len(exp) - 1, size=(args.batch_size))
-                if step % UPDATE_FREQ == 0:
-                    batch_state, batch_action, batch_reward, \
-                    batch_next_state, batch_isOver = exp.sample(batch_idx)
-                    agent.train(batch_state, batch_action, batch_reward, \
-                                batch_next_state, batch_isOver)
-        total_reward += reward
-        state = next_state
-        if isOver:
-            break
-    return total_reward
-
-
-def train_agent():
-    env = gym.make(args.env)
-    state_shape = env.observation_space.shape
-    exp = ReplayMemory(args.mem_size, state_shape)
-    action_dim = env.action_space.n
-    agent = Model(state_shape[0], action_dim, gamma=0.99)
-
-    while len(exp) < MEMORY_WARMUP_SIZE:
-        run_episode(agent, env, exp, train_or_test='train')
-
-    max_episode = 4000
-
-    # train
-    total_episode = 0
-    pbar = tqdm(total=max_episode)
-    recent_100_reward = []
-    for episode in xrange(max_episode):
-        # start epoch
-        total_reward = run_episode(agent, env, exp, train_or_test='train')
-        pbar.set_description('[train]exploration:{}'.format(agent.exploration))
-        pbar.update()
-
-        # recent 100 reward
-        total_reward = run_episode(agent, env, exp, train_or_test='test')
-        recent_100_reward.append(total_reward)
-        if len(recent_100_reward) > 100:
-            recent_100_reward = recent_100_reward[1:]
-        pbar.write("episode:{}    test_reward:{}".format(\
-                    episode, np.mean(recent_100_reward)))
-
-    pbar.close()
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--env', type=str, default='MountainCar-v0', \
-                        help='enviroment to train DQN model, e.g CartPole-v0')
-    parser.add_argument('--gamma', type=float, default=0.99, \
-                        help='discount factor for accumulated reward computation')
-    parser.add_argument('--mem_size', type=int, default=500000, \
-                        help='memory size for experience replay')
-    parser.add_argument('--batch_size', type=int, default=192, \
-                        help='batch size for training')
-    args = parser.parse_args()
-
-    train_agent()

fluid/DeepQNetwork/DQN_agent.py

+#-*- coding: utf-8 -*-
+
+import paddle.fluid as fluid
+from paddle.fluid.param_attr import ParamAttr
+import numpy as np
+import math
+from tqdm import tqdm
+from utils import fluid_flatten
+
+
+class DQNModel(object):
+    def __init__(self, state_dim, action_dim, gamma, hist_len, use_cuda=False):
+        self.img_height = state_dim[0]
+        self.img_width = state_dim[1]
+        self.action_dim = action_dim
+        self.gamma = gamma
+        self.exploration = 1.1
+        self.update_target_steps = 10000 // 4
+        self.hist_len = hist_len
+        self.use_cuda = use_cuda
+
+        self.global_step = 0
+        self._build_net()
+
+    def _get_inputs(self):
+        return fluid.layers.data(
+                   name='state',
+                   shape=[self.hist_len, self.img_height, self.img_width],
+                   dtype='float32'), \
+               fluid.layers.data(
+                   name='action', shape=[1], dtype='int32'), \
+               fluid.layers.data(
+                   name='reward', shape=[], dtype='float32'), \
+               fluid.layers.data(
+                   name='next_s',
+                   shape=[self.hist_len, self.img_height, self.img_width],
+                   dtype='float32'), \
+               fluid.layers.data(
+                   name='isOver', shape=[], dtype='bool')
+
+    def _build_net(self):
+        state, action, reward, next_s, isOver = self._get_inputs()
+        self.pred_value = self.get_DQN_prediction(state)
+        self.predict_program = fluid.default_main_program().clone()
+
+        reward = fluid.layers.clip(reward, min=-1.0, max=1.0)
+
+        action_onehot = fluid.layers.one_hot(action, self.action_dim)
+        action_onehot = fluid.layers.cast(action_onehot, dtype='float32')
+
+        pred_action_value = fluid.layers.reduce_sum(
+            fluid.layers.elementwise_mul(action_onehot, self.pred_value), dim=1)
+
+        targetQ_predict_value = self.get_DQN_prediction(next_s, target=True)
+        best_v = fluid.layers.reduce_max(targetQ_predict_value, dim=1)
+        best_v.stop_gradient = True
+
+        target = reward + (1.0 - fluid.layers.cast(
+            isOver, dtype='float32')) * self.gamma * best_v
+        cost = fluid.layers.square_error_cost(pred_action_value, target)
+        cost = fluid.layers.reduce_mean(cost)
+
+        self._sync_program = self._build_sync_target_network()
+
+        optimizer = fluid.optimizer.Adam(1e-3 * 0.5, epsilon=1e-3)
+        optimizer.minimize(cost)
+
+        # define program
+        self.train_program = fluid.default_main_program()
+
+        # fluid exe
+        place = fluid.CUDAPlace(0) if self.use_cuda else fluid.CPUPlace()
+        self.exe = fluid.Executor(place)
+        self.exe.run(fluid.default_startup_program())
+
+    def get_DQN_prediction(self, image, target=False):
+        image = image / 255.0
+
+        variable_field = 'target' if target else 'policy'
+
+        conv1 = fluid.layers.conv2d(
+            input=image,
+            num_filters=32,
+            filter_size=[5, 5],
+            stride=[1, 1],
+            padding=[2, 2],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv1'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv1_b'.format(variable_field)))
+        max_pool1 = fluid.layers.pool2d(
+            input=conv1, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv2 = fluid.layers.conv2d(
+            input=max_pool1,
+            num_filters=32,
+            filter_size=[5, 5],
+            stride=[1, 1],
+            padding=[2, 2],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv2'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv2_b'.format(variable_field)))
+        max_pool2 = fluid.layers.pool2d(
+            input=conv2, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv3 = fluid.layers.conv2d(
+            input=max_pool2,
+            num_filters=64,
+            filter_size=[4, 4],
+            stride=[1, 1],
+            padding=[1, 1],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv3'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv3_b'.format(variable_field)))
+        max_pool3 = fluid.layers.pool2d(
+            input=conv3, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv4 = fluid.layers.conv2d(
+            input=max_pool3,
+            num_filters=64,
+            filter_size=[3, 3],
+            stride=[1, 1],
+            padding=[1, 1],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv4'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv4_b'.format(variable_field)))
+
+        flatten = fluid_flatten(conv4)
+
+        out = fluid.layers.fc(
+            input=flatten,
+            size=self.action_dim,
+            param_attr=ParamAttr(name='{}_fc1'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_fc1_b'.format(variable_field)))
+        return out
+
+    def _build_sync_target_network(self):
+        vars = list(fluid.default_main_program().list_vars())
+        policy_vars = filter(
+            lambda x: 'GRAD' not in x.name and 'policy' in x.name, vars)
+        target_vars = filter(
+            lambda x: 'GRAD' not in x.name and 'target' in x.name, vars)
+        policy_vars.sort(key=lambda x: x.name)
+        target_vars.sort(key=lambda x: x.name)
+
+        sync_program = fluid.default_main_program().clone()
+        with fluid.program_guard(sync_program):
+            sync_ops = []
+            for i, var in enumerate(policy_vars):
+                sync_op = fluid.layers.assign(policy_vars[i], target_vars[i])
+                sync_ops.append(sync_op)
+        sync_program = sync_program.prune(sync_ops)
+        return sync_program
+
+    def act(self, state, train_or_test):
+        sample = np.random.random()
+        if train_or_test == 'train' and sample < self.exploration:
+            act = np.random.randint(self.action_dim)
+        else:
+            if np.random.random() < 0.01:
+                act = np.random.randint(self.action_dim)
+            else:
+                state = np.expand_dims(state, axis=0)
+                pred_Q = self.exe.run(self.predict_program,
+                                      feed={'state': state.astype('float32')},
+                                      fetch_list=[self.pred_value])[0]
+                pred_Q = np.squeeze(pred_Q, axis=0)
+                act = np.argmax(pred_Q)
+        if train_or_test == 'train':
+            self.exploration = max(0.1, self.exploration - 1e-6)
+        return act
+
+    def train(self, state, action, reward, next_state, isOver):
+        if self.global_step % self.update_target_steps == 0:
+            self.sync_target_network()
+        self.global_step += 1
+
+        action = np.expand_dims(action, -1)
+        self.exe.run(self.train_program,
+                     feed={
+                         'state': state.astype('float32'),
+                         'action': action.astype('int32'),
+                         'reward': reward,
+                         'next_s': next_state.astype('float32'),
+                         'isOver': isOver
+                     })
+
+    def sync_target_network(self):
+        self.exe.run(self._sync_program)

fluid/DeepQNetwork/DoubleDQN_agent.py

+#-*- coding: utf-8 -*-
+
+import paddle.fluid as fluid
+from paddle.fluid.param_attr import ParamAttr
+import numpy as np
+from tqdm import tqdm
+import math
+from utils import fluid_argmax, fluid_flatten
+
+
+class DoubleDQNModel(object):
+    def __init__(self, state_dim, action_dim, gamma, hist_len, use_cuda=False):
+        self.img_height = state_dim[0]
+        self.img_width = state_dim[1]
+        self.action_dim = action_dim
+        self.gamma = gamma
+        self.exploration = 1.1
+        self.update_target_steps = 10000 // 4
+        self.hist_len = hist_len
+        self.use_cuda = use_cuda
+
+        self.global_step = 0
+        self._build_net()
+
+    def _get_inputs(self):
+        return fluid.layers.data(
+                   name='state',
+                   shape=[self.hist_len, self.img_height, self.img_width],
+                   dtype='float32'), \
+               fluid.layers.data(
+                   name='action', shape=[1], dtype='int32'), \
+               fluid.layers.data(
+                   name='reward', shape=[], dtype='float32'), \
+               fluid.layers.data(
+                   name='next_s',
+                   shape=[self.hist_len, self.img_height, self.img_width],
+                   dtype='float32'), \
+               fluid.layers.data(
+                   name='isOver', shape=[], dtype='bool')
+
+    def _build_net(self):
+        state, action, reward, next_s, isOver = self._get_inputs()
+        self.pred_value = self.get_DQN_prediction(state)
+        self.predict_program = fluid.default_main_program().clone()
+
+        reward = fluid.layers.clip(reward, min=-1.0, max=1.0)
+
+        action_onehot = fluid.layers.one_hot(action, self.action_dim)
+        action_onehot = fluid.layers.cast(action_onehot, dtype='float32')
+
+        pred_action_value = fluid.layers.reduce_sum(
+            fluid.layers.elementwise_mul(action_onehot, self.pred_value), dim=1)
+
+        targetQ_predict_value = self.get_DQN_prediction(next_s, target=True)
+
+        next_s_predcit_value = self.get_DQN_prediction(next_s)
+        greedy_action = fluid_argmax(next_s_predcit_value)
+
+        predict_onehot = fluid.layers.one_hot(greedy_action, self.action_dim)
+        best_v = fluid.layers.reduce_sum(
+            fluid.layers.elementwise_mul(predict_onehot, targetQ_predict_value),
+            dim=1)
+        best_v.stop_gradient = True
+
+        target = reward + (1.0 - fluid.layers.cast(
+            isOver, dtype='float32')) * self.gamma * best_v
+        cost = fluid.layers.square_error_cost(pred_action_value, target)
+        cost = fluid.layers.reduce_mean(cost)
+
+        self._sync_program = self._build_sync_target_network()
+
+        optimizer = fluid.optimizer.Adam(1e-3 * 0.5, epsilon=1e-3)
+        optimizer.minimize(cost)
+
+        # define program
+        self.train_program = fluid.default_main_program()
+
+        # fluid exe
+        place = fluid.CUDAPlace(0) if self.use_cuda else fluid.CPUPlace()
+        self.exe = fluid.Executor(place)
+        self.exe.run(fluid.default_startup_program())
+
+    def get_DQN_prediction(self, image, target=False):
+        image = image / 255.0
+
+        variable_field = 'target' if target else 'policy'
+
+        conv1 = fluid.layers.conv2d(
+            input=image,
+            num_filters=32,
+            filter_size=[5, 5],
+            stride=[1, 1],
+            padding=[2, 2],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv1'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv1_b'.format(variable_field)))
+        max_pool1 = fluid.layers.pool2d(
+            input=conv1, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv2 = fluid.layers.conv2d(
+            input=max_pool1,
+            num_filters=32,
+            filter_size=[5, 5],
+            stride=[1, 1],
+            padding=[2, 2],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv2'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv2_b'.format(variable_field)))
+        max_pool2 = fluid.layers.pool2d(
+            input=conv2, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv3 = fluid.layers.conv2d(
+            input=max_pool2,
+            num_filters=64,
+            filter_size=[4, 4],
+            stride=[1, 1],
+            padding=[1, 1],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv3'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv3_b'.format(variable_field)))
+        max_pool3 = fluid.layers.pool2d(
+            input=conv3, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv4 = fluid.layers.conv2d(
+            input=max_pool3,
+            num_filters=64,
+            filter_size=[3, 3],
+            stride=[1, 1],
+            padding=[1, 1],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv4'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv4_b'.format(variable_field)))
+
+        flatten = fluid_flatten(conv4)
+
+        out = fluid.layers.fc(
+            input=flatten,
+            size=self.action_dim,
+            param_attr=ParamAttr(name='{}_fc1'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_fc1_b'.format(variable_field)))
+        return out
+
+    def _build_sync_target_network(self):
+        vars = list(fluid.default_main_program().list_vars())
+        policy_vars = filter(
+            lambda x: 'GRAD' not in x.name and 'policy' in x.name, vars)
+        target_vars = filter(
+            lambda x: 'GRAD' not in x.name and 'target' in x.name, vars)
+        policy_vars.sort(key=lambda x: x.name)
+        target_vars.sort(key=lambda x: x.name)
+
+        sync_program = fluid.default_main_program().clone()
+        with fluid.program_guard(sync_program):
+            sync_ops = []
+            for i, var in enumerate(policy_vars):
+                sync_op = fluid.layers.assign(policy_vars[i], target_vars[i])
+                sync_ops.append(sync_op)
+        sync_program = sync_program.prune(sync_ops)
+        return sync_program
+
+    def act(self, state, train_or_test):
+        sample = np.random.random()
+        if train_or_test == 'train' and sample < self.exploration:
+            act = np.random.randint(self.action_dim)
+        else:
+            if np.random.random() < 0.01:
+                act = np.random.randint(self.action_dim)
+            else:
+                state = np.expand_dims(state, axis=0)
+                pred_Q = self.exe.run(self.predict_program,
+                                      feed={'state': state.astype('float32')},
+                                      fetch_list=[self.pred_value])[0]
+                pred_Q = np.squeeze(pred_Q, axis=0)
+                act = np.argmax(pred_Q)
+        if train_or_test == 'train':
+            self.exploration = max(0.1, self.exploration - 1e-6)
+        return act
+
+    def train(self, state, action, reward, next_state, isOver):
+        if self.global_step % self.update_target_steps == 0:
+            self.sync_target_network()
+        self.global_step += 1
+
+        action = np.expand_dims(action, -1)
+        self.exe.run(self.train_program,
+                     feed={
+                         'state': state.astype('float32'),
+                         'action': action.astype('int32'),
+                         'reward': reward,
+                         'next_s': next_state.astype('float32'),
+                         'isOver': isOver
+                     })
+
+    def sync_target_network(self):
+        self.exe.run(self._sync_program)

fluid/DeepQNetwork/agent.py → fluid/DeepQNetwork/DuelingDQN_agent.py

 #-*- coding: utf-8 -*-
-#File: agent.py
 
 import paddle.fluid as fluid
 from paddle.fluid.param_attr import ParamAttr
 import numpy as np
 from tqdm import tqdm
 import math
+from utils import fluid_flatten
 
-UPDATE_TARGET_STEPS = 200
 
-
-class Model(object):
-    def __init__(self, state_dim, action_dim, gamma):
-        self.global_step = 0
-        self.state_dim = state_dim
+class DuelingDQNModel(object):
+    def __init__(self, state_dim, action_dim, gamma, hist_len, use_cuda=False):
+        self.img_height = state_dim[0]
+        self.img_width = state_dim[1]
         self.action_dim = action_dim
         self.gamma = gamma
-        self.exploration = 1.0
+        self.exploration = 1.1
+        self.update_target_steps = 10000 // 4
+        self.hist_len = hist_len
+        self.use_cuda = use_cuda
 
+        self.global_step = 0
         self._build_net()
 
     def _get_inputs(self):
-        return [fluid.layers.data(\
-                    name='state', shape=[self.state_dim], dtype='float32'),
-                fluid.layers.data(\
-                    name='action', shape=[1], dtype='int32'),
-                fluid.layers.data(\
-                    name='reward', shape=[], dtype='float32'),
-                fluid.layers.data(\
-                    name='next_s', shape=[self.state_dim], dtype='float32'),
-                fluid.layers.data(\
-                  name='isOver', shape=[], dtype='bool')]
+        return fluid.layers.data(
+                   name='state',
+                   shape=[self.hist_len, self.img_height, self.img_width],
+                   dtype='float32'), \
+               fluid.layers.data(
+                   name='action', shape=[1], dtype='int32'), \
+               fluid.layers.data(
+                   name='reward', shape=[], dtype='float32'), \
+               fluid.layers.data(
+                   name='next_s',
+                   shape=[self.hist_len, self.img_height, self.img_width],
+                   dtype='float32'), \
+               fluid.layers.data(
+                   name='isOver', shape=[], dtype='bool')
 
     def _build_net(self):
         state, action, reward, next_s, isOver = self._get_inputs()
         self.pred_value = self.get_DQN_prediction(state)
         self.predict_program = fluid.default_main_program().clone()
 
+        reward = fluid.layers.clip(reward, min=-1.0, max=1.0)
+
         action_onehot = fluid.layers.one_hot(action, self.action_dim)
         action_onehot = fluid.layers.cast(action_onehot, dtype='float32')
 
-        pred_action_value = fluid.layers.reduce_sum(\
-                    fluid.layers.elementwise_mul(action_onehot, self.pred_value), dim=1)
+        pred_action_value = fluid.layers.reduce_sum(
+            fluid.layers.elementwise_mul(action_onehot, self.pred_value), dim=1)
 
         targetQ_predict_value = self.get_DQN_prediction(next_s, target=True)
         best_v = fluid.layers.reduce_max(targetQ_predict_value, dim=1)
         best_v.stop_gradient = True
 
-        target = reward + (1.0 - fluid.layers.cast(\
+        target = reward + (1.0 - fluid.layers.cast(
             isOver, dtype='float32')) * self.gamma * best_v
-        cost = fluid.layers.square_error_cost(\
-            input=pred_action_value, label=target)
+        cost = fluid.layers.square_error_cost(pred_action_value, target)
         cost = fluid.layers.reduce_mean(cost)
 
         self._sync_program = self._build_sync_target_network()
 
-        optimizer = fluid.optimizer.Adam(1e-3)
+        optimizer = fluid.optimizer.Adam(1e-3 * 0.5, epsilon=1e-3)
         optimizer.minimize(cost)
 
         # define program
         self.train_program = fluid.default_main_program()
 
         # fluid exe
-        place = fluid.CUDAPlace(0)
+        place = fluid.CUDAPlace(0) if self.use_cuda else fluid.CPUPlace()
         self.exe = fluid.Executor(place)
         self.exe.run(fluid.default_startup_program())
 
-    def get_DQN_prediction(self, state, target=False):
+    def get_DQN_prediction(self, image, target=False):
+        image = image / 255.0
+
         variable_field = 'target' if target else 'policy'
-        # layer fc1
-        param_attr = ParamAttr(name='{}_fc1'.format(variable_field))
-        bias_attr = ParamAttr(name='{}_fc1_b'.format(variable_field))
-        fc1 = fluid.layers.fc(input=state,
-                              size=256,
-                              act='relu',
-                              param_attr=param_attr,
-                              bias_attr=bias_attr)
-
-        param_attr = ParamAttr(name='{}_fc2'.format(variable_field))
-        bias_attr = ParamAttr(name='{}_fc2_b'.format(variable_field))
-        fc2 = fluid.layers.fc(input=fc1,
-                              size=128,
-                              act='tanh',
-                              param_attr=param_attr,
-                              bias_attr=bias_attr)
-
-        param_attr = ParamAttr(name='{}_fc3'.format(variable_field))
-        bias_attr = ParamAttr(name='{}_fc3_b'.format(variable_field))
-        value = fluid.layers.fc(input=fc2,
-                                size=self.action_dim,
-                                param_attr=param_attr,
-                                bias_attr=bias_attr)
-
-        return value
+
+        conv1 = fluid.layers.conv2d(
+            input=image,
+            num_filters=32,
+            filter_size=[5, 5],
+            stride=[1, 1],
+            padding=[2, 2],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv1'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv1_b'.format(variable_field)))
+        max_pool1 = fluid.layers.pool2d(
+            input=conv1, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv2 = fluid.layers.conv2d(
+            input=max_pool1,
+            num_filters=32,
+            filter_size=[5, 5],
+            stride=[1, 1],
+            padding=[2, 2],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv2'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv2_b'.format(variable_field)))
+        max_pool2 = fluid.layers.pool2d(
+            input=conv2, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv3 = fluid.layers.conv2d(
+            input=max_pool2,
+            num_filters=64,
+            filter_size=[4, 4],
+            stride=[1, 1],
+            padding=[1, 1],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv3'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv3_b'.format(variable_field)))
+        max_pool3 = fluid.layers.pool2d(
+            input=conv3, pool_size=[2, 2], pool_stride=[2, 2], pool_type='max')
+
+        conv4 = fluid.layers.conv2d(
+            input=max_pool3,
+            num_filters=64,
+            filter_size=[3, 3],
+            stride=[1, 1],
+            padding=[1, 1],
+            act='relu',
+            param_attr=ParamAttr(name='{}_conv4'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_conv4_b'.format(variable_field)))
+
+        flatten = fluid_flatten(conv4)
+
+        value = fluid.layers.fc(
+            input=flatten,
+            size=1,
+            param_attr=ParamAttr(name='{}_value_fc'.format(variable_field)),
+            bias_attr=ParamAttr(name='{}_value_fc_b'.format(variable_field)))
+
+        advantage = fluid.layers.fc(
+            input=flatten,
+            size=self.action_dim,
+            param_attr=ParamAttr(name='{}_advantage_fc'.format(variable_field)),
+            bias_attr=ParamAttr(
+                name='{}_advantage_fc_b'.format(variable_field)))
+
+        Q = advantage + (value - fluid.layers.reduce_mean(
+            advantage, dim=1, keep_dim=True))
+        return Q
 
     def _build_sync_target_network(self):
-        vars = fluid.default_main_program().list_vars()
-        policy_vars = []
-        target_vars = []
-        for var in vars:
-            if 'GRAD' in var.name: continue
-            if 'policy' in var.name:
-                policy_vars.append(var)
-            elif 'target' in var.name:
-                target_vars.append(var)
-
-        policy_vars.sort(key=lambda x: x.name.split('policy_')[1])
-        target_vars.sort(key=lambda x: x.name.split('target_')[1])
+        vars = list(fluid.default_main_program().list_vars())
+        policy_vars = filter(
+            lambda x: 'GRAD' not in x.name and 'policy' in x.name, vars)
+        target_vars = filter(
+            lambda x: 'GRAD' not in x.name and 'target' in x.name, vars)
+        policy_vars.sort(key=lambda x: x.name)
+        target_vars.sort(key=lambda x: x.name)
 
         sync_program = fluid.default_main_program().clone()
         with fluid.program_guard(sync_program):
@@ -122,26 +166,30 @@ class Model(object):
         if train_or_test == 'train' and sample < self.exploration:
             act = np.random.randint(self.action_dim)
         else:
-            state = np.expand_dims(state, axis=0)
-            pred_Q = self.exe.run(self.predict_program,
-                                  feed={'state': state.astype('float32')},
-                                  fetch_list=[self.pred_value])[0]
-            pred_Q = np.squeeze(pred_Q, axis=0)
-            act = np.argmax(pred_Q)
-        self.exploration = max(0.1, self.exploration - 1e-6)
+            if np.random.random() < 0.01:
+                act = np.random.randint(self.action_dim)
+            else:
+                state = np.expand_dims(state, axis=0)
+                pred_Q = self.exe.run(self.predict_program,
+                                      feed={'state': state.astype('float32')},
+                                      fetch_list=[self.pred_value])[0]
+                pred_Q = np.squeeze(pred_Q, axis=0)
+                act = np.argmax(pred_Q)
+        if train_or_test == 'train':
+            self.exploration = max(0.1, self.exploration - 1e-6)
         return act
 
     def train(self, state, action, reward, next_state, isOver):
-        if self.global_step % UPDATE_TARGET_STEPS == 0:
+        if self.global_step % self.update_target_steps == 0:
             self.sync_target_network()
         self.global_step += 1
 
         action = np.expand_dims(action, -1)
         self.exe.run(self.train_program, \
-                  feed={'state': state, \
-                        'action': action, \
+                  feed={'state': state.astype('float32'), \
+                        'action': action.astype('int32'), \
                         'reward': reward, \
-                        'next_s': next_state, \
+                        'next_s': next_state.astype('float32'), \
                         'isOver': isOver})
 
     def sync_target_network(self):

fluid/DeepQNetwork/README.md

-<img src="mountain_car.gif" width="300" height="200">
+# Reproduce DQN, DoubleDQN, DuelingDQN model with fluid version of PaddlePaddle
 
-# Reproduce DQN model
- + DQN in:
++ DQN in:
 [Human-level Control Through Deep Reinforcement Learning](http://www.nature.com/nature/journal/v518/n7540/full/nature14236.html)
++ DoubleDQN in:
+[Deep Reinforcement Learning with Double Q-Learning](https://www.aaai.org/ocs/index.php/AAAI/AAAI16/paper/viewPaper/12389)
++ DuelingDQN in:
+[Dueling Network Architectures for Deep Reinforcement Learning](http://proceedings.mlr.press/v48/wangf16.html)
 
-# Mountain-CAR benchmark & performance
-[MountainCar-v0](https://gym.openai.com/envs/MountainCar-v0/)
+# Atari benchmark & performance
+## [Atari games introduction](https://gym.openai.com/envs/#atari)
 
-A car is on a one-dimensional track, positioned between two "mountains". The goal is to drive up the mountain on the right; however, the car's engine is not strong enough to scale the mountain in a single pass. Therefore, the only way to succeed is to drive back and forth to build up momentum.
++ Pong game result
+![DQN result](assets/dqn.png)
 
+# How to use
++ Dependencies:
+    + python2.7
+    + gym
+    + tqdm
+    + paddlepaddle-gpu==0.12.0
 
++ Start Training:
+    ```
+    # To train a model for Pong game with gpu (use DQN model as default)
+    python train.py --rom ./rom_files/pong.bin --use_cuda
 
-<img src="curve.png" >
+    # To train a model for Pong with DoubleDQN
+    python train.py --rom ./rom_files/pong.bin --use_cuda --alg DoubleDQN
 
+    # To train a model for Pong with DuelingDQN
+    python train.py --rom ./rom_files/pong.bin --use_cuda --alg DuelingDQN
+    ```
 
+To train more games, can install more rom files from [here](https://github.com/openai/atari-py/tree/master/atari_py/atari_roms)
 
-# How to use
-+ Dependencies:
-   + python2.7
-   + gym
-   + tqdm
-   + paddle-fluid
-+ Start Training:
-   ```
-   # use mountain-car enviroment as default
-   python DQN.py
++ Start Testing:
+    ```
+    # Play the game with saved model and calculate the average rewards
+    python play.py --rom ./rom_files/pong.bin --use_cuda --model_path ./saved_model/DQN-pong/stepXXXXX
 
-   # use other enviorment
-   python DQN.py --env CartPole-v0
-   ```
+    # Play the game with visualization
+    python play.py --rom ./rom_files/pong.bin --use_cuda --model_path ./saved_model/DQN-pong/stepXXXXX --viz 0.01
+    ```

fluid/DeepQNetwork/atari.py

+# -*- coding: utf-8 -*-
+
+import numpy as np
+import os
+import cv2
+import threading
+
+import gym
+from gym import spaces
+from gym.envs.atari.atari_env import ACTION_MEANING
+
+from ale_python_interface import ALEInterface
+
+__all__ = ['AtariPlayer']
+
+ROM_URL = "https://github.com/openai/atari-py/tree/master/atari_py/atari_roms"
+_ALE_LOCK = threading.Lock()
+"""
+The following AtariPlayer are copied or modified from tensorpack/tensorpack:
+    https://github.com/tensorpack/tensorpack/blob/master/examples/DeepQNetwork/atari.py
+"""
+
+
+class AtariPlayer(gym.Env):
+    """
+    A wrapper for ALE emulator, with configurations to mimic DeepMind DQN settings.
+    Info:
+        score: the accumulated reward in the current game
+        gameOver: True when the current game is Over
+    """
+
+    def __init__(self,
+                 rom_file,
+                 viz=0,
+                 frame_skip=4,
+                 nullop_start=30,
+                 live_lost_as_eoe=True,
+                 max_num_frames=0):
+        """
+        Args:
+            rom_file: path to the rom
+            frame_skip: skip every k frames and repeat the action
+            viz: visualization to be done.
+                Set to 0 to disable.
+                Set to a positive number to be the delay between frames to show.
+                Set to a string to be a directory to store frames.
+            nullop_start: start with random number of null ops.
+            live_losts_as_eoe: consider lost of lives as end of episode. Useful for training.
+            max_num_frames: maximum number of frames per episode.
+        """
+        super(AtariPlayer, self).__init__()
+        assert os.path.isfile(rom_file), \
+            "rom {} not found. Please download at {}".format(rom_file, ROM_URL)
+
+        try:
+            ALEInterface.setLoggerMode(ALEInterface.Logger.Error)
+        except AttributeError:
+            print "You're not using latest ALE"
+
+        # avoid simulator bugs: https://github.com/mgbellemare/Arcade-Learning-Environment/issues/86
+        with _ALE_LOCK:
+            self.ale = ALEInterface()
+            self.ale.setInt(b"random_seed", np.random.randint(0, 30000))
+            self.ale.setInt(b"max_num_frames_per_episode", max_num_frames)
+            self.ale.setBool(b"showinfo", False)
+
+            self.ale.setInt(b"frame_skip", 1)
+            self.ale.setBool(b'color_averaging', False)
+            # manual.pdf suggests otherwise.
+            self.ale.setFloat(b'repeat_action_probability', 0.0)
+
+            # viz setup
+            if isinstance(viz, str):
+                assert os.path.isdir(viz), viz
+                self.ale.setString(b'record_screen_dir', viz)
+                viz = 0
+            if isinstance(viz, int):
+                viz = float(viz)
+            self.viz = viz
+            if self.viz and isinstance(self.viz, float):
+                self.windowname = os.path.basename(rom_file)
+                cv2.startWindowThread()
+                cv2.namedWindow(self.windowname)
+
+            self.ale.loadROM(rom_file.encode('utf-8'))
+        self.width, self.height = self.ale.getScreenDims()
+        self.actions = self.ale.getMinimalActionSet()
+
+        self.live_lost_as_eoe = live_lost_as_eoe
+        self.frame_skip = frame_skip
+        self.nullop_start = nullop_start
+
+        self.action_space = spaces.Discrete(len(self.actions))
+        self.observation_space = spaces.Box(low=0,
+                                            high=255,
+                                            shape=(self.height, self.width),
+                                            dtype=np.uint8)
+        self._restart_episode()
+
+    def get_action_meanings(self):
+        return [ACTION_MEANING[i] for i in self.actions]
+
+    def _grab_raw_image(self):
+        """
+        :returns: the current 3-channel image
+        """
+        m = self.ale.getScreenRGB()
+        return m.reshape((self.height, self.width, 3))
+
+    def _current_state(self):
+        """
+        returns: a gray-scale (h, w) uint8 image
+        """
+        ret = self._grab_raw_image()
+        # avoid missing frame issue: max-pooled over the last screen
+        ret = np.maximum(ret, self.last_raw_screen)
+        if self.viz:
+            if isinstance(self.viz, float):
+                cv2.imshow(self.windowname, ret)
+                cv2.waitKey(int(self.viz * 1000))
+        ret = ret.astype('float32')
+        # 0.299,0.587.0.114. same as rgb2y in torch/image
+        ret = cv2.cvtColor(ret, cv2.COLOR_RGB2GRAY)
+        return ret.astype('uint8')  # to save some memory
+
+    def _restart_episode(self):
+        with _ALE_LOCK:
+            self.ale.reset_game()
+
+        # random null-ops start
+        n = np.random.randint(self.nullop_start)
+        self.last_raw_screen = self._grab_raw_image()
+        for k in range(n):
+            if k == n - 1:
+                self.last_raw_screen = self._grab_raw_image()
+            self.ale.act(0)
+
+    def reset(self):
+        if self.ale.game_over():
+            self._restart_episode()
+        return self._current_state()
+
+    def step(self, act):
+        oldlives = self.ale.lives()
+        r = 0
+        for k in range(self.frame_skip):
+            if k == self.frame_skip - 1:
+                self.last_raw_screen = self._grab_raw_image()
+            r += self.ale.act(self.actions[act])
+            newlives = self.ale.lives()
+            if self.ale.game_over() or \
+                    (self.live_lost_as_eoe and newlives < oldlives):
+                break
+
+        isOver = self.ale.game_over()
+        if self.live_lost_as_eoe:
+            isOver = isOver or newlives < oldlives
+
+        info = {'ale.lives': newlives}
+        return self._current_state(), r, isOver, info

fluid/DeepQNetwork/atari_wrapper.py

+# -*- coding: utf-8 -*-
+
+import numpy as np
+from collections import deque
+
+import gym
+from gym import spaces
+
+_v0, _v1 = gym.__version__.split('.')[:2]
+assert int(_v0) > 0 or int(_v1) >= 10, gym.__version__
+"""
+The following wrappers are copied or modified from openai/baselines:
+https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py
+"""
+
+
+class MapState(gym.ObservationWrapper):
+    def __init__(self, env, map_func):
+        gym.ObservationWrapper.__init__(self, env)
+        self._func = map_func
+
+    def observation(self, obs):
+        return self._func(obs)
+
+
+class FrameStack(gym.Wrapper):
+    def __init__(self, env, k):
+        """Buffer observations and stack across channels (last axis)."""
+        gym.Wrapper.__init__(self, env)
+        self.k = k
+        self.frames = deque([], maxlen=k)
+        shp = env.observation_space.shape
+        chan = 1 if len(shp) == 2 else shp[2]
+        self.observation_space = spaces.Box(low=0,
+                                            high=255,
+                                            shape=(shp[0], shp[1], chan * k),
+                                            dtype=np.uint8)
+
+    def reset(self):
+        """Clear buffer and re-fill by duplicating the first observation."""
+        ob = self.env.reset()
+        for _ in range(self.k - 1):
+            self.frames.append(np.zeros_like(ob))
+        self.frames.append(ob)
+        return self.observation()
+
+    def step(self, action):
+        ob, reward, done, info = self.env.step(action)
+        self.frames.append(ob)
+        return self.observation(), reward, done, info
+
+    def observation(self):
+        assert len(self.frames) == self.k
+        return np.stack(self.frames, axis=0)
+
+
+class _FireResetEnv(gym.Wrapper):
+    def __init__(self, env):
+        """Take action on reset for environments that are fixed until firing."""
+        gym.Wrapper.__init__(self, env)
+        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
+        assert len(env.unwrapped.get_action_meanings()) >= 3
+
+    def reset(self):
+        self.env.reset()
+        obs, _, done, _ = self.env.step(1)
+        if done:
+            self.env.reset()
+        obs, _, done, _ = self.env.step(2)
+        if done:
+            self.env.reset()
+        return obs
+
+    def step(self, action):
+        return self.env.step(action)
+
+
+def FireResetEnv(env):
+    if isinstance(env, gym.Wrapper):
+        baseenv = env.unwrapped
+    else:
+        baseenv = env
+    if 'FIRE' in baseenv.get_action_meanings():
+        return _FireResetEnv(env)
+    return env
+
+
+class LimitLength(gym.Wrapper):
+    def __init__(self, env, k):
+        gym.Wrapper.__init__(self, env)
+        self.k = k
+
+    def reset(self):
+        # This assumes that reset() will really reset the env.
+        # If the underlying env tries to be smart about reset
+        # (e.g. end-of-life), the assumption doesn't hold.
+        ob = self.env.reset()
+        self.cnt = 0
+        return ob
+
+    def step(self, action):
+        ob, r, done, info = self.env.step(action)
+        self.cnt += 1
+        if self.cnt == self.k:
+            done = True
+        return ob, r, done, info

fluid/DeepQNetwork/expreplay.py

-#-*- coding: utf-8 -*-
-#File: expreplay.py
+# -*- coding: utf-8 -*-
 
-from collections import namedtuple
 import numpy as np
+import copy
+from collections import deque, namedtuple
 
 Experience = namedtuple('Experience', ['state', 'action', 'reward', 'isOver'])
 
 
 class ReplayMemory(object):
-    def __init__(self, max_size, state_shape):
+    def __init__(self, max_size, state_shape, context_len):
         self.max_size = int(max_size)
         self.state_shape = state_shape
+        self.context_len = int(context_len)
 
-        self.state = np.zeros((self.max_size, ) + state_shape, dtype='float32')
+        self.state = np.zeros((self.max_size, ) + state_shape, dtype='uint8')
         self.action = np.zeros((self.max_size, ), dtype='int32')
         self.reward = np.zeros((self.max_size, ), dtype='float32')
         self.isOver = np.zeros((self.max_size, ), dtype='bool')
 
         self._curr_size = 0
         self._curr_pos = 0
+        self._context = deque(maxlen=context_len - 1)
 
     def append(self, exp):
+        """append a new experience into replay memory
+        """
         if self._curr_size < self.max_size:
             self._assign(self._curr_pos, exp)
             self._curr_size += 1
         else:
             self._assign(self._curr_pos, exp)
         self._curr_pos = (self._curr_pos + 1) % self.max_size
+        if exp.isOver:
+            self._context.clear()
+        else:
+            self._context.append(exp)
+
+    def recent_state(self):
+        """ maintain recent state for training"""
+        lst = list(self._context)
+        states = [np.zeros(self.state_shape, dtype='uint8')] * \
+                    (self._context.maxlen - len(lst))
+        states.extend([k.state for k in lst])
+        return states
+
+    def sample(self, idx):
+        """ return state, action, reward, isOver,
+            note that some frames in state may be generated from last episode,
+            they should be removed from state
+            """
+        state = np.zeros(
+            (self.context_len + 1, ) + self.state_shape, dtype=np.uint8)
+        state_idx = np.arange(idx, idx + self.context_len + 1) % self._curr_size
+
+        # confirm that no frame was generated from last episode
+        has_last_episode = False
+        for k in range(self.context_len - 2, -1, -1):
+            to_check_idx = state_idx[k]
+            if self.isOver[to_check_idx]:
+                has_last_episode = True
+                state_idx = state_idx[k + 1:]
+                state[k + 1:] = self.state[state_idx]
+                break
+
+        if not has_last_episode:
+            state = self.state[state_idx]
+
+        real_idx = (idx + self.context_len - 1) % self._curr_size
+        action = self.action[real_idx]
+        reward = self.reward[real_idx]
+        isOver = self.isOver[real_idx]
+        return state, reward, action, isOver
+
+    def __len__(self):
+        return self._curr_size
 
     def _assign(self, pos, exp):
         self.state[pos] = exp.state
-        self.action[pos] = exp.action
         self.reward[pos] = exp.reward
+        self.action[pos] = exp.action
         self.isOver[pos] = exp.isOver
 
-    def __len__(self):
-        return self._curr_size
-
-    def sample(self, batch_idx):
-        # index mapping to avoid sampling lastest state
+    def sample_batch(self, batch_size):
+        """sample a batch from replay memory for training
+        """
+        batch_idx = np.random.randint(
+            self._curr_size - self.context_len - 1, size=batch_size)
         batch_idx = (self._curr_pos + batch_idx) % self._curr_size
-        next_idx = (batch_idx + 1) % self._curr_size
-
-        state = self.state[batch_idx]
-        reward = self.reward[batch_idx]
-        action = self.action[batch_idx]
-        next_state = self.state[next_idx]
-        isOver = self.isOver[batch_idx]
-        return (state, action, reward, next_state, isOver)
+        batch_exp = [self.sample(i) for i in batch_idx]
+        return self._process_batch(batch_exp)
+
+    def _process_batch(self, batch_exp):
+        state = np.asarray([e[0] for e in batch_exp], dtype='uint8')
+        reward = np.asarray([e[1] for e in batch_exp], dtype='float32')
+        action = np.asarray([e[2] for e in batch_exp], dtype='int8')
+        isOver = np.asarray([e[3] for e in batch_exp], dtype='bool')
+        return [state, action, reward, isOver]

fluid/DeepQNetwork/play.py

+#-*- coding: utf-8 -*-
+
+import argparse
+import os
+import numpy as np
+import paddle.fluid as fluid
+
+from train import get_player
+from tqdm import tqdm
+
+
+def predict_action(exe, state, predict_program, feed_names, fetch_targets,
+                   action_dim):
+    if np.random.randint(100) == 0:
+        act = np.random.randint(action_dim)
+    else:
+        state = np.expand_dims(state, axis=0)
+        pred_Q = exe.run(predict_program,
+                         feed={feed_names[0]: state.astype('float32')},
+                         fetch_list=fetch_targets)[0]
+        pred_Q = np.squeeze(pred_Q, axis=0)
+        act = np.argmax(pred_Q)
+    return act
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--use_cuda', action='store_true', help='if set, use cuda')
+    parser.add_argument('--rom', type=str, required=True, help='atari rom')
+    parser.add_argument(
+        '--model_path', type=str, required=True, help='dirname to load model')
+    parser.add_argument(
+        '--viz',
+        type=float,
+        default=0,
+        help='''viz: visualization setting:
+                Set to 0 to disable;
+                Set to a positive number to be the delay between frames to show.
+             ''')
+    args = parser.parse_args()
+
+    env = get_player(args.rom, viz=args.viz)
+
+    place = fluid.CUDAPlace(0) if args.use_cuda else fluid.CPUPlace()
+    exe = fluid.Executor(place)
+    inference_scope = fluid.core.Scope()
+    with fluid.scope_guard(inference_scope):
+        [predict_program, feed_names,
+         fetch_targets] = fluid.io.load_inference_model(args.model_path, exe)
+
+        episode_reward = []
+        for _ in tqdm(xrange(30), desc='eval agent'):
+            state = env.reset()
+            total_reward = 0
+            while True:
+                action = predict_action(exe, state, predict_program, feed_names,
+                                        fetch_targets, env.action_space.n)
+                state, reward, isOver, info = env.step(action)
+                total_reward += reward
+                if isOver:
+                    break
+            episode_reward.append(total_reward)
+        eval_reward = np.mean(episode_reward)
+        print('Average reward of 30 epidose: {}'.format(eval_reward))

fluid/DeepQNetwork/train.py

+#-*- coding: utf-8 -*-
+
+from DQN_agent import DQNModel
+from DoubleDQN_agent import DoubleDQNModel
+from DuelingDQN_agent import DuelingDQNModel
+from atari import AtariPlayer
+import paddle.fluid as fluid
+import gym
+import argparse
+import cv2
+from tqdm import tqdm
+from expreplay import ReplayMemory, Experience
+import numpy as np
+import os
+
+from datetime import datetime
+from atari_wrapper import FrameStack, MapState, FireResetEnv, LimitLength
+from collections import deque
+
+UPDATE_FREQ = 4
+
+#MEMORY_WARMUP_SIZE = 2000
+MEMORY_SIZE = 1e6
+MEMORY_WARMUP_SIZE = MEMORY_SIZE // 20
+IMAGE_SIZE = (84, 84)
+CONTEXT_LEN = 4
+ACTION_REPEAT = 4  # aka FRAME_SKIP
+UPDATE_FREQ = 4
+
+
+def run_train_episode(agent, env, exp):
+    total_reward = 0
+    state = env.reset()
+    step = 0
+    while True:
+        step += 1
+        context = exp.recent_state()
+        context.append(state)
+        context = np.stack(context, axis=0)
+        action = agent.act(context, train_or_test='train')
+        next_state, reward, isOver, _ = env.step(action)
+        exp.append(Experience(state, action, reward, isOver))
+        # train model
+        # start training 
+        if len(exp) > MEMORY_WARMUP_SIZE:
+            if step % UPDATE_FREQ == 0:
+                batch_all_state, batch_action, batch_reward, batch_isOver = exp.sample_batch(
+                    args.batch_size)
+                batch_state = batch_all_state[:, :CONTEXT_LEN, :, :]
+                batch_next_state = batch_all_state[:, 1:, :, :]
+                agent.train(batch_state, batch_action, batch_reward,
+                            batch_next_state, batch_isOver)
+        total_reward += reward
+        state = next_state
+        if isOver:
+            break
+    return total_reward, step
+
+
+def get_player(rom, viz=False, train=False):
+    env = AtariPlayer(
+        rom,
+        frame_skip=ACTION_REPEAT,
+        viz=viz,
+        live_lost_as_eoe=train,
+        max_num_frames=60000)
+    env = FireResetEnv(env)
+    env = MapState(env, lambda im: cv2.resize(im, IMAGE_SIZE))
+    if not train:
+        # in training, context is taken care of in expreplay buffer
+        env = FrameStack(env, CONTEXT_LEN)
+    return env
+
+
+def eval_agent(agent, env):
+    episode_reward = []
+    for _ in tqdm(xrange(30), desc='eval agent'):
+        state = env.reset()
+        total_reward = 0
+        step = 0
+        while True:
+            step += 1
+            action = agent.act(state, train_or_test='test')
+            state, reward, isOver, info = env.step(action)
+            total_reward += reward
+            if isOver:
+                break
+        episode_reward.append(total_reward)
+    eval_reward = np.mean(episode_reward)
+    return eval_reward
+
+
+def train_agent():
+    env = get_player(args.rom, train=True)
+    test_env = get_player(args.rom)
+    exp = ReplayMemory(args.mem_size, IMAGE_SIZE, CONTEXT_LEN)
+    action_dim = env.action_space.n
+
+    if args.alg == 'DQN':
+        agent = DQNModel(IMAGE_SIZE, action_dim, args.gamma, CONTEXT_LEN,
+                         args.use_cuda)
+    elif args.alg == 'DoubleDQN':
+        agent = DoubleDQNModel(IMAGE_SIZE, action_dim, args.gamma, CONTEXT_LEN,
+                               args.use_cuda)
+    elif args.alg == 'DuelingDQN':
+        agent = DuelingDQNModel(IMAGE_SIZE, action_dim, args.gamma, CONTEXT_LEN,
+                                args.use_cuda)
+    else:
+        print('Input algorithm name error!')
+        return
+
+    with tqdm(total=MEMORY_WARMUP_SIZE) as pbar:
+        while len(exp) < MEMORY_WARMUP_SIZE:
+            total_reward, step = run_train_episode(agent, env, exp)
+            pbar.update(step)
+
+    # train
+    test_flag = 0
+    save_flag = 0
+    pbar = tqdm(total=1e8)
+    recent_100_reward = []
+    total_step = 0
+    while True:
+        # start epoch
+        total_reward, step = run_train_episode(agent, env, exp)
+        total_step += step
+        pbar.set_description('[train]exploration:{}'.format(agent.exploration))
+        pbar.update(step)
+
+        if total_step // args.test_every_steps == test_flag:
+            pbar.write("testing")
+            eval_reward = eval_agent(agent, test_env)
+            test_flag += 1
+            print("eval_agent done, (steps, eval_reward): ({}, {})".format(
+                total_step, eval_reward))
+
+        if total_step // args.save_every_steps == save_flag:
+            save_flag += 1
+            save_path = os.path.join(args.model_dirname, '{}-{}'.format(
+                args.alg, os.path.basename(args.rom).split('.')[0]),
+                                     'step{}'.format(total_step))
+            fluid.io.save_inference_model(save_path, ['state'],
+                                          agent.pred_value, agent.exe,
+                                          agent.predict_program)
+    pbar.close()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--alg',
+        type=str,
+        default='DQN',
+        help='Reinforcement learning algorithm, support: DQN, DoubleDQN, DuelingDQN'
+    )
+    parser.add_argument(
+        '--use_cuda', action='store_true', help='if set, use cuda')
+    parser.add_argument(
+        '--gamma',
+        type=float,
+        default=0.99,
+        help='discount factor for accumulated reward computation')
+    parser.add_argument(
+        '--mem_size',
+        type=int,
+        default=1000000,
+        help='memory size for experience replay')
+    parser.add_argument(
+        '--batch_size', type=int, default=64, help='batch size for training')
+    parser.add_argument('--rom', help='atari rom', required=True)
+    parser.add_argument(
+        '--model_dirname',
+        type=str,
+        default='saved_model',
+        help='dirname to save model')
+    parser.add_argument(
+        '--save_every_steps',
+        type=int,
+        default=100000,
+        help='every steps number to save model')
+    parser.add_argument(
+        '--test_every_steps',
+        type=int,
+        default=100000,
+        help='every steps number to run test')
+    args = parser.parse_args()
+    train_agent()

fluid/DeepQNetwork/utils.py

+#-*- coding: utf-8 -*-
+#File: utils.py
+
+import paddle.fluid as fluid
+import numpy as np
+
+
+def fluid_argmax(x):
+    """
+    Get index of max value for the last dimension
+    """
+    _, max_index = fluid.layers.topk(x, k=1)
+    return max_index
+
+
+def fluid_flatten(x):
+    """
+    Flatten fluid variable along the first dimension
+    """
+    return fluid.layers.reshape(x, shape=[-1, np.prod(x.shape[1:])])