Update app.py

app.py

@@ -1,147 +1,234 @@
-import gymnasium as gym
-from gymnasium import spaces
 import numpy as np
 import gradio as gr
-from PIL import Image, ImageDraw
-import collections
-
+import gymnasium as gym
+from gymnasium import spaces
+from collections import deque
+
+GRID_SIZE = 8
+OBSTACLE_DENSITY = 0.2
+STEP_LIMIT = 100
+
+# =========================
+# Utility: BFS solvability
+# =========================
+def is_solvable(grid, start, goal):
+    q = deque([start])
+    visited = set([start])
+
+    while q:
+        x, y = q.popleft()
+        if (x, y) == goal:
+            return True
+
+        for dx, dy in [(-1,0),(1,0),(0,-1),(0,1)]:
+            nx, ny = x+dx, y+dy
+            if (
+                0 <= nx < GRID_SIZE and
+                0 <= ny < GRID_SIZE and
+                grid[nx, ny] != 1 and
+                (nx, ny) not in visited
+            ):
+                visited.add((nx, ny))
+                q.append((nx, ny))
+    return False
+
+
+# =========================
+# Gym Environment
+# =========================
 class GridWorldEnv(gym.Env):
-    def __init__(self, size=10, obstacle_density=0.2):
+    def __init__(self):
         super().__init__()
-        self.size = size
-        self.obstacle_density = obstacle_density
+
+        self.action_space = spaces.Discrete(4)
         self.observation_space = spaces.Box(low=0, high=1, shape=(4,), dtype=np.float32)
-        self.action_space = spaces.Discrete(4) # 0:UP, 1:RIGHT, 2:DOWN, 3:LEFT
-        self.reset()
 
-    def _is_solvable(self, start, goal, obstacles):
-        queue = collections.deque([tuple(start)])
-        visited = {tuple(start)}
-        while queue:
-            curr = queue.popleft()
-            if curr == tuple(goal): return True
-            for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
-                nxt = (curr[0] + dx, curr[1] + dy)
-                if 0 <= nxt[0] < self.size and 0 <= nxt[1] < self.size:
-                    if nxt not in obstacles and nxt not in visited:
-                        visited.add(nxt)
-                        queue.append(nxt)
-        return False
+        self.reset()
 
     def reset(self, seed=None, options=None):
-        super().reset(seed=seed)
         while True:
-            coords = [(r, c) for r in range(self.size) for c in range(self.size)]
-            np.random.shuffle(coords)
-            self.start_pos = np.array(coords[0])
-            self.goal_pos = np.array(coords[1])
-            self.obstacles = {coords[i] for i in range(2, len(coords)) if np.random.random() < self.obstacle_density}
-            if self._is_solvable(self.start_pos, self.goal_pos, self.obstacles): break
-
-        self.agent_pos = self.start_pos.copy()
-        self.visited_cells = {tuple(self.agent_pos)}
-        self.steps, self.total_reward, self.last_reward = 0, 0, 0
-        self.max_steps = self.size * self.size
+            self.grid = np.zeros((GRID_SIZE, GRID_SIZE), dtype=int)
+
+            # obstacles
+            for i in range(GRID_SIZE):
+                for j in range(GRID_SIZE):
+                    if np.random.rand() < OBSTACLE_DENSITY:
+                        self.grid[i, j] = 1
+
+            # start & goal
+            self.start = (np.random.randint(GRID_SIZE), np.random.randint(GRID_SIZE))
+            self.goal = (np.random.randint(GRID_SIZE), np.random.randint(GRID_SIZE))
+
+            self.grid[self.start] = 0
+            self.grid[self.goal] = 0
+
+            if self.start != self.goal and is_solvable(self.grid, self.start, self.goal):
+                break
+
+        self.agent = self.start
+        self.steps = 0
+        self.total_score = 0
+        self.visited = set([self.agent])
+
         return self._get_obs(), {}
 
     def _get_obs(self):
-        return np.array([self.agent_pos[0]/self.size, self.agent_pos[1]/self.size, 
-                         self.goal_pos[0]/self.size, self.goal_pos[1]/self.size], dtype=np.float32)
+        return np.array([
+            self.agent[0] / GRID_SIZE,
+            self.agent[1] / GRID_SIZE,
+            self.goal[0] / GRID_SIZE,
+            self.goal[1] / GRID_SIZE
+        ], dtype=np.float32)
 
     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, 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
