tiny mod

Author: lazyprogrammer

rl2/cartpole/q_learning.py

@@ -73,10 +73,7 @@ def __init__(self, env, feature_transformer):
 
   def predict(self, s):
     X = self.feature_transformer.transform(np.atleast_2d(s))
-    # print("X.shape", X.shape)
-    result = np.array([m.predict(X)[0] for m in self.models])
-    result = np.atleast_2d(result)
-    assert(len(result.shape) == 2)
+    result = np.stack([m.predict(X) for m in self.models]).T
     return result
 
   def update(self, s, a, G):

rl2/cartpole/td_lambda.py

@@ -51,9 +51,7 @@ def reset(self):
 
   def predict(self, s):
     X = self.feature_transformer.transform([s])
-    # assert(len(X.shape) == 2)
-    result = np.array([m.predict(X)[0] for m in self.models])
-    result = np.atleast_2d(result)
+    result = np.stack([m.predict(X) for m in self.models]).T
     return result
 
   def update(self, s, a, G, gamma, lambda_):

rl2/mountaincar/q_learning.py

@@ -75,9 +75,7 @@ def __init__(self, env, feature_transformer, learning_rate):
 
   def predict(self, s):
     X = self.feature_transformer.transform([s])
-    assert(len(X.shape) == 2)
-    result = np.array([m.predict(X)[0] for m in self.models])
-    result = np.atleast_2d(result)
+    result = np.stack([m.predict(X) for m in self.models]).T
     assert(len(result.shape) == 2)
     return result
 

rl2/mountaincar/td_lambda.py

@@ -52,8 +52,7 @@ def __init__(self, env, feature_transformer):
   def predict(self, s):
     X = self.feature_transformer.transform([s])
     assert(len(X.shape) == 2)
-    result = np.array([m.predict(X)[0] for m in self.models])
-    result = np.atleast_2d(result)
+    result = np.stack([m.predict(X) for m in self.models]).T
     assert(len(result.shape) == 2)
     return result