Hotfix/0.6.1 (#172)

cpnota · Sep 30, 2020 · 8085344 · 8085344
1 parent 31e5aa9
commit 8085344
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/all/environments/gym.py b/all/environments/gym.py
@@ -39,11 +39,15 @@ def name(self):
         return self._name
 
     def reset(self):
-        self._state = State.from_gym(self._env.reset(), device=self._device)
+        self._state = State.from_gym(self._env.reset(), dtype=self._env.observation_space.dtype, device=self._device)
         return self._state
 
     def step(self, action):
-        self._state = State.from_gym(self._env.step(self._convert(action)), device=self._device)
+        self._state = State.from_gym(
+            self._env.step(self._convert(action)),
+            dtype=self._env.observation_space.dtype,
+            device=self._device
+        )
         return self._state
 
     def render(self, **kwargs):