-            new_dist = np.sum(np.abs(self.agent_pos - self.goal_pos))
-            reward += (old_dist - new_dist) * 2 
-            if tuple(self.agent_pos) not in self.visited_cells:
-                reward += 0.5
-                self.visited_cells.add(tuple(self.agent_pos))
+        x, y = self.agent
+        old_dist = abs(x - self.goal[0]) + abs(y - self.goal[1])
+
+        moves = {
+            0: (-1, 0),  # up
+            1: (0, 1),   # right
+            2: (1, 0),   # down
+            3: (0, -1)   # left
+        }
+
+        dx, dy = moves[action]
+        nx, ny = x + dx, y + dy
+
+        reward = 0
+
+        # invalid move
+        if not (0 <= nx < GRID_SIZE and 0 <= ny < GRID_SIZE) or self.grid[nx, ny] == 1:
+            reward -= 2
+            nx, ny = x, y
         else:
-            reward -= 1.0
+            new_dist = abs(nx - self.goal[0]) + abs(ny - self.goal[1])
 
+            if new_dist < old_dist:
+                reward += 1
+            elif new_dist > old_dist:
+                reward -= 1
+
+            if (nx, ny) not in self.visited:
+                reward += 0.3
+
+        self.agent = (nx, ny)
+        self.visited.add(self.agent)
         self.steps += 1
-        if np.array_equal(self.agent_pos, self.goal_pos):
-            reward += 20.0
-            term = True
-        elif self.steps >= self.max_steps:
-            reward -= 5.0
-            trunc = True
-
-        self.last_reward, self.total_reward = reward, self.total_reward + reward
-        return self._get_obs(), reward, term, trunc, {}
-
-    def render_image(self):
-        cell = 40
-        img = Image.new("RGB", (self.size * cell, self.size * cell), "#F0F0F0")
-        draw = ImageDraw.Draw(img)
-        for r in range(self.size):
-            for c in range(self.size):
-                x0, y0, x1, y1 = c*cell, r*cell, (c+1)*cell, (r+1)*cell
-                if (r, c) == tuple(self.start_pos): draw.rectangle([x0, y0, x1, y1], fill="blue")
-                elif (r, c) == tuple(self.goal_pos): draw.rectangle([x0, y0, x1, y1], fill="green")
-                elif (r, c) in self.obstacles: draw.rectangle([x0, y0, x1, y1], fill="#333333")
-                draw.rectangle([x0, y0, x1, y1], outline="#CCCCCC")
-        r, c = self.agent_pos
-        draw.ellipse([c*cell+5, r*cell+5, (c+1)*cell-5, (r+1)*cell-5], fill="red")
-        return img
-
-# --- Gradio UI ---
+
+        done = False
+
+        if self.agent == self.goal:
+            reward += 10
+            done = True
+
+        if self.steps >= STEP_LIMIT:
+            reward -= 5
+            done = True
+
+        self.total_score += reward
+
+        return self._get_obs(), reward, done, False, {}
+
+# =========================
+# Rendering
+# =========================
+def render(env):
+    grid = np.ones((GRID_SIZE, GRID_SIZE, 3))  # white
+
+    for i in range(GRID_SIZE):
+        for j in range(GRID_SIZE):
+            if env.grid[i, j] == 1:
+                grid[i, j] = [0.2, 0.2, 0.2]  # obstacle
+
+    sx, sy = env.start
+    gx, gy = env.goal
+    ax, ay = env.agent
+
+    grid[sx, sy] = [0, 0, 1]  # blue start
+    grid[gx, gy] = [0, 1, 0]  # green goal
+    grid[ax, ay] = [1, 0, 0]  # red agent
+
+    return grid
+
+def scoreboard(env, last_reward):
+    ax, ay = env.agent
+    gx, gy = env.goal
+    dist = abs(ax - gx) + abs(ay - gy)
+
+    return f"""
+    **Score:** {env.total_score:.2f}  
+    **Last Reward:** {last_reward:.2f}  
+    **Steps:** {env.steps}  
+    **Agent:** {env.agent}  
+    **Goal:** {env.goal}  
+    **Manhattan Distance:** {dist}  
+    **Goal Reached:** {"✅" if env.agent == env.goal else "❌"}
+    """
+
+# =========================
+# App State
+# =========================
 env = GridWorldEnv()
 
