From 3b60dd065cdf5536e471bb8e27b9576591d4dca9 Mon Sep 17 00:00:00 2001
From: Pranjal Tandon <pranjal.tandon@gmail.com>
Date: Thu, 4 Oct 2018 11:23:30 +0530
Subject: [PATCH] Update sac.py

---
 sac.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/sac.py b/sac.py
index f86cb35..37e8aa1 100644
--- a/sac.py
+++ b/sac.py
@@ -38,15 +38,19 @@ class SAC(object):
 
         self.soft_q_criterion = nn.MSELoss()
 
-    def select_action(self, state, deterministic=False):
+    def select_action(self, state, eval=False):
         state = torch.FloatTensor(state).unsqueeze(0)
-        if deterministic == False:
-            _, _, x_t, _, _ = self.policy.evaluate(state)
-            action = torch.tanh(x_t)
+        if eval == False:
+            self.policy.train()
+            print("train")
+            _, _, action, _, _ = self.policy.evaluate(state)
         else:
-            _, _, _, x_t, _ = self.policy.evaluate(state)
-            action = torch.tanh(x_t)
-        action = action.detach().numpy()
+            self.policy.eval()
+            print("eval")
+            _, _, _, action, _ = self.policy.evaluate(state)
+
+        action = torch.tanh(action)
+        action = action.detach().cpu().numpy()
         return action[0]