Spaces:
Sleeping
Sleeping
| 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 | |