mem maze env ok 1.2

dreamer.py

@@ -210,17 +210,17 @@ def make_env(config, logger, mode, train_eps, eval_eps):
             task, mode if "train" in mode else "test", config.action_repeat
         )
         env = wrappers.OneHotAction(env)
-    elif suite == "mazegym":
+    elif suite == "MemoryMaze":
         import gym
-        if task == '9':
+        if task == '9x9':
             env = gym.make('memory_maze:MemoryMaze-9x9-v0')
-        elif task == '15':
+        elif task == '15x15':
             env = gym.make('memory_maze:MemoryMaze-15x15-v0')
         else:
             raise NotImplementedError(suite)
-        from envs.memmazeEnv import MZGymWrapper
-        env = MZGymWrapper(env)
-        env = wrappers.OneHotAction2(env)
+        from envs.memorymaze import MemoryMaze
+        env = MemoryMaze(env)
+        env = wrappers.OneHotAction(env)
     else:
         raise NotImplementedError(suite)
     env = wrappers.TimeLimit(env, config.time_limit)