simple_maze_env.py

import gymnasium as gym
from gymnasium import spaces
import numpy as np
import matplotlib.pyplot as plt


class MazeEnv(gym.Env):
    def __init__(self, pause=0.1):
        super(MazeEnv, self).__init__()

        self.pause = pause  # for render
        # Размеры лабиринта (5x5)
        self.maze_size = (5, 5)

        # Определяем перегородки (стены) между клетками
        self.walls = {
            ((1, 0), (1, 1)),
            ((2, 0), (3, 0)),
            ((4, 0), (4, 1)),
            ((1, 1), (1, 2)),
            ((2, 1), (2, 2)),
            ((2, 1), (3, 1)),
            ((3, 1), (3, 2)),
            ((0, 2), (0, 3)),
            ((1, 2), (1, 3)),
            ((2, 2), (2, 3)),
            ((3, 2), (3, 3)),
            ((0, 3), (1, 3)),
            ((2, 3), (2, 4)),
            ((3, 3), (3, 4)),
            ((4, 3), (4, 4)),
        }

        self.start_pos = np.array([0, 0])
        self.goal_pos = np.array([4, 4])

        # Действия: вверх (0), вниз (1), влево (2), вправо (3)
        self.action_space = spaces.Discrete(4)

        # Состояние: координаты агента (y, x)
        self.observation_space = spaces.Box(low=0, high=4, shape=(2,), dtype=np.int32)

        # Инициализация позиции агента
        self.agent_pos = self.start_pos.copy()

        # Создаём фигуру и оси для рендеринга (инициализация один раз)
        self.fig, self.ax = plt.subplots()
        self.render_initialized = False

    def reset(self):
        # Сброс позиции агента на старт
        self.agent_pos = self.start_pos.copy()
        return self.agent_pos

    def step(self, action):
        # Сохраним старую позицию
        old_pos = self.agent_pos.copy()

        # Определяем новую позицию
        if action == 0:  # Вниз - Move south (down)
            new_pos = [min(self.agent_pos[0] + 1, 4), self.agent_pos[1]]
        elif action == 1:  # Вверх - Move north (up)
            new_pos = [max(self.agent_pos[0] - 1, 0), self.agent_pos[1]]
        elif action == 2:  # Вправо - Move east (right)
            new_pos = [self.agent_pos[0], min(self.agent_pos[1] + 1, 4)]
        elif action == 3:  # Влево - Move west (left)
            new_pos = [self.agent_pos[0], max(self.agent_pos[1] - 1, 0)]
        else:
            new_pos = old_pos  # Неверное действие

        # Проверяем наличие стены между клетками
        if not self.is_wall_between(old_pos, new_pos):
            self.agent_pos = np.array(new_pos)

        # Проверка завершения: достигли цели или нет
        done = np.array_equal(self.agent_pos, self.goal_pos)

        # Вознаграждение: 1 за достижение цели, -0.004 в других случаях
        reward = 1 if done else -0.004

        return self.agent_pos, reward, done, {}

    def is_wall_between(self, pos1, pos2):
        """Проверка наличия стены между двумя клетками"""
        return (tuple(pos1), tuple(pos2)) in self.walls or (
            tuple(pos2),
            tuple(pos1),
        ) in self.walls

    def save_settings_axis(self):
        self.xlim = self.ax.get_xlim()
        self.ylim = self.ax.get_ylim()
        self.aspect = self.ax.get_aspect()

    def set_settings_axis(self):
        self.ax.set_xlim(self.xlim)
        self.ax.set_ylim(self.ylim)
        self.ax.set_aspect(self.aspect)

    def render(self, mode="human"):
        # Если рендер не был инициализирован, создаем оси и отображаем лабиринт
        if not self.render_initialized:
            self.ax.set_xlim(0, self.maze_size[1])
            self.ax.set_ylim(0, self.maze_size[0])
            self.ax.set_aspect("equal")
            self.ax.invert_yaxis()  # Инвертируем ось Y, чтобы (0,0) был в верхнем левом углу
            self.save_settings_axis()
            self.render_initialized = True

        # Очищаем старые объекты
        self.ax.clear()
        self.set_settings_axis()

        # Отображаем клетки лабиринта
        for y in range(self.maze_size[0]):
            for x in range(self.maze_size[1]):
                self.ax.add_patch(
                    plt.Rectangle((x, y), 1, 1, fill=None, edgecolor="yellow")
                )

        # Отображаем стены
        for p1, p2 in self.walls:
            y1, x1 = p1
            y2, x2 = p2
            # print('y1, x1', y1, x1, 'y2, x2', y2, x2)
            if y1 == y2:  # Вертикальная стена
                self.ax.plot([x2, x2], [y1, y2 + 1], color="black", linewidth=3)
            else:  # Горизонтальная стена
                self.ax.plot([x1, x2 + 1], [y2, y2], color="black", linewidth=3)

        # Отображаем агента и цель
        self.ax.scatter(
            self.agent_pos[1] + 0.5,
            self.agent_pos[0] + 0.5,
            c="red",
            marker="o",
            s=200,
            label="Agent",
        )
        self.ax.scatter(
            self.goal_pos[1] + 0.5,
            self.goal_pos[0] + 0.5,
            c="green",
            marker="*",
            s=200,
            label="Goal",
        )

        # Обновляем изображение
        plt.draw()
        plt.pause(self.pause)  # Пауза для визуализации шагов

    def close(self):
        # Закрываем окно рендеринга при завершении
        plt.close(self.fig)


def main():
    # Пример использования среды и анимации
    # env = MazeEnv()
    env = MazeEnv(0.01)
    obs = env.reset()
    env.render()

    # Запланируем шаги агента до достижения цели
    # while True:
    for _ in range(200):
        action = env.action_space.sample()  # Выбираем случайное действие
        obs, reward, done, _ = env.step(action)
        env.render()
        if done:
            print("Goal reached!")
            break

    env.close()


if __name__ == "__main__":
    main()