Merge pull request #241 from cpnota/bugfix/atari-n-envs-1

fix n-envs=1 atari bug
cpnota · Apr 19, 2021 · 480071a · 480071a
2 parents ba0b7ec + da2620b
commit 480071a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/all/bodies/atari.py b/all/bodies/atari.py
@@ -19,7 +19,7 @@ def process_state(self, state):
         if 'life_lost' not in state:
             return state
 
-        if len(state) == 1:
+        if len(state.shape) == 0:
             if state['life_lost']:
                 return state.update('mask', 0.)
             return state