Merge branch 'master' of https://github.com/ShangtongZhang/DeepRL

deep_rl/agent/OptionCritic_agent.py

@@ -0,0 +1,115 @@
+#######################################################################
+# 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 .BaseAgent import *
+
+class OptionCriticAgent(BaseAgent):
+    def __init__(self, config):
+        BaseAgent.__init__(self, config)
+        self.config = config
+        self.task = config.task_fn()
+        self.network = config.network_fn(self.task.state_dim, self.task.action_dim)
+        self.target_network = config.network_fn(self.task.state_dim, self.task.action_dim)
+        self.optimizer = config.optimizer_fn(self.network.parameters())
+        self.target_network.load_state_dict(self.network.state_dict())
+        self.policy = config.policy_fn()
+
+        self.episode_rewards = np.zeros(config.num_workers)
+        self.last_episode_rewards = np.zeros(config.num_workers)
+
+        self.total_steps = 0
+        states = self.config.state_normalizer(self.task.reset())
+        self.q_options, self.betas, self.log_pi = self.network.predict(states)
+        self.options = np.asarray([self.policy.sample(q) for q in self.q_options.detach().cpu().numpy()])
+        self.is_initial_betas = np.ones(self.config.num_workers)
+        self.prev_options = np.copy(self.options)
+
+    def iteration(self):
+        config = self.config
+        rollout = []
+
+        q_options, betas, options, log_pi = self.q_options, self.betas, self.options, self.log_pi
+        for _ in range(config.rollout_length):
+            var_options = self.network.tensor(options).long()
+            worker_index = self.network.tensor(np.arange(config.num_workers)).long()
+            intra_log_pi = log_pi[worker_index, var_options, :]
+            dist = torch.distributions.Categorical(intra_log_pi.exp())
+            actions = dist.sample()
+            next_states, rewards, terminals, _ = self.task.step(actions.cpu().detach().numpy().flatten())
+            next_states = config.state_normalizer(next_states)
+            self.episode_rewards += rewards
+            rewards = config.reward_normalizer(rewards)
+            q_options_next, betas_next, log_pi_next = self.network.predict(next_states)
+            rollout.append([q_options, betas, options, self.prev_options, rewards, 1 - terminals, np.copy(self.is_initial_betas), intra_log_pi, actions])
+            self.is_initial_betas = np.asarray(terminals, dtype=np.float32)
+
+            np_q_options_next = q_options_next.cpu().detach().numpy()
+            np_betas_next = betas_next.gather(1, var_options.unsqueeze(1)).cpu().detach().numpy().flatten()
+            options_next = np.copy(options)
+            dice = np.random.rand(len(options_next))
+            for j in range(len(dice)):
+                if dice[j] < np_betas_next[j]:
+                    options_next[j] = self.policy.sample(np_q_options_next[j])
+            for i, terminal in enumerate(terminals):
+                if terminals[i]:
+                    self.last_episode_rewards[i] = self.episode_rewards[i]
+                    self.episode_rewards[i] = 0
+            self.prev_options = options
+            options = options_next
+            q_options = q_options_next
+            betas = betas_next
+            log_pi = log_pi_next
+
+            self.policy.update_epsilon()
+            self.total_steps += config.num_workers
+            if self.total_steps / config.num_workers % config.target_network_update_freq == 0:
+                self.target_network.load_state_dict(self.network.state_dict())
+
+        self.options = options
+        self.q_options = q_options
+        self.betas = betas
+        self.log_pi = log_pi
+
+        target_q_options, _, _ = self.target_network.predict(next_states)
+        prev_options = self.network.tensor(self.prev_options).long().unsqueeze(1)
+        betas_prev_options = betas.gather(1, prev_options)
+
+        returns = (1 - betas_prev_options) * target_q_options.gather(1, prev_options) +\
+                  betas_prev_options * torch.max(target_q_options, dim=1, keepdim=True)[0]
+        returns = returns.detach()
+
+        processed_rollout = [None] * (len(rollout))
+        for i in reversed(range(len(rollout))):
+            q_options, betas, options, prev_options, rewards, terminals, is_initial_betas, log_pi, actions = rollout[i]
+            options = self.network.tensor(options).unsqueeze(1).long()
+            prev_options = self.network.tensor(prev_options).unsqueeze(1).long()
+            terminals = self.network.tensor(terminals).unsqueeze(1)
+            rewards = self.network.tensor(rewards).unsqueeze(1)
+            is_initial_betas = self.network.tensor(is_initial_betas).unsqueeze(1)
+            returns = rewards + config.discount * terminals * returns
+
+            q_omg = q_options.gather(1, options)
+            log_action_prob = log_pi.gather(1, actions.unsqueeze(1))
+            entropy_loss = (log_pi.exp() * log_pi).sum(-1).unsqueeze(1)
+
+            q_prev_omg = q_options.gather(1, prev_options)
+            v_prev_omg = torch.max(q_options, dim=1, keepdim=True)[0]
+            advantage_omg = q_prev_omg - v_prev_omg
+            advantage_omg.add_(config.termination_regularizer)
+            betas = betas.gather(1, prev_options)
+            betas = betas * (1 - is_initial_betas)
+            processed_rollout[i] = [q_omg, returns, betas, advantage_omg.detach(), log_action_prob, entropy_loss]
+
+        q_omg, returns, beta_omg, advantage_omg, log_action_prob, entropy_loss = map(lambda x: torch.cat(x, dim=0), zip(*processed_rollout))
+        pi_loss = -log_action_prob * (returns - q_omg.detach()) + config.entropy_weight * entropy_loss
+        pi_loss = pi_loss.mean()
+        q_loss = 0.5 * (q_omg - returns).pow(2).mean()
+        beta_loss = (advantage_omg * beta_omg).mean()
+        self.optimizer.zero_grad()
+        (pi_loss + q_loss + beta_loss).backward()
+        nn.utils.clip_grad_norm_(self.network.parameters(), config.gradient_clip)
+        self.optimizer.step()

