diff --git a/sgail-ppo-options-setobs2.py b/sgail-ppo-options-setobs2.py
index db15496..b600059 100644
--- a/sgail-ppo-options-setobs2.py
+++ b/sgail-ppo-options-setobs2.py
@@ -108,9 +108,9 @@ def training_function(config):
         value=value,
         v_opt=v_opt,
         v_iters=config['value']['iterations_per_epoch'],
-        epochs=2, # 200 FIXME
-        rollout_episodes=6, #60, FIXME
-        rollout_steps=6, #60, FIXME
+        epochs=200,
+        rollout_episodes=60,
+        rollout_steps=60,
         gamma=0.99,
         gae_lambda=0.9,
         clip_ratio=config['policy']['clip_ratio'],
@@ -124,7 +124,6 @@ def training_function(config):
     # save model
     torch.save(policy.state_dict(), 'policy_final.pt')
 
-
 analysis = tune.run(
     training_function,
     config={
@@ -135,7 +134,7 @@ analysis = tune.run(
             'clip_ratio': 0.2, #tune.grid_search([0.2]),
             'iterations_per_epoch': 100, #tune.grid_search([100]),
             'hidden_layer_size': tune.grid_search([10, 20, 40]),
-            'n_hidden_layers': 1, #tune.grid_search([1, 2, 3]), #FIXME
+            'n_hidden_layers': tune.grid_search([2, 3, 4]),
             'activation':tune.grid_search([torch.nn.LeakyReLU, torch.nn.Tanh]),
             'option': tune.grid_search(list(range(len(option_list))))
         },
diff --git a/src/core/policy.py b/src/core/policy.py
index 22aeba4..ba4b3f6 100644
--- a/src/core/policy.py
+++ b/src/core/policy.py
@@ -36,15 +36,20 @@ class BasePolicy(nn.Module):
 
 class Policy(BasePolicy):
     
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, hidden_layer_size=50, n_hidden_layers=2, activation=nn.Tanh, **kwargs):
         super().__init__(*args, **kwargs)
-        self.nn = nn.Sequential(
-            nn.LazyLinear(50),
-            nn.Tanh(),
-            nn.LazyLinear(50),
-            nn.Tanh(),
-            nn.LazyLinear(2 * self.action_dim),
-        )
+        layers = sum([[nn.LazyLinear(hidden_layer_size), 
+            activation()] for _ in range(n_hidden_layers)],[])
+        self.nn = nn.Sequential(*layers, nn.LazyLinear(2 *self.action_dim))
+        
+        # old
+        # self.nn = nn.Sequential(
+        #    nn.LazyLinear(50),
+        #    nn.Tanh(),
+        #    nn.LazyLinear(50),
+        #    nn.Tanh(),
+        #    nn.LazyLinear(2 * self.action_dim),
+        #)
 
     def forward(self, states):
         return self.nn(states)
@@ -53,17 +58,10 @@ class DiscretePolicy(BasePolicy):
 
     def __init__(self, *args, hidden_layer_size=50, n_hidden_layers=2, activation=nn.Tanh, **kwargs):
         super().__init__(*args, **kwargs)
-        layers = [nn.LazyLinear(hidden_layer_size), activation()] * n_hidden_layers
+        layers = sum([[nn.LazyLinear(hidden_layer_size), 
+            activation()] for _ in range(n_hidden_layers)],[])
         self.nn = nn.Sequential(*layers, nn.LazyLinear(self.action_dim))
 
-        #self.nn = nn.Sequential(
-        #    nn.LazyLinear(hidden_layer_size),
-        #    nn.Tanh(),
-        #    nn.LazyLinear(hidden_layer_size),
-        #    nn.Tanh(),
-        #    nn.LazyLinear(self.action_dim),
-        #)
-
     def forward(self, states):
         return self.nn(states)