From d52182882a03d41825d245e15e2e9352c9649d0a Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Tue, 29 Aug 2017 21:52:26 -0600
Subject: [PATCH] Implementation of hybrid reward architecture

---
 README.md            |   5 +-
 agent/A2C_agent.py   |  98 +++++++++++++++
 agent/DQN_agent.py   |  79 +++++++++----
 agent/MSDQN_agent.py | 139 ++++++++++++++++++++++
 agent/__init__.py    |   4 +-
 component/replay.py  |  47 +++++++-
 hybrid.py            | 275 +++++++++++++++++++++++++++++++++++++++++++
 network/network.py   |   7 +-
 utils/config.py      |  10 ++
 9 files changed, 636 insertions(+), 28 deletions(-)
 create mode 100644 agent/A2C_agent.py
 create mode 100644 agent/MSDQN_agent.py
 create mode 100644 hybrid.py

diff --git a/README.md b/README.md
index c8f0c40..b9c2cb8 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,7 @@ Implemented algorithms:
 * Async N-Step Q-Learning
 * Continuous A3C
 * Deep Deterministic Policy Gradient (DDPG)
+* Hybrid Reward Architecture (HRA)
 
 # Curves
 > Curves for CartPole are trivial so I didn't place it here.
@@ -51,7 +52,8 @@ Sometimes _Bipedal Walker_ may run into _NAN_, I'm still not able to totally sol
 * Tensorflow (We need tensorboard)
 
 # Usage
