diff --git a/mujoco_maze/agent_model.py b/mujoco_maze/agent_model.py index 57c7de5..cefaa49 100644 --- a/mujoco_maze/agent_model.py +++ b/mujoco_maze/agent_model.py @@ -17,6 +17,13 @@ class AgentModel(ABC, MujocoEnv, EzPickle): MujocoEnv.__init__(self, file_path, frame_skip) EzPickle.__init__(self) + def close(self): + if self.viewer is not None and hasattr(self.viewer, "window"): + import glfw + + glfw.destroy_window(self.viewer.window) + super().close() + @abstractmethod def _get_obs(self) -> np.ndarray: """Returns the observation from the model. diff --git a/mujoco_maze/maze_env.py b/mujoco_maze/maze_env.py index d9fe7e8..0e8cbc0 100644 --- a/mujoco_maze/maze_env.py +++ b/mujoco_maze/maze_env.py @@ -453,3 +453,6 @@ class MazeEnv(gym.Env): done = self._task.termination(next_obs) info["position"] = self.wrapped_env.get_xy() return next_obs, inner_reward + outer_reward, done, info + + def close(self): + self.wrapped_env.close() diff --git a/tests/test_envs.py b/tests/test_envs.py index 45346c4..812aa29 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -24,7 +24,7 @@ def test_point_maze(maze_id): @pytest.mark.parametrize("v", [0, 1]) def test_maze_args(v): - env = gym.make(f"PointTRoom-v{v}", task_kwargs={"goals": [(-2.0, 4.0)]}) + env = gym.make(f"PointTRoom-v{v}", task_kwargs={"goals": [(-2.0, -4.0)]}) assert env.reset().shape == (7,) s, _, _, _ = env.step(env.action_space.sample()) assert s.shape == (7,)