Update QRDQN and C51

Author: ShangtongZhang

deep_rl/agent/QuantileRegressionDQN_agent.py

@@ -8,34 +8,19 @@
 from ..component import *
 from ..utils import *
 from .BaseAgent import *
+from .DQN_agent import *
 
 
-class QuantileRegressionDQNActor(BaseActor):
+class QuantileRegressionDQNActor(DQNActor):
     def __init__(self, config):
-        BaseActor.__init__(self, config)
-        self.config = config
-        self.start()
-
-    def _transition(self):
-        if self._state is None:
-            self._state = self._task.reset()
-        config = self.config
-        with config.lock:
-            q_values = self._network(config.state_normalizer(self._state)).mean(-1)
-        q_values = to_np(q_values).flatten()
-        if self._total_steps < config.exploration_steps \
-                or np.random.rand() < config.random_action_prob():
-            action = np.random.randint(0, len(q_values))
-        else:
-            action = np.argmax(q_values)
-        next_state, reward, done, info = self._task.step([action])
-        entry = [self._state[0], action, reward[0], next_state[0], int(done[0]), info]
-        self._total_steps += 1
-        self._state = next_state
-        return entry
-
-
-class QuantileRegressionDQNAgent(BaseAgent):
+        super().__init__(config)
+
+    def compute_q(self, prediction):
+        q_values = prediction['quantile'].mean(-1)
+        return to_np(q_values)
+
+
+class QuantileRegressionDQNAgent(DQNAgent):
     def __init__(self, config):
         BaseAgent.__init__(self, config)
         self.config = config
@@ -53,63 +38,40 @@ def __init__(self, config):
         self.actor.set_network(self.network)
 
         self.total_steps = 0
-        self.batch_indices = range_tensor(self.replay.batch_size)
+        self.batch_indices = range_tensor(config.batch_size)
 
         self.quantile_weight = 1.0 / self.config.num_quantiles
         self.cumulative_density = tensor(
             (2 * np.arange(self.config.num_quantiles) + 1) / (2.0 * self.config.num_quantiles)).view(1, -1)
 
-    def close(self):
-        close_obj(self.replay)
-        close_obj(self.actor)
-
     def eval_step(self, state):
         self.config.state_normalizer.set_read_only()
         state = self.config.state_normalizer(state)
-        q = self.network(state).mean(-1)
+        q = self.network(state)['quantile'].mean(-1)
         action = np.argmax(to_np(q).flatten())
         self.config.state_normalizer.unset_read_only()
         return [action]
 
-    def step(self):
-        config = self.config
-        transitions = self.actor.step()
-        experiences = []
-        for state, action, reward, next_state, done, info in transitions:
-            self.record_online_return(info)
-            self.total_steps += 1
-            reward = config.reward_normalizer(reward)
-            experiences.append([state, action, reward, next_state, done])
-        self.replay.feed_batch(experiences)
-
-        if self.total_steps > self.config.exploration_steps:
-            experiences = self.replay.sample()
-            states, actions, rewards, next_states, terminals = experiences
-            states = self.config.state_normalizer(states)
-            next_states = self.config.state_normalizer(next_states)
-
-            quantiles_next = self.target_network(next_states).detach()
-            a_next = torch.argmax(quantiles_next.sum(-1), dim=-1)
-            quantiles_next = quantiles_next[self.batch_indices, a_next, :]
-
-            rewards = tensor(rewards).unsqueeze(-1)
-            terminals = tensor(terminals).unsqueeze(-1)
-            quantiles_next = rewards + self.config.discount * (1 - terminals) * quantiles_next
-
-            quantiles = self.network(states)
-            actions = tensor(actions).long()
-            quantiles = quantiles[self.batch_indices, actions, :]
-
-            quantiles_next = quantiles_next.t().unsqueeze(-1)
-            diff = quantiles_next - quantiles
-            loss = huber(diff) * (self.cumulative_density - (diff.detach() < 0).float()).abs()
-
-            self.optimizer.zero_grad()
-            loss.mean(0).mean(1).sum().backward()
-            nn.utils.clip_grad_norm_(self.network.parameters(), self.config.gradient_clip)
-            with config.lock:
-                self.optimizer.step()
-
-        if self.total_steps / self.config.sgd_update_frequency % \
-                self.config.target_network_update_freq == 0:
-            self.target_network.load_state_dict(self.network.state_dict())
+    def compute_loss(self, transitions):
+        states = self.config.state_normalizer(transitions.state)
+        next_states = self.config.state_normalizer(transitions.next_state)
+
+        quantiles_next = self.target_network(next_states)['quantile'].detach()
+        a_next = torch.argmax(quantiles_next.sum(-1), dim=-1)
+        quantiles_next = quantiles_next[self.batch_indices, a_next, :]
+
+        rewards = tensor(transitions.reward).unsqueeze(-1)
+        masks = tensor(transitions.mask).unsqueeze(-1)
+        quantiles_next = rewards + self.config.discount * masks * quantiles_next
+
+        quantiles = self.network(states)['quantile']
+        actions = tensor(transitions.action).long()
+        quantiles = quantiles[self.batch_indices, actions, :]
+
+        quantiles_next = quantiles_next.t().unsqueeze(-1)
+        diff = quantiles_next - quantiles
+        loss = huber(diff) * (self.cumulative_density - (diff.detach() < 0).float()).abs()
+        return loss.sum(-1).mean(1)
+
+    def reduce_loss(self, loss):
+        return loss.mean()

deep_rl/agent/Rainbow_agent.py

