Major refactor

agent/DDPG_agent.py

@@ -0,0 +1,146 @@
+#######################################################################
+# 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 pickle
+
+class DDPGAgent:
+    def __init__(self,
+                 task_fn,
+                 actor_network_fn,
+                 critic_network_fn,
+                 actor_optimizer_fn,
+                 critic_optimizer_fn,
+                 replay_fn,
+                 discount,
+                 step_limit,
+                 tau,
+                 exploration_steps,
+                 random_process_fn,
+                 test_interval,
+                 test_repetitions,
+                 noise_decay_steps,
+                 tag,
+                 logger):
+        self.task = task_fn()
+        self.actor = actor_network_fn()
+        self.critic = critic_network_fn()
+        self.target_actor = actor_network_fn()
+        self.target_critic = critic_network_fn()
+        self.target_actor.load_state_dict(self.actor.state_dict())
+        self.target_critic.load_state_dict(self.critic.state_dict())
+        self.actor_opt = actor_optimizer_fn(self.actor.parameters())
+        self.critic_opt = critic_optimizer_fn(self.critic.parameters())
+        self.replay = replay_fn()
+        self.step_limit = step_limit
+        self.tau = tau
+        self.logger = logger
+        self.discount = discount
+        self.exploration_steps = exploration_steps
+        self.random_process = random_process_fn()
+        self.criterion = nn.MSELoss()
+        self.test_interval = test_interval
+        self.test_repetitions = test_repetitions
+        self.total_steps = 0
+        self.tag = tag
+        self.epsilon = 1.0
+        self.d_epsilon = 1.0 / noise_decay_steps
+
+    def soft_update(self, target, src):
+        for target_param, param in zip(target.parameters(), src.parameters()):
+            target_param.data.copy_(
+                target_param.data * (1.0 - self.tau) + param.data * self.tau
+            )
+
+    def episode(self, deterministic=False):
+        self.random_process.reset_states()
+        state = self.task.reset()
+
+        steps = 0
+        total_reward = 0.0
+        while not self.step_limit or steps < self.step_limit:
+            action = self.actor.predict(np.stack([state])).flatten()
+            self.logger.histo_summary('action', action, self.total_steps)
+            if not deterministic:
+                if self.total_steps < self.exploration_steps:
+                    action = self.task.random_action()
+                else:
+                    action += max(self.epsilon, 0) * self.random_process.sample()
+            self.logger.histo_summary('noised action', action, self.total_steps)
+            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
+                self.epsilon -= self.d_epsilon
+            steps += 1
+            total_reward += reward
+            state = next_state
+
+            if done:
+                break
+
+            if not deterministic and self.total_steps > self.exploration_steps:
+                experiences = self.replay.sample()
+                states, actions, rewards, next_states, terminals = experiences
+                q_next = self.target_critic.predict(next_states, self.target_actor.predict(next_states))
+                terminals = self.critic.to_torch_variable(terminals).unsqueeze(1)
+                rewards = self.critic.to_torch_variable(rewards).unsqueeze(1)
+                q_next = self.discount * q_next * (1 - terminals)
+                q_next.add_(rewards)
+                q_next = Variable(q_next.data)
+                q = self.critic.predict(states, actions)
+                critic_loss = self.criterion(q, q_next)
+
+                self.critic.zero_grad()
+                critic_loss.backward()
+                self.critic_opt.step()
+
+                actor_loss = -self.critic.predict(states, self.actor.predict(states, False))
+                actor_loss = actor_loss.mean()
+
+                self.actor.zero_grad()
+                actor_loss.backward()
+                self.actor_opt.step()
+
+                self.soft_update(self.target_actor, self.actor)
+                self.soft_update(self.target_critic, self.critic)
+
+        return total_reward
+
+    def save(self, file_name):
+        with open(file_name, 'wb') as f:
+            pickle.dump(self.actor.state_dict(), f)
+
+    def run(self):
+        window_size = 100
+        ep = 0
+        rewards = []
+        avg_test_rewards = []
+        while True:
+            ep += 1
+            reward = self.episode()
+            rewards.append(reward)
+            avg_reward = np.mean(rewards[-window_size:])
+            self.logger.info('episode %d, reward %f, avg reward %f, total steps %d' % (
+                ep, reward, avg_reward, self.total_steps))
+
+            if self.test_interval and ep % self.test_interval == 0:
+                self.logger.info('Testing...')
+                self.save('data/%sddpg-model-%s.bin' % (self.tag, self.task.name))
+                test_rewards = []
+                for _ in range(self.test_repetitions):
+                    test_rewards.append(self.episode(True))
+                avg_reward = np.mean(test_rewards)
+                avg_test_rewards.append(avg_reward)
+                self.logger.info('Avg reward %f(%f)' % (
+                    avg_reward, np.std(test_rewards) / np.sqrt(self.test_repetitions)))
+                with open('data/%sddpg-statistics-%s.bin' % (self.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

agent/DQN_agent.py

@@ -0,0 +1,144 @@
+#######################################################################
+# 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
+
+class DQNAgent:
+    def __init__(self,
+                 task_fn,
+                 network_fn,
+                 optimizer_fn,
+                 policy_fn,
+                 replay_fn,
+                 discount,
+                 step_limit,
+                 target_network_update_freq,
+                 explore_steps,
+                 history_length,
+                 double_q,
+                 test_interval,
+                 test_repetitions,
+                 tag,
+                 logger):
+        self.learning_network = network_fn(optimizer_fn)
+        self.target_network = network_fn(optimizer_fn)
+        self.target_network.load_state_dict(self.learning_network.state_dict())
+        self.task = task_fn()
+        self.step_limit = step_limit
+        self.replay = replay_fn()
+        self.discount = discount
+        self.target_network_update_freq = target_network_update_freq
+        self.policy = policy_fn()
+        self.total_steps = 0
+        self.explore_steps = explore_steps
+        self.history_length = history_length
+        self.logger = logger
+        self.test_interval = test_interval
+        self.test_repetitions = test_repetitions
+        self.history_buffer = None
+        self.double_q = double_q
+        self.tag = tag
+
+    def episode(self, deterministic=False):
+        episode_start_time = time.time()
+        state = self.task.reset()
+        if self.history_buffer is None:
+            self.history_buffer = [np.zeros_like(state)] * self.history_length
+        else:
+            self.history_buffer.pop(0)
+            self.history_buffer.append(state)
+        state = np.vstack(self.history_buffer)
+        total_reward = 0.0
+        steps = 0
+        while not self.step_limit or steps < self.step_limit:
+            value = self.learning_network.predict(np.stack([self.task.normalize_state(state)]), True)
+            if deterministic:
+                action = np.argmax(value.flatten())
+            elif self.total_steps < self.explore_steps:
+                action = np.random.randint(0, len(value.flatten()))
+            else:
+                action = self.policy.sample(value.flatten())
+            next_state, reward, done, info = self.task.step(action)
+            self.history_buffer.pop(0)
+            self.history_buffer.append(next_state)
+            next_state = np.vstack(self.history_buffer)
+            if not deterministic:
+                self.replay.feed([state, action, reward, next_state, int(done)])
+                self.total_steps += 1
+            total_reward += reward
+            steps += 1
+            state = next_state
+            if done:
+                break
+            if not deterministic and self.total_steps > self.explore_steps:
+                experiences = self.replay.sample()
+                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.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.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)
+                self.learning_network.zero_grad()
+                loss.backward()
+                self.learning_network.optimizer.step()
+            if not deterministic and self.total_steps % self.target_network_update_freq == 0:
+                self.target_network.load_state_dict(self.learning_network.state_dict())
+            if not deterministic and self.total_steps > self.explore_steps:
+                self.policy.update_epsilon()
+        episode_time = time.time() - episode_start_time
+        self.logger.debug('episode steps %d, episode time %f, time per step %f' %
+                          (steps, episode_time, episode_time / float(steps)))
+        return total_reward
+
+    def save(self, file_name):
+        with open(file_name, 'wb') as f:
+            pickle.dump(self.learning_network.state_dict(), f)
+
+    def run(self):
+        window_size = 100
+        ep = 0
+        rewards = []
+        avg_test_rewards = []
+        while True:
+            ep += 1
+            reward = self.episode()
+            rewards.append(reward)
+            avg_reward = np.mean(rewards[-window_size:])
+            self.logger.info('episode %d, epsilon %f, reward %f, avg reward %f, total steps %d' % (
+                ep, self.policy.epsilon, reward, avg_reward, self.total_steps))
+
+            if self.test_interval and ep % self.test_interval == 0:
+                self.logger.info('Testing...')
+                self.save('data/%sdqn-model-%s.bin' % (self.tag, self.task.name))
+                test_rewards = []
+                for _ in range(self.test_repetitions):
+                    test_rewards.append(self.episode(True))
+                avg_reward = np.mean(test_rewards)
+                avg_test_rewards.append(avg_reward)
+                self.logger.info('Avg reward %f(%f)' % (
+                    avg_reward, np.std(test_rewards) / np.sqrt(self.test_repetitions)))
+                with open('data/%sdqn-statistics-%s.bin' % (self.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

agent/__init__.py

@@ -0,0 +1,3 @@
+from async_agent import *
+from DDPG_agent import *
+from DQN_agent import *

agent/async_agent.py

@@ -0,0 +1,92 @@
+#######################################################################
+# 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                                              #
+#######################################################################
+
+import numpy as np
+import torch.multiprocessing as mp
+from network import *
+from utils import *
+from component import *
+from async_worker import *
+import pickle
+import os
+import time
+
+class AsyncAgent:
+    def __init__(self, config):
+        self.config = config
+        learning_network = config.network_fn()
+        learning_network.share_memory()
+        target_network = config.network_fn()
+        target_network.share_memory()
+        target_network.load_state_dict(learning_network.state_dict())
+
+        self.task = config.task_fn()
+
+        self.config.learning_network = learning_network
+        self.config.target_network = target_network
+        self.config.steps_lock = mp.Lock()
+        self.config.network_lock = mp.Lock()
+        self.config.total_steps = mp.Value('i', 0)
+        self.config.stop_signal = mp.Value('i', False)
+
+    def train(self, id):
+        worker = self.config.worker(self.config)
+        episode = 0
+        rewards = []
+        while not self.config.stop_signal.value:
+            steps, reward = worker.episode()
+            rewards.append(reward)
+            if len(rewards) > 100: rewards.pop(0)
+            self.config.logger.debug('worker %d, episode %d, return %f, avg return %f, episode steps %d, total steps %d' % (
+                id, episode, rewards[-1], np.mean(rewards[-100:]), steps, self.config.total_steps.value))
+
+    def save(self, file_name):
+        with open(file_name, 'wb') as f:
+            pickle.dump(self.config.learning_network.state_dict(), f)
+
+    def evaluate(self, id):
+        test_rewards = []
+        test_points = []
+        worker = self.config.worker(self.config)
+        while True:
+            steps = self.config.total_steps.value
+            if steps % self.config.test_interval == 0:
+                worker.worker_network.load_state_dict(self.config.learning_network.state_dict())
+                self.save('data/%s-%s-model-%s.bin' % (
+                    self.config.tag, self.config.worker.__name__, self.task.name))
+                rewards = np.zeros(self.config.test_repetitions)
+                for i in range(self.config.test_repetitions):
+                    rewards[i] = worker.episode(deterministic=True)[1]
+                self.config.logger.info('total steps: %d, averaged return per episode: %f(%f)' %\
+                      (steps, np.mean(rewards), np.std(rewards) / np.sqrt(self.config.test_repetitions)))
+                test_rewards.append(np.mean(rewards))
+                test_points.append(steps)
+                with open('data/%s-%s-statistics-%s.bin' % (
+                    self.config.tag, self.config.worker.__name__, self.task.name
+                ), 'wb') as f:
+                    pickle.dump([test_points, test_rewards], f)
+                if np.mean(rewards) > self.task.success_threshold:
+                    self.config.stop_signal.value = True
+                    break
+
+    def run(self):
+        os.environ['OMP_NUM_THREADS'] = '1'
+        procs = [mp.Process(target=self.train, args=(i, )) for i in range(self.config.num_workers)]
+        procs.append(mp.Process(target=self.evaluate, args=(self.config.num_workers, )))
+        for p in procs: p.start()
+        while True:
+            time.sleep(1)
+            for i, p in enumerate(procs):
+                if not p.is_alive() and not self.config.stop_signal.value:
+                    self.config.logger.warning('Worker %d exited unexpectedly.' % i)
+                    p.terminate()
+                    procs[i] = mp.Process(target=self.train, args=(i, ))
+                    procs[i].start()
+                    self.config.logger.warning('Worker %d restarted.' % i)
+                    break
+            if self.config.stop_signal.value:
+                break
+        for p in procs: p.join()