Wrapping the BurgerBot

import numpy as np
import gymnasium as gym
from gymnasium import spaces

# ── World constants — identical to the RL course's GridWorld ─────────────────
GRID_ROWS, GRID_COLS = 7, 7
OBSTACLES = {(0,3),(1,1),(1,5),(3,0),(3,3),(4,5),(5,1)}
START, GOAL = (0,0), (6,6)
ACTIONS   = ["forward","turn_left","turn_right","stop"]
HEADINGS  = ["north","east","south","west"]
HEADING_DELTA = {"north":(-1,0),"east":(0,1),"south":(1,0),"west":(0,-1)}
NEAR, MEDIUM = 1, 3

class BurgerBotEnv(gym.Env):
    """
    The BurgerBot 7x7 grid-world as a Gymnasium environment.

    Observation: np.array([row, col, heading_index, sensor_level], dtype=int64)
        heading_index: 0=north, 1=east, 2=south, 3=west
        sensor_level:  0=near (<=1 clear cell ahead),
                       1=medium (2-3 clear cells),
                       2=far (>=4 clear cells)

    Actions: 0=forward, 1=turn_left, 2=turn_right, 3=stop

    Rewards:
        Hit obstacle:  -10.0  (terminated=True)
        Reach goal:    +50.0  (terminated=True)
        Move forward:   +1.0
        Turn or stop:   -0.1
        Time limit:      0.0  (truncated=True)
    """

    metadata = {"render_modes": ["ansi"]}

    def __init__(self, max_steps=150, render_mode=None):
        super().__init__()
        self.max_steps = max_steps
        self.render_mode = render_mode

        # MultiDiscrete([7, 7, 4, 3]) matches (row, col, heading, sensor)
        self.observation_space = spaces.MultiDiscrete([GRID_ROWS, GRID_COLS, 4, 3])
        self.action_space = spaces.Discrete(4)

    # ── Private helpers ───────────────────────────────────────────────────────

    def _sensor(self):
        """Count clear cells directly ahead, capped at 5. Return binned level."""
        dr, dc = HEADING_DELTA[self.heading]
        r, c = self.row + dr, self.col + dc
        n = 0
        while 0 <= r < GRID_ROWS and 0 <= c < GRID_COLS:
            if (r, c) in OBSTACLES:
                break
            n += 1
            if n >= 5:
                break
            r += dr; c += dc
        if n <= NEAR:   return 0   # near
        if n <= MEDIUM: return 1   # medium
        return 2                   # far

    def _obs(self):
        """Return the current observation as a numpy array."""
        return np.array(
            [self.row, self.col, HEADINGS.index(self.heading), self._sensor()],
            dtype=np.int64
        )

    # ── Gymnasium contract ────────────────────────────────────────────────────

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed)
        self.row, self.col = START
        self.heading = "east"
        self.steps = 0
        return self._obs(), {}

    def step(self, action):
        self.steps += 1
        a = ACTIONS[action]
        hit = False
        moved = False

        if a == "forward":
            dr, dc = HEADING_DELTA[self.heading]
            nr, nc = self.row + dr, self.col + dc
            if (0 <= nr < GRID_ROWS and 0 <= nc < GRID_COLS
                    and (nr, nc) not in OBSTACLES):
                self.row, self.col = nr, nc
                moved = True
            else:
                hit = True
        elif a == "turn_left":
            self.heading = HEADINGS[(HEADINGS.index(self.heading) - 1) % 4]
        elif a == "turn_right":
            self.heading = HEADINGS[(HEADINGS.index(self.heading) + 1) % 4]
        # "stop" does nothing

        reached = (self.row, self.col) == GOAL

        if hit:
            reward = -10.0
        elif reached:
            reward = 50.0
        elif moved:
            reward = 1.0
        else:
            reward = -0.1

        terminated = reached or hit
        truncated  = self.steps >= self.max_steps

        return self._obs(), reward, terminated, truncated, {}

    def render(self):
        if self.render_mode != "ansi":
            return
        heading_char = {"north":"^","east":">","south":"v","west":"<"}
        out = ""
        for r in range(GRID_ROWS):
            for c in range(GRID_COLS):
                if (r, c) == (self.row, self.col):
                    out += heading_char[self.heading] + " "
                elif (r, c) in OBSTACLES:
                    out += "O "
                elif (r, c) == GOAL:
                    out += "G "
                else:
                    out += ". "
            out += "\n"
        return out

from gymnasium.utils.env_checker import check_env

env = BurgerBotEnv()
check_env(env)
print("BurgerBotEnv passed check_env!")
env.close()

env = BurgerBotEnv(render_mode="ansi")
obs, info = env.reset(seed=7)
print("Initial state:")
print(env.render())

total_reward = 0
done = False

for step in range(30):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    total_reward += reward
    done = terminated or truncated

    if done:
        print(f"Episode ended at step {step + 1}.")
        print(f"Terminated: {terminated}, Truncated: {truncated}")
        break

print(f"Total reward from random agent: {total_reward:.1f}")
print("\nFinal grid state:")
print(env.render())
env.close()