@@ -7,4 +7,3 @@
 from .PPO_agent import *
 from .OptionCritic_agent import *
 from .TD3_agent import *
-from .Rainbow_agent import *

deep_rl/agent/__init__.py

@@ -51,7 +51,7 @@ def forward(self, x):
         pre_prob = self.fc_categorical(phi).view((-1, self.action_dim, self.num_atoms))
         prob = F.softmax(pre_prob, dim=-1)
         log_prob = F.log_softmax(pre_prob, dim=-1)
-        return prob, log_prob
+        return dict(prob=prob, log_prob=log_prob)
 
 
 class RainbowNet(nn.Module, BaseNet):
@@ -99,7 +99,7 @@ def forward(self, x):
         phi = self.body(tensor(x))
         quantiles = self.fc_quantiles(phi)
         quantiles = quantiles.view((-1, self.action_dim, self.num_quantiles))
-        return quantiles
+        return dict(quantile=quantiles)
 
 
 class OptionCriticNet(nn.Module, BaseNet):

deep_rl/network/network_heads.py

@@ -65,6 +65,7 @@ def __init__(self):
         self.decaying_lr = False
         self.shared_repr = False
         self.noisy_linear = False
+        self.n_step = 1
 
     @property
     def eval_env(self):

deep_rl/utils/config.py

@@ -108,8 +108,11 @@ def quantile_regression_dqn_feature(**kwargs):
     config.optimizer_fn = lambda params: torch.optim.RMSprop(params, 0.001)
     config.network_fn = lambda: QuantileNet(config.action_dim, config.num_quantiles, FCBody(config.state_dim))
 
-    # config.replay_fn = lambda: Replay(memory_size=int(1e4), batch_size=10)
-    config.replay_fn = lambda: AsyncReplay(memory_size=int(1e4), batch_size=10)
+    config.batch_size = 10
+    replay_kwargs = dict(
+        memory_size=int(1e4),
+        batch_size=config.batch_size)
+    config.replay_fn = lambda: ReplayWrapper(UniformReplay, replay_kwargs, async=True)
 
     config.random_action_prob = LinearSchedule(1.0, 0.1, 1e4)
     config.discount = 0.99
@@ -136,8 +139,13 @@ def quantile_regression_dqn_pixel(**kwargs):
     config.network_fn = lambda: QuantileNet(config.action_dim, config.num_quantiles, NatureConvBody())
     config.random_action_prob = LinearSchedule(1.0, 0.01, 1e6)
 
-    # config.replay_fn = lambda: Replay(memory_size=int(1e6), batch_size=32)
-    config.replay_fn = lambda: AsyncReplay(memory_size=int(1e6), batch_size=32)
+    config.batch_size = 32
+    replay_kwargs = dict(
+        memory_size=int(1e6),
+        batch_size=config.batch_size,
+        history_length=4,
+    )
+    config.replay_fn = lambda: ReplayWrapper(UniformReplay, replay_kwargs, async=True)
 
     config.state_normalizer = ImageNormalizer()
     config.reward_normalizer = SignNormalizer()
@@ -164,8 +172,11 @@ def categorical_dqn_feature(**kwargs):
     config.network_fn = lambda: CategoricalNet(config.action_dim, config.categorical_n_atoms, FCBody(config.state_dim))
     config.random_action_prob = LinearSchedule(1.0, 0.1, 1e4)
 
-    # config.replay_fn = lambda: Replay(memory_size=10000, batch_size=10)
-    config.replay_fn = lambda: AsyncReplay(memory_size=10000, batch_size=10)
+    config.batch_size = 10
+    replay_kwargs = dict(
+        memory_size=int(1e4),
+        batch_size=config.batch_size)
+    config.replay_fn = lambda: ReplayWrapper(UniformReplay, replay_kwargs, async=True)
 
     config.discount = 0.99
     config.target_network_update_freq = 200
@@ -193,8 +204,13 @@ def categorical_dqn_pixel(**kwargs):
     config.network_fn = lambda: CategoricalNet(config.action_dim, config.categorical_n_atoms, NatureConvBody())
     config.random_action_prob = LinearSchedule(1.0, 0.01, 1e6)
 
-    # config.replay_fn = lambda: Replay(memory_size=int(1e6), batch_size=32)
-    config.replay_fn = lambda: AsyncReplay(memory_size=int(1e6), batch_size=32)
+    config.batch_size = 32
+    replay_kwargs = dict(
+        memory_size=int(1e6),
+        batch_size=config.batch_size,
+        history_length=4,
+    )
+    config.replay_fn = lambda: ReplayWrapper(UniformReplay, replay_kwargs, async=True)
 
     config.discount = 0.99
     config.state_normalizer = ImageNormalizer()
@@ -605,6 +621,7 @@ def td3_continuous(**kwargs):
     mkdir('tf_log')
     set_one_thread()
     random_seed()
+    # -1 is CPU, a positive integer is the index of GPU
     select_device(-1)
     # select_device(0)
 
@@ -627,7 +644,7 @@ def td3_continuous(**kwargs):
     game = 'BreakoutNoFrameskip-v4'
     # dqn_pixel(game=game, n_step=1, replay_cls=UniformReplay, async_replay=True)
     # quantile_regression_dqn_pixel(game=game)
-    # categorical_dqn_pixel(game=game)
+    categorical_dqn_pixel(game=game)
     # rainbow_pixel(game=game)
     # a2c_pixel(game=game)
     # n_step_dqn_pixel(game=game)