-Detailed usage and all training parameters can be found in ```main.py```
+Detailed usage and all training parameters can be found in ```main.py```,
+For HRA, you may want to look into ```hybrid.py```.
 And you need to create following directories before running the program:
 ```
 cd DeepRL
@@ -68,6 +70,7 @@ mkdir data log evaluation_log
 * [Deterministic Policy Gradient Algorithms](http://proceedings.mlr.press/v32/silver14.pdf)
 * [Continuous control with deep reinforcement learning](https://arxiv.org/abs/1509.02971)
 * [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://arxiv.org/abs/1506.02438)
+* [Hybrid Reward Architecture for Reinforcement Learning](https://arxiv.org/abs/1706.04208)
 * [transedward/pytorch-dqn](https://github.com/transedward/pytorch-dqn)
 * [ikostrikov/pytorch-a3c](https://github.com/ikostrikov/pytorch-a3c)
 * [ghliu/pytorch-ddpg](https://github.com/ghliu/pytorch-ddpg)
diff --git a/agent/A2C_agent.py b/agent/A2C_agent.py
new file mode 100644
index 0000000..aac5806
--- /dev/null
+++ b/agent/A2C_agent.py
@@ -0,0 +1,98 @@
+#######################################################################
+# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
+# Permission given to modify the code as long as you keep this        #
+# declaration at the top                                              #
+#######################################################################
+
+from network import *
+from component import *
+from utils import *
+import numpy as np
+import time
+import os
+import pickle
+import torch
+
+class A2CAgent:
+    def __init__(self, config):
+        self.config = config
+        self.learning_network = config.network_fn()
+        self.optimizer = config.optimizer_fn(self.learning_network.parameters())
+        self.task = config.task_fn()
+        self.replay = config.replay_fn()
+        self.policy = config.policy_fn()
+        self.total_steps = 0
+
+    def episode(self, deterministic=False):
+        state = self.task.reset()
+        total_reward = 0.0
+        steps = 0
+        while not self.config.max_episode_length or steps < self.config.max_episode_length:
+            prob = self.learning_network.predict(np.stack([state]), True)
+            action = self.policy.sample(prob, deterministic=deterministic)
+            next_state, reward, done, info = self.task.step(action)
+            if not deterministic:
+                self.replay.feed([state, action, reward, next_state, int(done)])
+                self.total_steps += 1
+            total_reward += np.sum(reward * self.config.reward_weight)
+            steps += 1
+            state = next_state
+            if done:
+                break
+            if not deterministic and self.total_steps > self.config.min_memory_size:
+                experiences = self.replay.sample()
+                states, actions, rewards, next_states, terminals = experiences
+                prob, log_prob, value = self.learning_network.predict(states, False)
+                _, _, v_next = self.learning_network.predict(next_states, False)
+                terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
+                rewards = self.learning_network.to_torch_variable(rewards).unsqueeze(1)
+                actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
+                target = rewards + self.config.discount * v_next * (1 - terminals)
+                target = target.detach()
+                advantage = target - value
+                value_loss = 0.5 * advantage.pow(2).mean()
+                policy_loss = -(log_prob.gather(1, actions) * Variable(advantage.data)).mean()
+                kl_loss = (prob * log_prob).sum(1).mean()
+
+                self.optimizer.zero_grad()
+                (value_loss + policy_loss + self.config.entropy_weight * kl_loss).backward()
+                torch.nn.utils.clip_grad_norm(self.learning_network.parameters(), self.config.gradient_clip)
+                self.optimizer.step()
+
+        return total_reward, steps
+
+    def run(self):
+        window_size = 100
+        ep = 0
+        rewards = []
+        steps = []
+        avg_test_rewards = []
+        while True:
+            ep += 1
+            reward, step = self.episode()
+            rewards.append(reward)
+            steps.append(step)
+            avg_reward = np.mean(rewards[-window_size:])
+            self.config.logger.info('episode %d, reward %f, avg reward %f, total steps %d, episode step %d' % (
+                ep, reward, avg_reward, self.total_steps, step))
+
+            if self.config.episode_limit and ep > self.config.episode_limit:
+                return rewards, steps, avg_test_rewards
+
+            if self.config.test_interval and ep % self.config.test_interval == 0:
+                self.config.logger.info('Testing...')
+                with open('data/%s-dqn-model-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
+                    pickle.dump(self.learning_network.state_dict(), f)
+                test_rewards = []
+                for _ in range(self.config.test_repetitions):
+                    reward, step = self.episode(True)
+                    test_rewards.append(reward)
+                avg_reward = np.mean(test_rewards)
+                avg_test_rewards.append(avg_reward)
+                self.config.logger.info('Avg reward %f(%f)' % (
+                    avg_reward, np.std(test_rewards) / np.sqrt(self.config.test_repetitions)))
+                with open('data/%sdqn-statistics-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
+                    pickle.dump({'rewards': rewards,
+                                 'test_rewards': avg_test_rewards}, f)
+                if avg_reward > self.task.success_threshold:
+                    break
diff --git a/agent/DQN_agent.py b/agent/DQN_agent.py
index 42f5351..a8d70f2 100644
--- a/agent/DQN_agent.py
+++ b/agent/DQN_agent.py
@@ -11,6 +11,7 @@ import numpy as np
 import time
 import os
 import pickle
+import torch
 
 class DQNAgent:
     def __init__(self, config):
@@ -36,13 +37,14 @@ class DQNAgent:
         total_reward = 0.0
         steps = 0
         while not self.config.max_episode_length or steps < self.config.max_episode_length:
-            value = self.learning_network.predict(np.stack([self.task.normalize_state(state)]), True)
+            value = self.learning_network.predict(np.stack([self.task.normalize_state(state)]), False)
+            value = value.cpu().data.numpy().flatten()
             if deterministic:
-                action = np.argmax(value.flatten())
+                action = np.argmax(value)
             elif self.total_steps < self.config.exploration_steps:
-                action = np.random.randint(0, len(value.flatten()))
+                action = np.random.randint(0, len(value))
             else:
-                action = self.policy.sample(value.flatten())
+                action = self.policy.sample(value)
             next_state, reward, done, info = self.task.step(action)
             self.history_buffer.pop(0)
             self.history_buffer.append(next_state)
@@ -50,7 +52,7 @@ class DQNAgent:
             if not deterministic:
                 self.replay.feed([state, action, reward, next_state, int(done)])
                 self.total_steps += 1
-            total_reward += reward
+            total_reward += np.sum(reward * self.config.reward_weight)
             steps += 1
             state = next_state
             if done:
@@ -60,20 +62,41 @@ class DQNAgent:
                 states, actions, rewards, next_states, terminals = experiences
                 states = self.task.normalize_state(states)
                 next_states = self.task.normalize_state(next_states)
-                q_next = self.target_network.predict(next_states).detach()
-                if self.config.double_q:
-                    _, best_actions = self.learning_network.predict(next_states).detach().max(1)
-                    q_next = q_next.gather(1, best_actions)
+                if self.config.hybrid_reward:
+                    q_next = self.target_network.predict(next_states, True)
+                    target = []
+                    for q_next_ in q_next:
+                        if self.config.target_type == self.config.q_target:
+                            target.append(q_next_.detach().max(1)[0])
+                        elif self.config.target_type == self.config.expected_sarsa_target:
+                            target.append(q_next_.detach().mean(1))
+                    target = torch.cat(target, dim=1).detach()
+                    terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
+                    rewards = self.learning_network.to_torch_variable(rewards)
+                    target = self.config.discount * target * (1 - terminals.expand_as(target))
+                    target.add_(rewards)
+                    q = self.learning_network.predict(states, True)
+                    q_action = []
+                    actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
+                    for q_ in q:
+                        q_action.append(q_.gather(1, actions))
+                    q_action = torch.cat(q_action, dim=1)
+                    loss = self.learning_network.criterion(q_action, target)
                 else:
-                    q_next, _ = q_next.max(1)
-                terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
-                rewards = self.learning_network.to_torch_variable(rewards).unsqueeze(1)
-                q_next = self.config.discount * q_next * (1 - terminals)
-                q_next.add_(rewards)
-                actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
-                q = self.learning_network.predict(states)
-                q = q.gather(1, actions)
-                loss = self.learning_network.criterion(q, q_next)
+                    q_next = self.target_network.predict(next_states, False).detach()
+                    if self.config.double_q:
+                        _, best_actions = self.learning_network.predict(next_states).detach().max(1)
+                        q_next = q_next.gather(1, best_actions)
+                    else:
+                        q_next, _ = q_next.max(1)
+                    terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
+                    rewards = self.learning_network.to_torch_variable(rewards).unsqueeze(1)
+                    q_next = self.config.discount * q_next * (1 - terminals)
+                    q_next.add_(rewards)
+                    actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
+                    q = self.learning_network.predict(states, False)
+                    q = q.gather(1, actions)
+                    loss = self.learning_network.criterion(q, q_next)
                 self.learning_network.zero_grad()
                 loss.backward()
                 self.learning_network.optimizer.step()
@@ -84,20 +107,29 @@ class DQNAgent:
         episode_time = time.time() - episode_start_time
         self.config.logger.debug('episode steps %d, episode time %f, time per step %f' %
                           (steps, episode_time, episode_time / float(steps)))
-        return total_reward
+        return total_reward, steps
 
     def run(self):
         window_size = 100
         ep = 0
         rewards = []
+        steps = []
         avg_test_rewards = []
         while True:
             ep += 1
-            reward = self.episode()
+            reward, step = self.episode()
+            steps.append(step)
             rewards.append(reward)
             avg_reward = np.mean(rewards[-window_size:])
-            self.config.logger.info('episode %d, epsilon %f, reward %f, avg reward %f, total steps %d' % (
-                ep, self.policy.epsilon, reward, avg_reward, self.total_steps))
+            self.config.logger.info('episode %d, epsilon %f, reward %f, avg reward %f, total steps %d, episode step %d' % (
+                ep, self.policy.epsilon, reward, avg_reward, self.total_steps, step))
+            if self.config.episode_limit and ep > self.config.episode_limit:
+                return rewards, steps
+
+            if ep % 100 == 0:
+                with open('data/%s-dqn-statistics-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
+                    pickle.dump({'rewards': rewards,
+                                 'steps': steps}, f)
 
             if self.config.test_interval and ep % self.config.test_interval == 0:
                 self.config.logger.info('Testing...')
@@ -110,8 +142,9 @@ class DQNAgent:
                 avg_test_rewards.append(avg_reward)
                 self.config.logger.info('Avg reward %f(%f)' % (
                     avg_reward, np.std(test_rewards) / np.sqrt(self.config.test_repetitions)))
-                with open('data/%sdqn-statistics-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
+                with open('data/%s-dqn-statistics-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
                     pickle.dump({'rewards': rewards,
+                                 'steps': steps,
                                  'test_rewards': avg_test_rewards}, f)
                 if avg_reward > self.task.success_threshold:
                     break
\ No newline at end of file
diff --git a/agent/MSDQN_agent.py b/agent/MSDQN_agent.py
new file mode 100644
index 0000000..9a69161
--- /dev/null
+++ b/agent/MSDQN_agent.py
@@ -0,0 +1,139 @@
+#######################################################################
+# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
+# Permission given to modify the code as long as you keep this        #
+# declaration at the top                                              #
+#######################################################################
+
+from network import *
+from component import *
+from utils import *
+import numpy as np
+import time
+import os
+import pickle
+import torch
+
+class MSDQNAgent:
+    def __init__(self, config):
+        self.config = config
+        self.learning_network = config.network_fn(config.optimizer_fn)
+        self.target_network = config.network_fn(config.optimizer_fn)
+        self.target_network.load_state_dict(self.learning_network.state_dict())
+        self.task = config.task_fn()
+        self.replay = config.replay_fn()
+        self.policy = config.policy_fn()
+        self.total_steps = 0
+
+    def episode(self, deterministic=False):
+        episode_start_time = time.time()
+        state = self.task.reset()
+        total_reward = 0.0
+        steps = 0
+        while not self.config.max_episode_length or steps < self.config.max_episode_length:
+            value = self.learning_network.predict(np.stack([state]), True)
+            value = value.cpu().data.numpy().flatten()
+            if deterministic:
+                action = np.argmax(value)
+            elif self.total_steps < self.config.exploration_steps:
+                action = np.random.randint(0, len(value))
+            else:
+                action = self.policy.sample(value)
+            next_state, reward, done, info = self.task.step(action)
+            if not deterministic:
+                self.replay.feed([state, action, reward, next_state, int(done)])
+                self.total_steps += 1
+            total_reward += np.sum(reward * self.config.reward_weight)
+            steps += 1
+            state = next_state
+            if done:
+                break
+            if not deterministic and self.total_steps > self.config.exploration_steps:
+                experiences = self.replay.sample()
+                states, actions, rewards, next_states, terminals = experiences
+                if self.config.hybrid_reward:
+                    q_next = self.target_network.predict(next_states, False)
+                    target = []
+                    for q_next_ in q_next:
+                        if self.config.target_type == self.config.q_target:
+                            target.append(q_next_.detach().max(1)[0])
+                        elif self.config.target_type == self.config.expected_sarsa_target:
+                            target.append(q_next_.detach().mean(1))
+                    target = torch.cat(target, dim=1).detach()
+                    terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
+                    rewards = self.learning_network.to_torch_variable(rewards)
+                    target = self.config.discount * target * (1 - terminals.expand_as(target))
+                    target.add_(rewards)
+                    q = self.learning_network.predict(states, False)
+                    q_action = []
+                    actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
+                    for q_ in q:
+                        q_action.append(q_.gather(1, actions))
+                    q_action = torch.cat(q_action, dim=1)
+                    loss = self.learning_network.criterion(q_action, target)
+                else:
+                    q_next = self.target_network.predict(next_states, True).detach()
+                    if self.config.double_q:
+                        _, best_actions = self.learning_network.predict(next_states).detach().max(1)
+                        q_next = q_next.gather(1, best_actions)
+                    else:
+                        q_next, _ = q_next.max(1)
+                    terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
+                    rewards = np.sum(rewards * self.config.reward_weight, axis=1)
+                    rewards = self.learning_network.to_torch_variable(rewards).unsqueeze(1)
+                    q_next = self.config.discount * q_next * (1 - terminals)
+                    q_next.add_(rewards)
+                    actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
+                    q = self.learning_network.predict(states, True)
+                    q = q.gather(1, actions)
+                    loss = self.learning_network.criterion(q, q_next)
+                self.learning_network.zero_grad()
+                loss.backward()
+                self.learning_network.optimizer.step()
+            if not deterministic and self.total_steps % self.config.target_network_update_freq == 0:
+                self.target_network.load_state_dict(self.learning_network.state_dict())
+            if not deterministic and self.total_steps > self.config.exploration_steps:
+                self.policy.update_epsilon()
+        episode_time = time.time() - episode_start_time
+        self.config.logger.debug('episode steps %d, episode time %f, time per step %f' %
+                          (steps, episode_time, episode_time / float(steps)))
+        return total_reward, steps
+
+    def run(self):
+        window_size = 100
+        ep = 0
+        rewards = []
+        steps = []
+        avg_test_rewards = []
+        while True:
+            ep += 1
+            reward, step = self.episode()
+            steps.append(step)
+            rewards.append(reward)
+            avg_reward = np.mean(rewards[-window_size:])
+            self.config.logger.info('episode %d, epsilon %f, reward %f, avg reward %f, total steps %d, episode step %d' % (
+                ep, self.policy.epsilon, reward, avg_reward, self.total_steps, step))
+            if self.config.episode_limit and ep > self.config.episode_limit:
+                return rewards, steps
+
+            if ep % 100 == 0:
+                with open('data/%s-dqn-statistics-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
+                    pickle.dump({'rewards': rewards,
+                                 'steps': steps}, f)
+
+            if self.config.test_interval and ep % self.config.test_interval == 0:
+                self.config.logger.info('Testing...')
+                with open('data/%s-dqn-model-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
+                    pickle.dump(self.learning_network.state_dict(), f)
+                test_rewards = []
+                for _ in range(self.config.test_repetitions):
+                    test_rewards.append(self.episode(True))
+                avg_reward = np.mean(test_rewards)
+                avg_test_rewards.append(avg_reward)
+                self.config.logger.info('Avg reward %f(%f)' % (
+                    avg_reward, np.std(test_rewards) / np.sqrt(self.config.test_repetitions)))
+                with open('data/%s-dqn-statistics-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
+                    pickle.dump({'rewards': rewards,
+                                 'steps': steps,
+                                 'test_rewards': avg_test_rewards}, f)
+                if avg_reward > self.task.success_threshold:
+                    break
\ No newline at end of file
diff --git a/agent/__init__.py b/agent/__init__.py
index 447d3c1..1540617 100644
--- a/agent/__init__.py
+++ b/agent/__init__.py
@@ -1,3 +1,5 @@
 from async_agent import *
 from DDPG_agent import *
-from DQN_agent import *
\ No newline at end of file
+from DQN_agent import *
+from A2C_agent import *
+from MSDQN_agent import *
diff --git a/component/replay.py b/component/replay.py
index 07a0e2b..07dbd1b 100644
--- a/component/replay.py
+++ b/component/replay.py
@@ -49,6 +49,50 @@ class Replay:
                 self.next_states[sampled_indices],
                 self.terminals[sampled_indices]]
 
+class HybridRewardReplay:
+    def __init__(self, memory_size, batch_size, dtype=np.float32):
+        self.memory_size = memory_size
+        self.batch_size = batch_size
+        self.dtype = dtype
+
+        self.states = None
+        self.actions = np.empty(self.memory_size, dtype=np.int8)
+        self.rewards = None
+        self.next_states = None
+        self.terminals = np.empty(self.memory_size, dtype=np.int8)
+
+        self.pos = 0
+        self.full = False
+
+
+    def feed(self, experience):
+        state, action, reward, next_state, done = experience
+
+        if self.states is None:
+            self.rewards = np.empty((self.memory_size, ) + reward.shape, dtype=self.dtype)
+            self.states = np.empty((self.memory_size, ) + state.shape, dtype=self.dtype)
+            self.next_states = np.empty((self.memory_size, ) + state.shape, dtype=self.dtype)
+
+        self.states[self.pos][:] = state
+        self.actions[self.pos] = action
+        self.rewards[self.pos][:] = reward
+        self.next_states[self.pos][:] = next_state
+        self.terminals[self.pos] = done
+
+        self.pos += 1
+        if self.pos == self.memory_size:
+            self.full = True
+            self.pos = 0
+
+    def sample(self):
+        upper_bound = self.memory_size if self.full else self.pos
+        sampled_indices = np.random.randint(0, upper_bound, size=self.batch_size)
+        return [self.states[sampled_indices],
+                self.actions[sampled_indices],
+                self.rewards[sampled_indices],
+                self.next_states[sampled_indices],
+                self.terminals[sampled_indices]]
+
 class HighDimActionReplay:
     def __init__(self, memory_size, batch_size, dtype=np.float32):
         self.memory_size = memory_size
@@ -91,4 +135,5 @@ class HighDimActionReplay:
                 self.actions[sampled_indices],
                 self.rewards[sampled_indices],
                 self.next_states[sampled_indices],
-                self.terminals[sampled_indices]]
\ No newline at end of file
+                self.terminals[sampled_indices]]
+
diff --git a/hybrid.py b/hybrid.py
new file mode 100644
index 0000000..f784136
--- /dev/null
+++ b/hybrid.py
@@ -0,0 +1,275 @@
+import logging
+from agent import *
+from component import *
+from utils import *
+import argparse
+
+class FruitHRFCNet(nn.Module, VanillaNet):
+    def __init__(self, state_dim, action_dim, head_weights, optimizer_fn=None, gpu=True):
+        super(FruitHRFCNet, self).__init__()
+        hidden_size = 250
+        self.fc1 = nn.Linear(state_dim, hidden_size)
+        self.fc2 = nn.ModuleList([nn.Linear(hidden_size, action_dim) for _ in head_weights])
+        self.criterion = nn.MSELoss()
+        self.head_weights = head_weights
+        BasicNet.__init__(self, optimizer_fn, gpu)
+
+    def forward(self, x, heads_only):
+        x = self.to_torch_variable(x)
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.fc1(x))
+        head_q = [fc(x) for fc in self.fc2]
+        if not heads_only:
+            q = [h * w for h, w in zip(head_q, self.head_weights)]
+            q = torch.stack(q, dim=0)
+            q = q.sum(0).squeeze(0)
+            return q
+        else:
+            return head_q
+
+    def predict(self, x, heads_only):
+        return self.forward(x, heads_only)
+
+class FruitMultiStatesFCNet(nn.Module, BasicNet):
+    def __init__(self, state_dim, action_dim, head_weights, optimizer_fn=None, gpu=True):
+        super(FruitMultiStatesFCNet, self).__init__()
+        hidden_size = 250
+        self.fc1 = nn.ModuleList([nn.Linear(state_dim, hidden_size) for _ in head_weights])
+        self.fc2 = nn.ModuleList([nn.Linear(hidden_size, action_dim) for _ in head_weights])
+        self.criterion = nn.MSELoss()
+        self.head_weights = head_weights
+        self.state_dim = state_dim
+        self.n_heads = head_weights.shape[0]
+        BasicNet.__init__(self, optimizer_fn, gpu)
+
+    def predict(self, x, merge):
+        head_q = []
+        for i in range(self.n_heads):
+            q = self.to_torch_variable(x[:, i, :])
+            q = self.fc1[i](q)
+            q = F.relu(q)
+            q = self.fc2[i](q)
+            head_q.append(q)
+        if merge:
+            q = [q * w for q, w in zip(head_q, self.head_weights)]
+            q = torch.stack(q, dim=0)
+            q = q.sum(0).squeeze(0)
+            return q
+        return head_q
+
+FRUIT_EPISDOE_LENGTH = 100
+MAX_EPISODES = 7000
+
+class Fruit(BasicTask):
+    def __init__(self, hybrid_reward=False, pseudo_reward=False, atomic_state=True):
+        self.hybrid_reward = hybrid_reward
+        self.atomic_state = atomic_state
+        self.pseudo_reward = pseudo_reward
+        self.name = "Fruit"
+        self.success_threshold = 5
+        self.width = 10
+        self.height = 10
+        self.possible_fruits = 10
+        self.actual_fruits = 5
+        xs = np.random.randint(0, self.width, size=self.possible_fruits)
+        ys = np.random.randint(0, self.height, size=self.possible_fruits)
+        self.possible_locations = list(zip(xs, ys))
+        self.x = 0
+        self.y = 0
+        self.indices = np.arange(self.possible_fruits)
+        self.taken = []
+        self.remaining_fruits = 0
+
+    def get_nearest(self):
+        def distance(i):
+            x, y = self.possible_locations[i]
+            return np.abs(self.x - x) + np.abs(self.y - y)
+        pool = []
+        for i in range(self.possible_fruits):
+            if not self.taken[i]:
+                pool.append([i, distance(i)])
+        pool = sorted(pool, key=lambda x:x[1])
+        return pool[0][0]
+
+    def encode_pos(self, x, y):
+        return '{:04b}'.format(x) + '{:04b}'.format(y)
+
+    def encode_atomic_state(self):
+        offset = 8 * self.possible_fruits
+        state = np.copy(self.base_state)
+        str = self.encode_pos(self.x, self.y)
+        for i in range(len(str)):
+            state[offset + i] = int(str[i])
+        offset += 8
+        for i in range(len(self.taken)):
+            state[offset + i] = self.taken[i]
+        return state
+
+    def encode_decomposed_state(self):
+        state_size = (4 + 4) * 2 + 1
+        base_state = np.zeros(state_size)
+        str = self.encode_pos(self.x, self.y)
+        for i in range(len(str)):
+            base_state[i] = int(str[i])
+        states = []
+        for i in range(self.possible_fruits):
+            states.append(np.copy(base_state))
+            str = self.encode_pos(*self.possible_locations[i])
+            for j in range(len(str)):
+                states[-1][8 + j] = int(str[j])
+            states[-1][-1] = self.taken[i]
+        return np.asarray(states)
+
+    def encode_state(self):
+        if self.atomic_state:
+            return self.encode_atomic_state()
+        return self.encode_decomposed_state()
+
+    def reset(self):
+        self.x = np.random.randint(0, self.width)
+        self.y = np.random.randint(0, self.height)
+        np.random.shuffle(self.indices)
+        self.taken = np.ones(self.possible_fruits, dtype=np.bool)
+        self.taken[self.indices[: self.actual_fruits]] = False
+        self.remaining_fruits = self.actual_fruits
+        state_size = (4 + 4) * (self.possible_fruits + 1) + self.possible_fruits
+        self.base_state = np.zeros(state_size)
+        offset = 0
+        for x, y in self.possible_locations:
+            str = self.encode_pos(x, y)
+            for i in range(len(str)):
+                self.base_state[offset + i] = int(str[i])
+            offset += 8
+        return self.encode_state()
+
+    def step(self, action):
+        # action = action[0]
+        if action == 0:
+            self.x -= 1
+        elif action == 1:
+            self.x += 1
+        elif action == 2:
+            self.y -= 1
+        elif action == 3:
+            self.y += 1
+        else:
+            assert False
+        self.x = min(max(self.x, 0), self.width - 1)
+        self.y = min(max(self.y, 0), self.height - 1)
+        try:
+            pos = self.possible_locations.index((self.x, self.y))
+        except ValueError:
+            pos = -1
+        if self.hybrid_reward:
+            reward = np.zeros(self.possible_fruits)
+            if pos >= 0 and not self.taken[pos]:
+                reward[pos] = 1
+                self.taken[pos] = True
+                self.remaining_fruits -= 1
+            if self.pseudo_reward:
+                pseudo_reward = np.zeros(self.possible_fruits)
+                if pos >= 0:
+                    pseudo_reward[pos] = 1
+                reward = (reward, pseudo_reward)
+        else:
+            reward = 0.0
+            if pos >= 0 and not self.taken[pos]:
+                reward = 1.0
+                self.taken[pos] = True
+                self.remaining_fruits -= 1
+        return self.encode_state(), reward, not self.remaining_fruits, self.taken
+
+BATCH_SIZE = 15
+
+def dqn_fruit(args):
+    config = Config()
+    config.task_fn = lambda: Fruit()
+    config.optimizer_fn = lambda params: torch.optim.Adam(params, args.lr)
+    # config.optimizer_fn = lambda params: torch.optim.SGD(params, lr)
+    config.reward_weight = np.ones(10) / 10
+    config.hybrid_reward = False
+    config.network_fn = lambda optimizer_fn: FruitHRFCNet(
+        98, 4, config.reward_weight, optimizer_fn)
+    config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1)
+    config.replay_fn = lambda: Replay(memory_size=10000, batch_size=BATCH_SIZE)
+    config.discount = 0.95
+    config.target_network_update_freq = 200
+    config.max_episode_length = FRUIT_EPISDOE_LENGTH
+    config.exploration_steps = 200
+    config.logger = Logger('./log', gym.logger)
+    config.history_length = 1
+    config.test_interval = 0
+    config.test_repetitions = 10
+    config.episode_limit = 5000
+    config.tag = 'vanilla-%f' % (args.lr)
+    config.double_q = False
+    agent = DQNAgent(config)
+    return agent
+
+def hrdqn_fruit(args):
+    config = Config()
+    config.task_fn = lambda: Fruit(hybrid_reward=True)
+    config.hybrid_reward = True
+    config.reward_weight = np.ones(10) / 10
+    config.optimizer_fn = lambda params: torch.optim.Adam(params, args.lr)
+    # config.optimizer_fn = lambda params: torch.optim.SGD(params, 0.01)
+    config.network_fn = lambda optimizer_fn: FruitHRFCNet(
+        98, 4, config.reward_weight, optimizer_fn)
+    config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1)
+    config.replay_fn = lambda: HybridRewardReplay(memory_size=10000, batch_size=BATCH_SIZE)
+    config.discount = 0.95
+    config.target_network_update_freq = 200
+    config.max_episode_length = FRUIT_EPISDOE_LENGTH
+    config.exploration_steps = 200
+    config.logger = Logger('./log', gym.logger)
+    config.history_length = 1
+    config.test_interval = 0
+    config.test_repetitions = 10
+    config.target_type = config.expected_sarsa_target
+    config.tag = 'expected_sarsa-%f-%s' % (args.lr, args.tag)
+    # config.target_type = config.q_target
+    # config.tag = 'q-%f-%s' % (args.lr, args.tag)
+    config.double_q = False
+    config.episode_limit = 5000
+    agent = DQNAgent(config)
+    return agent
+
+def hrmsdqn_fruit(args):
+    config = Config()
+    config.task_fn = lambda: Fruit(hybrid_reward=True, atomic_state=False)
+    config.hybrid_reward = True
+    # config.hybrid_reward = False
+    config.reward_weight = np.ones(10) / 10
+    # config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
+    config.optimizer_fn = lambda params: torch.optim.SGD(params, 0.1, momentum=0.9)
+    config.network_fn = lambda optimizer_fn: FruitMultiStatesFCNet(
+        17, 4, config.reward_weight, optimizer_fn)
+    config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1)
+    config.replay_fn = lambda: HybridRewardReplay(memory_size=10000, batch_size=BATCH_SIZE)
+    config.discount = 0.95
+    config.target_network_update_freq = 200
+    config.max_episode_length = FRUIT_EPISDOE_LENGTH
+    config.exploration_steps = 200
+    config.logger = Logger('./log', gym.logger)
+    config.history_length = 1
+    config.test_interval = 0
+    config.test_repetitions = 10
+    config.target_type = config.expected_sarsa_target
+    config.tag = 'expected_sarsa-%f-%s' % (args.lr, args.tag)
+    # config.target_type = config.q_target
+    # config.tag = 'q-%f-%s' % (args.lr, args.tag)
+    config.double_q = False
+    config.episode_limit = 5000
+    agent = MSDQNAgent(config)
+    return agent
+
+if __name__ == '__main__':
+    # gym.logger.setLevel(logging.DEBUG)
+    gym.logger.setLevel(logging.INFO)
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--lr', type=float, default=0.001)
+    parser.add_argument('--tag', type=str, default='none')
+    args = parser.parse_args()
+    agent = hrdqn_fruit(args)
+    # agent = hrmsdqn_fruit(args)
+    agent.run()
diff --git a/network/network.py b/network/network.py
index e995f3d..39be784 100644
--- a/network/network.py
+++ b/network/network.py
@@ -43,12 +43,15 @@ class VanillaNet(BasicNet):
     def predict(self, x, to_numpy=False):
         y = self.forward(x)
         if to_numpy:
-            y = y.cpu().data.numpy()
+            if type(y) is list:
+                y = [y_.cpu().data.numpy() for y_ in y]
+            else:
+                y = y.cpu().data.numpy()
         return y
 
 # Base class for actor critic method
 class ActorCriticNet(BasicNet):
-    def predict(self, x):
+    def predict(self, x, _):
         phi = self.forward(x, True)
         pre_prob = self.fc_actor(phi)
         prob = F.softmax(pre_prob)
diff --git a/utils/config.py b/utils/config.py
index c9c88fc..186ef61 100644
--- a/utils/config.py
+++ b/utils/config.py
@@ -5,6 +5,8 @@
 #######################################################################
 
 class Config:
+    q_target = 0
+    expected_sarsa_target = 1
     def __init__(self):
         self.task_fn = None
         self.optimizer_fn = None
@@ -37,3 +39,11 @@ class Config:
         self.reward_shift_fn = lambda r: r
         self.state_shift_fn = lambda s: s
         self.action_shift_fn = lambda a: a
+        self.reward_weight = 1
+        self.hybrid_reward = False
+        self.target_type = self.q_target
+        self.episode_limit = 0
+        self.min_memory_size = 200
+        self.master_fn = None
+        self.master_optimizer_fn = None
+        self.num_heads = 10