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

network/shallow_network.py

@@ -55,3 +55,21 @@ class ActorCriticFCNet(nn.Module, ActorCriticNet):
         x = F.relu(self.fc1(x))
         phi = self.fc2(x)
         return phi
+
+class CategoricalFCNet(nn.Module, CategoricalNet):
+    def __init__(self, state_dim, n_actions, n_atoms, gpu=0):
+        super(CategoricalFCNet, self).__init__()
+        self.n_actions = n_actions
+        self.n_atoms = n_atoms
+
+        hidden_size = 64
+        self.fc1 = nn.Linear(state_dim, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, hidden_size)
+        self.fc_categorical = nn.Linear(hidden_size, n_actions * n_atoms)
+        BasicNet.__init__(self, gpu)
+
+    def forward(self, x):
+        x = self.variable(x)
+        phi = F.relu(self.fc1(x))
+        phi = F.relu(self.fc2(phi))
+        return phi