deep_rl/agent/__init__.py

@@ -5,3 +5,4 @@ from .CategoricalDQN_agent import *
 from .NStepDQN_agent import *
 from .QuantileRegressionDQN_agent import *
 from .PPO_agent import *
+from .OptionCritic_agent import *

deep_rl/network/network_heads.py

@@ -89,6 +89,26 @@ class QuantileNet(nn.Module, BaseNet):
             quantiles = quantiles.cpu().detach().numpy()
         return quantiles
 
+class OptionCriticNet(nn.Module, BaseNet):
+    def __init__(self, body, action_dim, num_options, gpu=-1):
+        super(OptionCriticNet, self).__init__()
+        self.fc_q = layer_init(nn.Linear(body.feature_dim, num_options))
+        self.fc_pi = layer_init(nn.Linear(body.feature_dim, num_options * action_dim))
+        self.fc_beta = layer_init(nn.Linear(body.feature_dim, num_options))
+        self.num_options = num_options
+        self.action_dim = action_dim
+        self.body = body
+        self.set_gpu(gpu)
+
+    def predict(self, x):
+        phi = self.body(self.tensor(x))
+        q = self.fc_q(phi)
+        beta = F.sigmoid(self.fc_beta(phi))
+        pi = self.fc_pi(phi)
+        pi = pi.view(-1, self.num_options, self.action_dim)
+        log_pi = F.log_softmax(pi, dim=-1)
+        return q, beta, log_pi
+
 class GaussianActorNet(nn.Module, BaseNet):
     def __init__(self, action_dim, body, gpu=-1):
         super(GaussianActorNet, self).__init__()

deep_rl/network/network_utils.py

@@ -7,6 +7,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import numpy as np
 
 class BaseNet:
     def set_gpu(self, gpu):

deep_rl/utils/config.py

@@ -60,6 +60,7 @@ class Config:
         self.test_interval = 0
         self.test_repetitions = 10
         self.evaluation_env = None
+        self.termination_regularizer = 0
 
     def add_argument(self, *args, **kwargs):
         self.parser.add_argument(*args, **kwargs)

deep_rl/utils/plot.py

@@ -1,7 +1,7 @@
 # Adapted from  https://github.com/openai/baselines/blob/master/baselines/results_plotter.py
 
+from ..component.bench import load_monitor_log
 import numpy as np
-from ..component import *
 import os
 import re
 
@@ -44,7 +44,7 @@ class Plotter:
     def load_results(self, dirs, max_timesteps=1e8, x_axis=X_TIMESTEPS, episode_window=100):
         tslist = []
         for dir in dirs:
-            ts = component.load_monitor_log(dir)
+            ts = load_monitor_log(dir)
             ts = ts[ts.l.cumsum() <= max_timesteps]
             tslist.append(ts)
         xy_list = [self.ts2xy(ts, x_axis) for ts in tslist]