From d3156ecb069d9f8a1bc8abcca9e631347cafbd40 Mon Sep 17 00:00:00 2001
From: NM512 <morihira3513@gmailcom>
Date: Wed, 17 May 2023 22:18:10 +0900
Subject: [PATCH] bug fix for atari env

---
 envs/atari.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/envs/atari.py b/envs/atari.py
index 025ca2d..4d32c49 100644
--- a/envs/atari.py
+++ b/envs/atari.py
@@ -1,3 +1,4 @@
+import gym
 import numpy as np
 
 
@@ -64,6 +65,16 @@ class Atari:
         self._done = True
         self._step = 0
 
+    @property
+    def observation_space(self):
+        img_shape = self._size + ((1,) if self._gray else (3,))
+        print(self._env.observation_space)
+        return gym.spaces.Dict(
+            {
+                "image": gym.spaces.Box(0, 255, img_shape, np.uint8),
+            }
+        )
+
     @property
     def action_space(self):
         space = self._env.action_space