Fixed trainer.py for GPU execution and modified main.py to run in Ubuntu

Author: ziritrion

main.py

@@ -4,8 +4,8 @@
 from trainer import Trainer
 
 from pyvirtualdisplay import Display
-#display = Display(visible=0, size=(1400, 900))
-#display.start()
+display = Display(visible=0, size=(1400, 900))
+display.start()
 
 # if gpu is to be used
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

params/policy-params.dl

@@ -66,7 +66,7 @@ def episode_train(self, iteration):
             policy_loss.append(-log_prob * advantage)
 
         # calculate critic (value) loss using L1 smooth loss
-            value_losses.append(F.smooth_l1_loss(baseline, torch.tensor([G])))
+            value_losses.append(F.smooth_l1_loss(baseline, torch.tensor([G]).to(self.device)))
 
         # Update policy:
         self.optimizer.zero_grad()