Update app.py

app.py

@@ -1,10 +1,10 @@
 # app.py
-# Gridworld RL (Q-learning) with real-time animation + NON-FLICKERING learning curve
-# Keeps your original visualization + layout as much as possible.
-#
-# Buttons:
-# - Train (epsilon decay): learns (explore -> exploit)
-# - Play learned policy (epsilon=0): shows shortest safe route deterministically
+# Gridworld RL (Q-learning) with:
+# ✅ Original visualization + layout (as much as possible)
+# ✅ Non-flickering learning curve (always visible)
+# ✅ Option 1: Obstacle density slider (auto-generate more/less blocks)
+# ✅ Train uses epsilon decay (converges); Play shows deterministic route (epsilon=0)
+# ✅ Same obstacle layout is reused for Play (stored in state)
 
 import time
 import numpy as np
@@ -13,6 +13,7 @@ import matplotlib.pyplot as plt
 from matplotlib.patches import Rectangle, FancyBboxPatch
 from io import BytesIO
 from PIL import Image
+from collections import deque
 
 # -----------------------------
 # Gridworld Environment
@@ -25,13 +26,71 @@ ACTION_DELTAS = {
     3: (0, -1),  # left
 }
 
+def _neighbors(r, c, n):
+    if r > 0: yield (r - 1, c)
+    if r < n - 1: yield (r + 1, c)
+    if c > 0: yield (r, c - 1)
+    if c < n - 1: yield (r, c + 1)
+
+def _has_path(size, start, goal, blocked):
+    """BFS to ensure there's at least one safe path from start to goal."""
+    q = deque([start])
+    seen = {start}
+    while q:
+        cur = q.popleft()
+        if cur == goal:
+            return True
+        r, c = cur
+        for nr, nc in _neighbors(r, c, size):
+            nxt = (nr, nc)
+            if nxt in seen or nxt in blocked:
+                continue
+            seen.add(nxt)
+            q.append(nxt)
+    return False
+
+def generate_obstacles(size, start, goal, density, wall_ratio=0.7, max_tries=60, rng=None):
+    """
+    Generate walls + lava with a given density, retrying until there is a safe path.
+    Lava is treated as blocked (terminal negative), so we keep at least one safe route.
+    """
+    rng = rng or np.random.default_rng()
+    density = float(np.clip(density, 0.0, 0.60))
+
+    # If density is too high, repeatedly try; if impossible, gradually reduce density
+    cur_density = density
+    for _ in range(max_tries):
+        walls = set()
+        lava = set()
+
+        for r in range(size):
+            for c in range(size):
+                cell = (r, c)
+                if cell == start or cell == goal:
+                    continue
+                if rng.random() < cur_density:
+                    if rng.random() < wall_ratio:
+                        walls.add(cell)
+                    else:
+                        lava.add(cell)
+
+        blocked = walls | lava
+        if _has_path(size, start, goal, blocked):
+            return walls, lava
+
+        # If no path, soften the environment a bit and try again
+        cur_density = max(0.0, cur_density - 0.02)
+
+    # Fallback: empty obstacles (always solvable)
+    return set(), set()
+
 class GridWorld:
     def __init__(self, size=5, start=(0, 0), goal=None, lava=None, walls=None):
         self.size = int(size)
         self.start = start
         self.goal = goal if goal is not None else (self.size - 1, self.size - 1)
-        self.lava = set(lava or [(1, 3), (2, 3), (3, 1)])
-        self.walls = set(walls or [(1, 1), (2, 1)])
+        self.lava = set(lava or [])
+        self.walls = set(walls or [])
         self.reset()
 
     def reset(self):
@@ -59,7 +118,7 @@ class GridWorld:
         if self.pos in self.lava:
             return self.pos, -10.0, True
 
-        return self.pos, -0.1, False  # step penalty -> shortest path is optimal
+        return self.pos, -0.1, False  # small step penalty -> shortest safe path is optimal
 
 # -----------------------------
 # Q-Learning Agent
@@ -91,7 +150,7 @@ class QAgent:
         self.Q[r1, c1, a] += self.alpha * td_error
 
 # -----------------------------
-# Rendering helpers (original look)
+# Rendering helpers (ORIGINAL look)
 # -----------------------------
 def fig_to_pil(fig):
     buf = BytesIO()
