Categorical DQN cart pole

network/base_network.py

@@ -91,4 +91,13 @@ class DuelingNet(BasicNet):
         q = value.expand_as(advantange) + (advantange - advantange.mean(1).expand_as(advantange))
         if to_numpy:
             return q.cpu().data.numpy()
-        return q
+        return q
+
+class CategoricalNet(BasicNet):
+    def predict(self, x, to_numpy=False):
+        phi = self.forward(x)
+        pre_prob = self.fc_categorical(phi).view((-1, self.n_actions, self.n_atoms))
+        prob = F.softmax(pre_prob, dim=-1)
+        if to_numpy:
+            return pre_prob.cpu().data.numpy()
+        return prob