Merge da7dc01 into 1a322f7

chainer · Nov 9, 2017 · 10f5a0d · 10f5a0d
2 parents 1a322f7 + da7dc01
commit 10f5a0d
Showing 1 changed file with 16 additions and 7 deletions.
diff --git a/chainerrl/policies/gaussian_policy.py b/chainerrl/policies/gaussian_policy.py
@@ -164,14 +164,23 @@ def __init__(self, n_input_channels, action_size, var,
         else:
             self.var = var
         layers = []
-        layers.append(L.Linear(n_input_channels, n_hidden_channels))
-        for _ in range(n_hidden_layers - 1):
+        if n_hidden_layers > 0:
+            # Input to hidden
+            layers.append(L.Linear(n_input_channels, n_hidden_channels))
             layers.append(self.nonlinearity)
-            layers.append(L.Linear(n_hidden_channels, n_hidden_channels))
-        # The last layer is used to compute the mean
-        layers.append(
-            L.Linear(n_hidden_channels, action_size,
-                     initialW=LeCunNormal(mean_wscale)))
+            for _ in range(n_hidden_layers - 1):
+                # Hidden to hidden
+                layers.append(L.Linear(n_hidden_channels, n_hidden_channels))
+                layers.append(self.nonlinearity)
+            # The last layer is used to compute the mean
+            layers.append(
+                L.Linear(n_hidden_channels, action_size,
+                         initialW=LeCunNormal(mean_wscale)))
+        else:
+            # There's only one layer for computing the mean
+            layers.append(
+                L.Linear(n_input_channels, action_size,
+                         initialW=LeCunNormal(mean_wscale)))
 
         if self.bound_mean:
             layers.append(lambda x: bound_by_tanh(