@@ -108,8 +167,10 @@ def draw_grid(env: GridWorld, agent: QAgent = None, show_q=False, episode=None,
     ax.set_aspect("equal")
     ax.axis("off")
 
+    # Background
     ax.add_patch(Rectangle((0, 0), n, n, facecolor="#0b1020"))
 
+    # Draw cells
     for r in range(n):
         for c in range(n):
             x, y = c, n - 1 - r  # invert y so (0,0) is top-left visually
@@ -133,6 +194,7 @@ def draw_grid(env: GridWorld, agent: QAgent = None, show_q=False, episode=None,
                 )
             )
 
+            # overlay Q hint (optional)
             if show_q and agent is not None and (r, c) not in env.walls:
                 best_a = int(np.argmax(agent.Q[r, c]))
                 qv = float(np.max(agent.Q[r, c]))
@@ -141,6 +203,7 @@ def draw_grid(env: GridWorld, agent: QAgent = None, show_q=False, episode=None,
                 ax.text(x + 0.5, y + 0.30, f"{qv:+.2f}", ha="center", va="center",
                         fontsize=9, color="#a9b7e6", alpha=0.55)
 
+    # Icons
     def put_icon(rc, icon, color="#ffffff"):
         r, c = rc
         x, y = c + 0.5, (n - 1 - r) + 0.5
@@ -151,8 +214,11 @@ def draw_grid(env: GridWorld, agent: QAgent = None, show_q=False, episode=None,
         put_icon(rc, "🔥")
     for rc in env.walls:
         put_icon(rc, "🧱")
+
+    # Agent
     put_icon(env.pos, "🤖")
 
+    # Header overlay
     title = "Gridworld RL • Q-learning"
     sub = []
     if episode is not None:
@@ -171,7 +237,7 @@ def draw_grid(env: GridWorld, agent: QAgent = None, show_q=False, episode=None,
 # -----------------------------
 # Learning curve chart (no flicker)
 # -----------------------------
-def moving_average(x, window=20):
+def moving_average(x, window=25):
     if len(x) < 2:
         return np.array(x, dtype=float)
     w = max(2, min(int(window), len(x)))
@@ -181,8 +247,6 @@ def moving_average(x, window=20):
 def draw_learning_curve(returns, successes, window=25):
     fig, ax = plt.subplots(figsize=(5.4, 4.6))
     ax.set_facecolor("#0b1020")
-
-    # dark-friendly axes styling
     for spine in ax.spines.values():
         spine.set_color("#2a355f")
     ax.tick_params(colors="#c9d6ff")
@@ -201,7 +265,8 @@ def draw_learning_curve(returns, successes, window=25):
         ma = moving_average(returns, window=window)
         if len(ma) > 0:
             xs_ma = np.arange(len(returns) - len(ma) + 1, len(returns) + 1)
-            ax.plot(xs_ma, ma, linewidth=2.5, alpha=0.95, label=f"Moving avg ({min(int(window), len(returns))})")
+            ax.plot(xs_ma, ma, linewidth=2.5, alpha=0.95,
+                    label=f"Moving avg ({min(int(window), len(returns))})")
 
     ax2 = ax.twinx()
     ax2.tick_params(colors="#c9d6ff")
@@ -221,15 +286,31 @@ def draw_learning_curve(returns, successes, window=25):
     return fig_to_pil(fig)
 
 # -----------------------------
-# Training + Playback
+# Training + Playback (store env layout so Play matches Train)
 # -----------------------------
-def init_env_agent(size, alpha, gamma):
-    env = GridWorld(size=size, start=(0, 0), goal=(size - 1, size - 1))
+def make_env_and_agent(grid_size, obstacle_density, alpha, gamma):
+    size = int(grid_size)
+    start = (0, 0)
+    goal = (size - 1, size - 1)
+
+    rng = np.random.default_rng()  # new layout each train run
+    walls, lava = generate_obstacles(size, start, goal, density=float(obstacle_density), wall_ratio=0.7, rng=rng)
+
+    env = GridWorld(size=size, start=start, goal=goal, walls=walls, lava=lava)
     agent = QAgent(size=size, alpha=alpha, gamma=gamma)
-    return env, agent
+
+    env_state = {
+        "size": size,
+        "start": start,
+        "goal": goal,
+        "walls": sorted(list(walls)),
+        "lava": sorted(list(lava)),
+    }
+    return env, agent, env_state
 
 def train_stream(
     grid_size,
+    obstacle_density,
     alpha,
     gamma,
     eps_start,
@@ -241,7 +322,7 @@ def train_stream(
     show_q_overlay,
     curve_window,
 ):
-    env, agent = init_env_agent(grid_size, alpha, gamma)
+    env, agent, env_state = make_env_and_agent(grid_size, obstacle_density, alpha, gamma)
     eps = float(eps_start)
 
     returns = []
@@ -250,10 +331,9 @@ def train_stream(
     # initial
     frame = draw_grid(env, agent, show_q=show_q_overlay, episode=0, step_i=0, total_reward=0.0)
     last_curve = draw_learning_curve(returns, successes, window=int(curve_window))
-    status = "Klaar om te trainen. De learning curve blijft nu stabiel in beeld (geen knipperen)."
-    yield frame, last_curve, agent.Q, status
+    status = f"Klaar om te trainen. Obstacle density={float(obstacle_density):.2f}. (Curve knippert niet.)"
+    yield frame, last_curve, agent.Q, env_state, status
 
-    # only redraw chart every N steps (but ALWAYS output the last image)
     CURVE_UPDATE_EVERY_STEPS = 8
 
     for ep in range(1, int(episodes) + 1):
@@ -278,7 +358,7 @@ def train_stream(
 
             frame = draw_grid(env, agent, show_q=show_q_overlay, episode=ep, step_i=t, total_reward=total_r)
             status = f"Train • ep {ep}/{episodes} • step {t}/{max_steps} • return {total_r:+.2f} • eps {eps:.3f}"
-            yield frame, last_curve, agent.Q, status
+            yield frame, last_curve, agent.Q, env_state, status
 
             if speed > 0:
                 time.sleep(float(speed))
@@ -289,33 +369,40 @@ def train_stream(
         returns.append(total_r)
         successes.append(reached_goal_this_ep)
         last_curve = draw_learning_curve(returns, successes, window=int(curve_window))
-        yield frame, last_curve, agent.Q, f"Episode {ep} klaar • return {total_r:+.2f} • success={reached_goal_this_ep} • eps {eps:.3f}"
+        yield frame, last_curve, agent.Q, env_state, f"Episode {ep} klaar • return {total_r:+.2f} • success={reached_goal_this_ep} • eps {eps:.3f}"
 
         eps = max(float(eps_end), eps * float(eps_decay))
 
     frame = draw_grid(env, agent, show_q=show_q_overlay, episode=episodes, step_i=None, total_reward=None)
     last_curve = draw_learning_curve(returns, successes, window=int(curve_window))
     status = "Training klaar ✅ Klik nu op ‘Play learned policy’ om de strakke kortste veilige route te zien (epsilon=0)."
-    yield frame, last_curve, agent.Q, status
+    yield frame, last_curve, agent.Q, env_state, status
 
-def play_stream(q_table, grid_size, max_steps, speed, show_q_overlay):
-    if q_table is None:
-        env = GridWorld(size=grid_size)
-        agent = QAgent(size=grid_size)
+def play_stream(q_table, env_state, speed, show_q_overlay, max_steps):
+    if q_table is None or env_state is None:
+        # show something reasonable
+        env = GridWorld(size=5, start=(0, 0), goal=(4, 4), walls=[], lava=[])
+        agent = QAgent(size=5)
         frame = draw_grid(env, agent, show_q=show_q_overlay, episode=None, step_i=None, total_reward=None)
         curve = draw_learning_curve([], [], window=25)
-        yield frame, curve, "Nog geen Q-table. Klik eerst op Train."
+        yield frame, curve, "Nog geen training gedaan. Klik eerst op Train."
         return
 
-    env = GridWorld(size=grid_size)
-    agent = QAgent(size=grid_size)
+    size = int(env_state["size"])
+    start = tuple(env_state["start"])
+    goal = tuple(env_state["goal"])
+    walls = [tuple(x) for x in env_state["walls"]]
+    lava = [tuple(x) for x in env_state["lava"]]
+
+    env = GridWorld(size=size, start=start, goal=goal, walls=walls, lava=lava)
+    agent = QAgent(size=size)
     agent.Q = np.array(q_table, dtype=np.float32)
 
     s = env.reset()
     total_r = 0.0
 
-    frame = draw_grid(env, agent, show_q=show_q_overlay, episode="PLAY", step_i=0, total_reward=total_r)
     curve = draw_learning_curve([], [], window=25)  # keep curve visible (static) during play
+    frame = draw_grid(env, agent, show_q=show_q_overlay, episode="PLAY", step_i=0, total_reward=total_r)
     yield frame, curve, "Play • epsilon=0.0 (deterministisch) • toont de geleerde route"
 
     for t in range(1, int(max_steps) + 1):
@@ -335,34 +422,41 @@ def play_stream(q_table, grid_size, max_steps, speed, show_q_overlay):
     if env.pos == env.goal:
         end = f"🏁 Goal bereikt! return {total_r:+.2f} (korter pad = minder step-penalty)."
     elif env.pos in env.lava:
-        end = f"🔥 In lava beland. Tip: train langer of eps_start hoger."
+        end = "🔥 In lava beland. Tip: train langer of zet density lager."
     else:
-        end = f"Play klaar. Tip: train langer of max_steps omhoog."
+        end = "Play klaar. Tip: train langer of max_steps omhoog."
 
     frame = draw_grid(env, agent, show_q=show_q_overlay, episode="PLAY", step_i=None, total_reward=total_r)
     yield frame, curve, end
 
 # -----------------------------
-# Gradio UI (original layout)
+# Gradio UI (original layout) + obstacle density slider
 # -----------------------------
 with gr.Blocks(theme=gr.themes.Soft(), title="RL Gridworld (Q-learning)") as demo:
     gr.Markdown(
         """
 # 🤖 Reinforcement Learning in een Gridworld (real-time animatie)
 
+- **Obstacle density**: hoeveel 🧱/🔥 er in het grid staan (meer = moeilijker)
 - **Train**: agent leert (epsilon decays: eerst ontdekken, later benutten)
 - **Play learned policy**: toont wat hij geleerd heeft (**epsilon=0**)
 
-Rechts zie je een **learning curve** (return + moving average + success rate) die nu **niet knippert**.
+Rechts zie je een **learning curve** (return + moving average + success rate) die **niet knippert**.
         """
     )
 
     q_state = gr.State(None)
+    env_state = gr.State(None)
 
     with gr.Row():
         with gr.Column(scale=1):
             grid_size = gr.Slider(4, 10, value=5, step=1, label="Grid size")
 
+            obstacle_density = gr.Slider(
+                0.0, 0.45, value=0.15, step=0.05,
+                label="Obstacle density (meer blokken/gevaar)"
+            )
+
             with gr.Accordion("RL parameters (defaults = goede convergentie)", open=True):
                 alpha = gr.Slider(0.01, 1.0, value=0.45, step=0.01, label="Alpha (learning rate)")
                 gamma = gr.Slider(0.0, 0.999, value=0.97, step=0.001, label="Gamma (discount)")
@@ -373,7 +467,7 @@ Rechts zie je een **learning curve** (return + moving average + success rate) di
                 eps_decay = gr.Slider(0.90, 0.999, value=0.985, step=0.001, label="Epsilon decay per episode")
 
             episodes = gr.Slider(1, 400, value=200, step=1, label="Episodes")
-            max_steps = gr.Slider(5, 200, value=60, step=1, label="Max steps per episode")
+            max_steps_train = gr.Slider(5, 200, value=60, step=1, label="Max steps per episode")
 
             with gr.Accordion("Visuals & snelheid", open=True):
                 speed = gr.Slider(0.0, 0.3, value=0.02, step=0.01, label="Animatie vertraging (sec/frame)")
@@ -384,7 +478,7 @@ Rechts zie je een **learning curve** (return + moving average + success rate) di
                 train_btn = gr.Button("🚀 Train (epsilon decay)", variant="primary")
                 play_btn = gr.Button("▶️ Play learned policy (epsilon=0)")
 
-            status = gr.Textbox(label="Status", value="Klik Train om te zien hoe de agent leert.", interactive=False)
+            status = gr.Textbox(label="Status", value="Kies density en klik Train.", interactive=False)
 
         with gr.Column(scale=1):
             frame_out = gr.Image(label="Live animatie", type="pil", height=520)
@@ -393,17 +487,19 @@ Rechts zie je een **learning curve** (return + moving average + success rate) di
     train_btn.click(
         fn=train_stream,
         inputs=[
-            grid_size, alpha, gamma,
+            grid_size,
+            obstacle_density,
+            alpha, gamma,
             eps_start, eps_end, eps_decay,
-            episodes, max_steps,
+            episodes, max_steps_train,
             speed, show_q_overlay, curve_window
         ],
-        outputs=[frame_out, curve_out, q_state, status],
+        outputs=[frame_out, curve_out, q_state, env_state, status],
     )
 
     play_btn.click(
         fn=play_stream,
-        inputs=[q_state, grid_size, max_steps, speed, show_q_overlay],
+        inputs=[q_state, env_state, speed, show_q_overlay, max_steps_train],
         outputs=[frame_out, curve_out, status],
     )