-def format_stats():
-    dist = np.sum(np.abs(env.agent_pos - env.goal_pos))
-    data = [
-        ["Total Score", f"{env.total_reward:.2f}"],
-        ["Last Reward", f"{env.last_reward:.2f}"],
-        ["Step Counter", f"{env.steps} / {env.max_steps}"],
-        ["Agent Position", str(tuple(env.agent_pos))],
-        ["Goal Position", str(tuple(env.goal_pos))],
-        ["Dist to Goal", str(int(dist))],
-        ["Status", "GOAL REACHED" if np.array_equal(env.agent_pos, env.goal_pos) else "In Progress"]
-    ]
-    return env.render_image(), data
-
-def play(action):
-    _, _, term, trunc, _ = env.step(action)
-    img, data = format_stats()
-    msg = "🎉 Success!" if term else ("⏰ Timeout!" if trunc else "Navigate to the green cell.")
-    return img, data, msg
-
-with gr.Blocks(css=".gradio-container {background-color: #f8f9fa}") as demo:
-    gr.Markdown("## 📦 Warehouse GridWorld\n**Goal:** Move the red agent to the green exit. Avoid dark obstacles.")
-    
+# =========================
+# Controls
+# =========================
+def step_env(action):
+    obs, reward, done, _, _ = env.step(action)
+    img = render(env)
+    board = scoreboard(env, reward)
+    return img, board
+
+def reset_env():
+    env.reset()
+    return render(env), scoreboard(env, 0)
+
+# =========================
+# Gradio UI
+# =========================
+with gr.Blocks() as demo:
+    gr.Markdown("# 🏭 GridWorld Warehouse Navigation")
+
     with gr.Row():
-        with gr.Column(scale=2):
-            board = gr.Image(label="Warehouse Map", interactive=False)
-            msg_box = gr.Markdown("Click buttons to move.")
-        
-        with gr.Column(scale=1):
-            gr.Markdown("### 📊 Live Scoreboard")
-            stats_table = gr.Dataframe(headers=["Metric", "Value"], col_count=2, interactive=False)
-            
-            with gr.Group():
-                with gr.Row(): gr.Button("").click(None) # Spacer
-                with gr.Row(): up = gr.Button("⬆️ UP")
-                with gr.Row():
-                    left = gr.Button("⬅️ LEFT")
-                    down = gr.Button("⬇️ DOWN")
-                    right = gr.Button("➡️ RIGHT")
-            
-            reset = gr.Button("🔄 Reset Random Maze", variant="stop")
-
-    # Bindings
-    demo.load(format_stats, None, [board, stats_table])
-    up.click(lambda: play(0), None, [board, stats_table, msg_box])
-    right.click(lambda: play(1), None, [board, stats_table, msg_box])
-    down.click(lambda: play(2), None, [board, stats_table, msg_box])
-    left.click(lambda: play(3), None, [board, stats_table, msg_box])
-    reset.click(lambda: (env.reset(), format_stats())[1], None, [board, stats_table])
-
-if __name__ == "__main__":
-    demo.launch()
+        image = gr.Image(value=render(env), interactive=False)
+        board = gr.Markdown(scoreboard(env, 0))
 
+    with gr.Row():
+        up = gr.Button("⬆️")
+    with gr.Row():
+        left = gr.Button("⬅️")
+        down = gr.Button("⬇️")
+        right = gr.Button("➡️")
+
+    reset = gr.Button("🔄 Reset")
+
+    up.click(lambda: step_env(0), outputs=[image, board])
+    right.click(lambda: step_env(1), outputs=[image, board])
+    down.click(lambda: step_env(2), outputs=[image, board])
+    left.click(lambda: step_env(3), outputs=[image, board])
+    reset.click(reset_env, outputs=[image, board])
+
+    # Keyboard support
+    demo.load(
+        None,
+        None,
+        None,
+        _js="""
+        () => {
+            document.addEventListener("keydown", (e) => {
+                const map = {
+                    ArrowUp: 0,
+                    ArrowRight: 1,
+                    ArrowDown: 2,
+                    ArrowLeft: 3
+                };
+                if (map[e.key] !== undefined) {
+                    const buttons = document.querySelectorAll("button");
+                    buttons[map[e.key]].click();
+                }
+            });
+        }
+        """
+    )
+
+demo.launch()