From ef647949861711e947720641ecbd3f6dd8427c06 Mon Sep 17 00:00:00 2001
From: Arpit Garg <iarpitgarg@gmail.com>
Date: Thu, 2 Aug 2018 18:01:48 -0600
Subject: [PATCH] Corner case in A3C added where the game ends before N steps
 are taken

---
 CartPole-A3C.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/CartPole-A3C.py b/CartPole-A3C.py
index 1f3807c..9e90fce 100644
--- a/CartPole-A3C.py
+++ b/CartPole-A3C.py
@@ -194,6 +194,12 @@ def get_sample(memory, n):
 		self.R = ( self.R + r * GAMMA_N ) / GAMMA
 
 		if s_ is None:
+                        if len(self.memory) < N_STEP_RETURN:    # possible edge case - if an episode ends in <N steps
+                            cnt = N_STEP_RETURN - len(self.memory)
+                            while cnt:
+                                self.R /= GAMMA
+                                cnt -= 1
+
 			while len(self.memory) > 0:
 				n = len(self.memory)
 				s, a, r, s_ = get_sample(self.memory, n)
@@ -210,8 +216,6 @@ def get_sample(memory, n):
 
 			self.R = self.R - self.memory[0][2]
 			self.memory.pop(0)	
-	
-	# possible edge case - if an episode ends in <N steps, the computation is incorrect
 		
 #---------
 class Environment(threading.Thread):
@@ -300,4 +304,4 @@ def stop(self):
 	o.join()
 
 print("Training finished")
-env_test.run()
\ No newline at end of file
+env_test.run()