fixes for reshape

microsoft · Aug 1, 2019 · 0e4938e · 0e4938e
1 parent d10fe20
commit 0e4938e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/MalmoEnv/malmoenv/core.py b/MalmoEnv/malmoenv/core.py
@@ -268,11 +268,14 @@ def _peek_obs(self):
 
             obs = np.frombuffer(obs, dtype=np.uint8)
 
-        if obs is None or len(obs) == 0:
+        if obs is None or len(obs) == 0 or obs.size == 0:
             if self.reshape:
                 obs = np.zeros((self.height, self.width, self.depth), dtype=np.uint8)
             else: 
                 obs = np.zeros(self.height * self.width * self.depth, dtype=np.uint8)
+        elif self.reshape:
+            obs = obs.reshape((self.height, self.width, self.depth)).astype(np.uint8)
+
         return obs
 
     def _quit_episode(self):
@@ -327,8 +330,12 @@ def step(self, action):
             if (obs is None or len(obs) == 0) or turn:
                 time.sleep(0.1)
             obs = np.frombuffer(obs, dtype=np.uint8)
-            if self.reshape:
-                obs = obs.reshape((self.height, self.width, self.depth), dtype=np.uint8)
+
+        if self.reshape:
+            if obs.size == 0:
+                obs = np.zeros((self.height, self.width, self.depth), dtype=np.uint8)
+            else:
+                obs = obs.reshape((self.height, self.width, self.depth)).astype(np.uint8)
 
         return obs, reward, self.done, info
 

diff --git a/MalmoEnv/setup.py b/MalmoEnv/setup.py
@@ -24,7 +24,7 @@
 
 setuptools.setup(
     name="malmoenv",
-    version="0.0.4",
+    version="0.0.6",
     author="Andre Kramer",
     author_email="[email protected]",
     description="A gym environemnt for Malmo",