import gymnasium as gym import numpy as np import gradio as gr from PIL import Image, ImageDraw import collections import time class GridWorldEnv(gym.Env): def __init__(self, size=8, density=0.2): super().__init__() self.size = size self.density = density self.q_table = np.zeros([size * size, 4]) self.reset_maze() def _pos_to_state(self, pos): return int(pos[0] * self.size + pos[1]) def reset_maze(self): while True: coords = [(r, c) for r in range(self.size) for c in range(self.size)] np.random.shuffle(coords) self.start_pos, self.goal_pos = np.array(coords[0]), np.array(coords[1]) self.obstacles = {coords[i] for i in range(2, len(coords)) if np.random.random() < self.density} if self._is_solvable(): break self.reset_agent() def _is_solvable(self): q = collections.deque([tuple(self.start_pos)]) v = {tuple(self.start_pos)} while q: curr = q.popleft() if curr == tuple(self.goal_pos): return True for d in [(-1,0), (1,0), (0,-1), (0,1)]: nxt = (curr[0]+d[0], curr[1]+d[1]) if 0 <= nxt[0] < self.size and 0 <= nxt[1] < self.size and nxt not in self.obstacles and nxt not in v: v.add(nxt); q.append(nxt) return False def reset_agent(self): self.agent_pos = self.start_pos.copy() self.steps, self.total_reward = 0, 0 return self._pos_to_state(self.agent_pos) def step(self, action): old_dist = np.sum(np.abs(self.agent_pos - self.goal_pos)) move = [(-1, 0), (0, 1), (1, 0), (0, -1)][action] new_pos = self.agent_pos + move reward, term, trunc = -0.1, False, False if (0 <= new_pos[0] < self.size and 0 <= new_pos[1] < self.size and tuple(new_pos) not in self.obstacles): self.agent_pos = new_pos reward += (old_dist - np.sum(np.abs(self.agent_pos - self.goal_pos))) * 2 else: reward -= 1.0 self.steps += 1 if np.array_equal(self.agent_pos, self.goal_pos): reward += 50.0; term = True elif self.steps >= self.size * self.size: trunc = True self.total_reward += reward return self._pos_to_state(self.agent_pos), reward, term, trunc def render(self): cell = 60 img = Image.new("RGB", (self.size*cell, self.size*cell), "white") d = ImageDraw.Draw(img) for r in range(self.size): for c in range(self.size): x, y = c*cell, r*cell color = "#F0F0F0" if (r,c) == tuple(self.start_pos): color = "#D0E0FF" if (r,c) == tuple(self.goal_pos): color = "#D0FFD0" if (r,c) in self.obstacles: color = "#444444" d.rectangle([x, y, x+cell, y+cell], fill=color, outline="#CCC") if (r,c) not in self.obstacles and (r,c) != tuple(self.goal_pos): st = r*self.size + c if np.max(self.q_table[st]) != 0: cx, cy = x+cell//2, y+cell//2 ar = [[(0,-10),(-5,-5),(5,-5)],[(10,0),(5,-5),(5,5)],[(0,10),(-5,5),(5,5)],[(-10,0),(-5,-5),(-5,5)]] d.polygon([(cx+px, cy+py) for px, py in ar[np.argmax(self.q_table[st])]], fill="#06F") ar, ac = self.agent_pos d.ellipse([ac*cell+15, ar*cell+15, (ac+1)*cell-15, (ar+1)*cell-15], fill="#F44") return img env = GridWorldEnv() def get_stats(): return [["Score", f"{env.total_reward:.1f}"], ["Steps", f"{env.steps}"], ["Goal", str(tuple(env.goal_pos))]] def solve_logic(mode, lr, df, eps, episodes, progress=gr.Progress()): if mode == "New Brain": env.q_table.fill(0) for ep in progress.tqdm(range(int(episodes)), desc="Training"): state = env.reset_agent() done = False while not done: action = np.random.randint(4) if np.random.rand() < eps else np.argmax(env.q_table[state]) next_s, reward, term, trunc = env.step(action) env.q_table[state, action] += lr * (reward + df * np.max(env.q_table[next_s]) - env.q_table[state, action]) state, done = next_s, term or trunc if ep % 10 == 0: yield env.render(), get_stats() yield env.render(), get_stats() def run_best_path(): state = env.reset_agent() done = False while not done: action = np.argmax(env.q_table[state]) state, _, term, trunc = env.step(action) done = term or trunc yield env.render(), get_stats() time.sleep(0.1) with gr.Blocks(css=".gradio-container {background-color: #f8f9fa}") as demo: gr.Markdown("# 🤖 AI Warehouse Navigator") with gr.Row(): with gr.Column(scale=2): board = gr.Image(label="Warehouse Grid", interactive=False) with gr.Row(): up = gr.Button("⬆️"); left = gr.Button("⬅️"); down = gr.Button("⬇️"); right = gr.Button("➡️") with gr.Column(scale=1): stats_view = gr.Dataframe(headers=["Metric", "Value"], col_count=2, interactive=False) with gr.Accordion("⚙️ Hyperparameters", open=False): lr_sl = gr.Slider(0.01, 1.0, 0.1, label="Alpha (Learning Rate)") df_sl = gr.Slider(0.1, 1.0, 0.9, label="Gamma (Discount Factor)") eps_sl = gr.Slider(0.0, 1.0, 0.1, label="Epsilon (Exploration)") ep_num = gr.Number(100, label="Episodes") mode_st = gr.State("New Brain") solve_btn = gr.Button("🧠 Train New Brain", variant="primary") cont_btn = gr.Button("📈 Continue Training") best_btn = gr.Button("🎯 Run Best Path", variant="secondary") reset_btn = gr.Button("🔄 New Maze") # Manual Actions up.click(lambda: (env.step(0), env.render(), get_stats())[1:], None, [board, stats_view]) right.click(lambda: (env.step(1), env.render(), get_stats())[1:], None, [board, stats_view]) down.click(lambda: (env.step(2), env.render(), get_stats())[1:], None, [board, stats_view]) left.click(lambda: (env.step(3), env.render(), get_stats())[1:], None, [board, stats_view]) # RL Actions solve_btn.click(fn=solve_logic, inputs=[gr.State("New Brain"), lr_sl, df_sl, eps_sl, ep_num], outputs=[board, stats_view]) cont_btn.click(fn=solve_logic, inputs=[gr.State("Keep Brain"), lr_sl, df_sl, eps_sl, ep_num], outputs=[board, stats_view]) best_btn.click(fn=run_best_path, outputs=[board, stats_view]) reset_btn.click(lambda: (env.reset_maze(), env.render(), get_stats()), None, [board, board, stats_view]) demo.load(lambda: (env.render(), get_stats()), None, [board, stats_view]) if __name__ == "__main__": demo.queue().launch() # demo.queue() is mandatory for generators