RLGridWorld.py

"""
RL Gridworld App — Python port of RLGridworldApp.mlapp
Requires: numpy, matplotlib (pip install numpy matplotlib)
Run:      python rl_gridworld_app.py
"""

import tkinter as tk
from tkinter import ttk
from collections import deque

import numpy as np
import matplotlib
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg


# ─────────────────────────────────────────────────────────────────────────────
# Core RL Environment & Agent
# ─────────────────────────────────────────────────────────────────────────────

class GridworldEnv:
    """3×3 or 5×5 gridworld with a fixed reward structure.

    Actions  : 0=North, 1=South, 2=East, 3=West
    Goal     : Agent learns to walk west along the bottom row
               (forward reward +10, backward reward −30, step penalty −1).
    """

    STEP_PENALTY   = -1
    FWD_REWARD     =  10
    BWD_REWARD     = -30

    def __init__(self, n: int = 3):
        self.reset(n)

    # ------------------------------------------------------------------
    def reset(self, n: int):
        self.n = n
        self.qtable = np.zeros((n, n, 4))
        self.R = self._build_reward_matrix(n)
        self.start = [n // 2, n // 2]
        self.state = list(self.start)
        # Crawler joint angles: phi_gw[row, col, 0/1]
        self.phi_gw = self._build_phi_gw(n)
        np.random.seed(11)

    # ------------------------------------------------------------------
    def _build_reward_matrix(self, n):
        """Build NxN² reward matrix; special rewards only in the bottom row."""
        R = self.STEP_PENALTY * np.ones((n * n, n * n))
        # Bottom row: moving west = forward (+10), moving east = backward (−30)
        for col in range(n - 2, -1, -1):
            R[self.s2i(n-1, col+1), self.s2i(n-1, col)] = self.FWD_REWARD
            R[self.s2i(n-1, col),   self.s2i(n-1, col+1)] = self.BWD_REWARD
        return R

    def _build_phi_gw(self, n):
        phi = np.zeros((n, n, 2))
        phi[:, :, 0] = np.linspace(90, 0, n)[:, None]   # broadcast over cols
        phi[:, :, 1] = np.linspace(210, 310, n)[None, :]  # broadcast over rows
        return phi

    # ------------------------------------------------------------------
    def s2i(self, row, col) -> int:
        """(row, col) → flat state index (0-based)."""
        return row * self.n + col

    def move(self, state, action) -> list:
        r, c = state
        n = self.n
        if   action == 0: return [max(r-1, 0),   c]           # North
        elif action == 1: return [min(r+1, n-1),  c]           # South
        elif action == 2: return [r,  min(c+1, n-1)]           # East
        else:             return [r,  max(c-1, 0)]             # West

    def reward(self, state, next_state) -> float:
        return self.R[self.s2i(*state), self.s2i(*next_state)]

    # ------------------------------------------------------------------
    def q_update(self, state, action, next_state, discount):
        """In-place Bellman update (no learning-rate — matches original MATLAB)."""
        r = self.reward(state, next_state)
        self.qtable[state[0], state[1], action] = (
            r + discount * np.max(self.qtable[next_state[0], next_state[1], :])
        )
        return r


# ─────────────────────────────────────────────────────────────────────────────
# GUI Application
# ─────────────────────────────────────────────────────────────────────────────

class RLGridworldApp:

    ACTION_NAMES = ['N', 'S', 'E', 'W']
    ACTION_MARKERS = ['^', 'v', '>', '<']

    # Crawler geometry (matching original constants)
    BASE_L  = 0.3
    W_RAD   = 0.025
    W_OFF   = 0.3/2 - 0.04
    LINK1_L = 0.10
    LINK2_L = 0.10

    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("RL Gridworld — Q-Learning Demo")
        self.root.configure(bg="#ececec")

        self.env   = GridworldEnv(3)

        # Hyper-parameters (mirrored by sliders/spinboxes below)
        self.discount = tk.DoubleVar(value=0.95)
        self.epsilon  = tk.DoubleVar(value=0.10)

        # Runtime state
        self.auto_training = False
        self.simulating    = False
        self._train_job    = None
        self._sim_job      = None
        self.step_count    = 0
        self.cumulative_r  = 0.0
        self.crawler_pos   = 0.0
        self.rocky_ground  = tk.BooleanVar(value=False)
        self.update_plots  = tk.BooleanVar(value=True)
        self.train_n_steps = tk.BooleanVar(value=False)
        self.n_steps_field = tk.IntVar(value=100)
        self._n_steps_start = 0

        # Moving-average window (4·N² matches MATLAB)
        self._mov_win = deque(maxlen=4 * 9)
        self._mov_win.extend([0.0] * (4 * 9))

        # Reward history for the plot
        self._step_hist  = []
        self._avg_hist   = []
        self._mavg_hist  = []

        self._build_layout()
        self.initialize()

    # ═══════════════════════════════════════════════════════════════════
    # UI construction
    # ═══════════════════════════════════════════════════════════════════

    def _build_layout(self):
        left  = tk.Frame(self.root, width=195, bg="#ececec")
        right = tk.Frame(self.root, bg="#ececec")
        left.pack(side=tk.LEFT, fill=tk.Y, padx=(6, 0), pady=6)
        right.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=6, pady=6)
        left.pack_propagate(False)
        self._build_controls(left)
        self._build_plots(right)

    # ── Left panel ──────────────────────────────────────────────────────

    def _section(self, parent, title):
        f = ttk.LabelFrame(parent, text=title, padding=4)
        f.pack(fill=tk.X, pady=(0, 6), padx=2)
        return f

    def _build_controls(self, parent):
        # ── Environment ──────────────────────────────────
        env_f = self._section(parent, "Environment")
        self._size_var = tk.IntVar(value=9)
        tk.Radiobutton(env_f, text="3 × 3  (9 states)",  variable=self._size_var,
                       value=9,  bg="#ececec").pack(anchor=tk.W)
        tk.Radiobutton(env_f, text="5 × 5  (25 states)", variable=self._size_var,
                       value=25, bg="#ececec").pack(anchor=tk.W)
        tk.Checkbutton(env_f, text="Rocky ground", variable=self.rocky_ground,
                       bg="#ececec").pack(anchor=tk.W)
        tk.Button(env_f, text="⟳  Initialize", command=self.initialize,
                  bg="#4caf50", fg="white", font=("Arial", 10, "bold"),
                  relief=tk.FLAT, padx=4).pack(fill=tk.X, pady=(4, 0))

        # ── Q-Learning parameters ────────────────────────
        qp_f = self._section(parent, "Q-Learning Parameters")
        self._spinbox_row(qp_f, "Discount  γ :", self.discount, 0.0, 1.0, 0.05)
        self._spinbox_row(qp_f, "Epsilon   ε :", self.epsilon,  0.0, 1.0, 0.05)

        # ── Training control ─────────────────────────────
        tr_f = self._section(parent, "Training Control")
        tk.Checkbutton(tr_f, text="Update plots", variable=self.update_plots,
                       bg="#ececec").pack(anchor=tk.W)
        tk.Button(tr_f, text="Step once", command=self.train_single_step,
                  bg="#2196f3", fg="white", relief=tk.FLAT, padx=4).pack(fill=tk.X, pady=2)
        self._train_btn_text = tk.StringVar(value="▶  Train")
        tk.Button(tr_f, textvariable=self._train_btn_text, command=self.toggle_training,
                  bg="#ff9800", fg="white", font=("Arial", 10, "bold"),
                  relief=tk.FLAT, padx=4).pack(fill=tk.X, pady=2)
        nrow = tk.Frame(tr_f, bg="#ececec")
        nrow.pack(fill=tk.X)
        tk.Checkbutton(nrow, text="For N steps:", variable=self.train_n_steps,
                       bg="#ececec").pack(side=tk.LEFT)
        tk.Spinbox(nrow, from_=1, to=999999, textvariable=self.n_steps_field,
                   width=7).pack(side=tk.LEFT)

        # ── Simulate ─────────────────────────────────────
        sim_f = self._section(parent, "Simulate  (no Q update)")
        self._sim_btn_text = tk.StringVar(value="▶  Simulate")
        tk.Button(sim_f, textvariable=self._sim_btn_text, command=self.toggle_simulate,
                  bg="#9c27b0", fg="white", relief=tk.FLAT, padx=4).pack(fill=tk.X)

        # ── Manual control ───────────────────────────────
        mc_f = self._section(parent, "Manual Control  (bypasses ε-greedy)")
        bf = tk.Frame(mc_f, bg="#ececec")
        bf.pack()
        bkw = dict(width=4, relief=tk.RAISED, padx=2, pady=2)
        tk.Button(bf, text=" N↑",  command=lambda: self.manual_step(0), **bkw).grid(row=0, column=1, padx=2, pady=2)
        tk.Button(bf, text="W←",  command=lambda: self.manual_step(3), **bkw).grid(row=1, column=0, padx=2, pady=2)
        tk.Button(bf, text=" S↓",  command=lambda: self.manual_step(1), **bkw).grid(row=1, column=1, padx=2, pady=2)
        tk.Button(bf, text="E→",  command=lambda: self.manual_step(2), **bkw).grid(row=1, column=2, padx=2, pady=2)

        # ── Exploration lamps ────────────────────────────
        lamp_f = self._section(parent, "Current decision")
        self._explore_lbl  = self._lamp_row(lamp_f, "Exploration  (random)")
        self._exploit_lbl  = self._lamp_row(lamp_f, "Exploitation (greedy)")

        # ── Progress ─────────────────────────────────────
        prog_f = self._section(parent, "Training Progress")
        self._prog_vars = {}
        for key, label in [("steps",   "Steps"),
                            ("cum_r",   "Cumulative reward"),
                            ("avg_r",   "Avg reward / step"),
                            ("mov_avg", "Moving average")]:
            self._prog_vars[key] = self._kv_row(prog_f, label)

        # ── Q-value calc display ─────────────────────────
        qc_f = self._section(parent, "Q-Update")
        tk.Label(qc_f, text="Q(s,a) = r + γ · max Q(s',:)",
                 font=("Courier", 8), bg="#ececec").pack(anchor=tk.W)
        self._q_calc_var = tk.StringVar(value="—")
        tk.Label(qc_f, textvariable=self._q_calc_var, font=("Courier", 7),
                 bg="#ececec", wraplength=180, justify=tk.LEFT).pack(anchor=tk.W)

    def _spinbox_row(self, parent, label, var, lo, hi, inc):
        row = tk.Frame(parent, bg="#ececec")
        row.pack(fill=tk.X, pady=1)
        tk.Label(row, text=label, bg="#ececec", width=14, anchor=tk.W).pack(side=tk.LEFT)
        sb = tk.Spinbox(row, from_=lo, to=hi, increment=inc,
                        textvariable=var, width=6, format="%.2f")
        sb.pack(side=tk.LEFT)

    def _lamp_row(self, parent, label):
        row = tk.Frame(parent, bg="#ececec")
        row.pack(fill=tk.X)
        lbl = tk.Label(row, text="●", fg="gray", font=("Arial", 16), bg="#ececec")
        lbl.pack(side=tk.LEFT)
        tk.Label(row, text=label, bg="#ececec", font=("Arial", 9)).pack(side=tk.LEFT)
        return lbl

    def _kv_row(self, parent, label):
        row = tk.Frame(parent, bg="#ececec")
        row.pack(fill=tk.X)
        tk.Label(row, text=label + ":", bg="#ececec", font=("Arial", 8),
                 anchor=tk.W, width=18).pack(side=tk.LEFT)
        var = tk.StringVar(value="0")
        tk.Label(row, textvariable=var, bg="#ececec",
                 font=("Arial", 8, "bold")).pack(side=tk.RIGHT)
        return var

    # ── Right panel (plots) ──────────────────────────────────────────────

    def _build_plots(self, parent):
        self.fig = plt.figure(figsize=(13, 8.5), facecolor="#f5f5f5")
        gs = self.fig.add_gridspec(2, 3, hspace=0.45, wspace=0.32,
                                   left=0.05, right=0.97,
                                   top=0.94, bottom=0.07)
        self.ax_gw      = self.fig.add_subplot(gs[0, 0])
        self.ax_q       = self.fig.add_subplot(gs[0, 1])
        self.ax_p       = self.fig.add_subplot(gs[0, 2])
        self.ax_r       = self.fig.add_subplot(gs[1, 0:2])
        self.ax_crawler = self.fig.add_subplot(gs[1, 2])

        for ax in (self.ax_gw, self.ax_q, self.ax_p, self.ax_r, self.ax_crawler):
            ax.set_facecolor("white")

        self.canvas = FigureCanvasTkAgg(self.fig, master=parent)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

    # ═══════════════════════════════════════════════════════════════════
    # Initialise / reset
    # ═══════════════════════════════════════════════════════════════════

    def initialize(self):
        self._cancel_timers()
        self.auto_training = False
        self.simulating    = False
        self._train_btn_text.set("▶  Train")
        self._sim_btn_text.set("▶  Simulate")

        n = 3 if self._size_var.get() == 9 else 5
        self.env.reset(n)
        self.step_count   = 0
        self.cumulative_r = 0.0
        self.crawler_pos  = 0.0

        win = 4 * n * n
        self._mov_win = deque([0.0] * win, maxlen=win)
        self._step_hist.clear()
        self._avg_hist.clear()
        self._mavg_hist.clear()

        for v in self._prog_vars.values():
            v.set("0")
        self._q_calc_var.set("—")
        self._set_lamps(None)
        self._plot_all(self.env.state)

    # ═══════════════════════════════════════════════════════════════════
    # RL logic
    # ═══════════════════════════════════════════════════════════════════

    def _e_greedy(self):
        if np.random.rand() < self.epsilon.get():
            self._set_lamps("explore")
            return np.random.randint(4)
        else:
            self._set_lamps("exploit")
            return int(np.argmax(self.env.qtable[self.env.state[0], self.env.state[1], :]))

    def _do_train_step(self):
        """One Q-learning step; returns (new_state, reward)."""
        env    = self.env
        action = self._e_greedy()
        ns     = env.move(env.state, action)
        reward = env.q_update(env.state, action, ns, self.discount.get())
        self._record_progress(reward, action, ns)
        return ns, reward

    def _record_progress(self, reward, action, ns):
        self._mov_win.append(reward)
        self.cumulative_r += reward
        self.step_count   += 1
        avg  = self.cumulative_r / self.step_count
        mavg = np.mean(self._mov_win)
        self._step_hist.append(self.step_count)
        self._avg_hist.append(avg)
        self._mavg_hist.append(mavg)

        self._prog_vars["steps"].set(str(self.step_count))
        self._prog_vars["cum_r"].set(f"{self.cumulative_r:.1f}")
        self._prog_vars["avg_r"].set(f"{avg:.3f}")
        self._prog_vars["mov_avg"].set(f"{mavg:.3f}")

        anames = self.ACTION_NAMES
        best_q = np.max(self.env.qtable[ns[0], ns[1], :])
        s = self.env.state
        self._q_calc_var.set(
            f"Q([{s[0]+1},{s[1]+1}], {anames[action]}) = "
            f"{reward:.1f} + {self.discount.get():.2f}·{best_q:.3g}"
        )

    # ═══════════════════════════════════════════════════════════════════
    # Callbacks
    # ═══════════════════════════════════════════════════════════════════

    def train_single_step(self):
        ns, reward = self._do_train_step()
        self.env.state = ns
        self._plot_all(ns, reward)

    def manual_step(self, action):
        env    = self.env
        ns     = env.move(env.state, action)
        reward = env.q_update(env.state, action, ns, self.discount.get())
        self._record_progress(reward, action, ns)
        env.state = ns
        self._plot_all(ns, reward)

    # ── Auto-training timer ─────────────────────────────────────────────

    def toggle_training(self):
        if self.auto_training:
            self.auto_training = False
            self._train_btn_text.set("▶  Train")
            self._cancel_train_timer()
        else:
            self.auto_training = True
            self._train_btn_text.set("⏹  Stop")
            self._n_steps_start = self.step_count
            self._schedule_train()

    def _schedule_train(self):
        if self.auto_training:
            self._train_job = self.root.after(10, self._auto_train_tick)

    def _auto_train_tick(self):
        if not self.auto_training:
            return
        ns, reward = self._do_train_step()
        if self.update_plots.get():
            self._plot_all(ns, reward)
        else:
            # still advance the crawler without redrawing
            self.crawler_pos += (
                np.sign(reward - self.env.STEP_PENALTY) / (5 * self.env.n)
            )
        self.env.state = ns

        # N-steps guard
        if (self.train_n_steps.get() and
                self.step_count >= self._n_steps_start + self.n_steps_field.get()):
            self.auto_training = False
            self._train_btn_text.set("▶  Train")
            return

        self._schedule_train()

    # ── Simulate timer ──────────────────────────────────────────────────

    def toggle_simulate(self):
        if self.simulating:
            self.simulating = False
            self._sim_btn_text.set("▶  Simulate")
            self._cancel_sim_timer()
        else:
            self.simulating = True
            self._sim_btn_text.set("⏹  Stop")
            self._schedule_sim()

    def _schedule_sim(self):
        if self.simulating:
            self._sim_job = self.root.after(10, self._sim_tick)

    def _sim_tick(self):
        if not self.simulating:
            return
        env    = self.env
        action = self._e_greedy()
        ns     = env.move(env.state, action)
        reward = env.reward(env.state, ns)
        # No Q update — only visualise
        self._plot_gw(ns)
        self._plot_qtable(ns)
        self._plot_crawler(ns[0], ns[1], reward)
        self.canvas.draw_idle()
        env.state = ns
        self._schedule_sim()

    # ── Helpers ─────────────────────────────────────────────────────────

    def _cancel_timers(self):
        self._cancel_train_timer()
        self._cancel_sim_timer()

    def _cancel_train_timer(self):
        if self._train_job:
            self.root.after_cancel(self._train_job)
            self._train_job = None

    def _cancel_sim_timer(self):
        if self._sim_job:
            self.root.after_cancel(self._sim_job)
            self._sim_job = None

    def _set_lamps(self, mode):
        if mode == "explore":
            self._explore_lbl.config(fg="lime green")
            self._exploit_lbl.config(fg="gray")
        elif mode == "exploit":
            self._explore_lbl.config(fg="gray")
            self._exploit_lbl.config(fg="lime green")
        else:
            self._explore_lbl.config(fg="gray")
            self._exploit_lbl.config(fg="gray")

    # ═══════════════════════════════════════════════════════════════════
    # Plotting
    # ═══════════════════════════════════════════════════════════════════

    def _plot_all(self, state, reward=None):
        env = self.env
        if reward is None:
            # Use step-penalty as a neutral default (e.g., on init)
            reward = env.STEP_PENALTY
        self._plot_gw(state)
        self._plot_qtable(state)
        self._plot_policy()
        self._plot_rewards_chart()
        self._plot_crawler(state[0], state[1], reward)
        self.canvas.draw_idle()

    # ── Grid helper ─────────────────────────────────────────────────────

    def _draw_grid(self, ax, n):
        for i in range(n + 1):
            v = i + 0.5
            ax.axhline(v, color="black", lw=1.8)
            ax.axvline(v, color="black", lw=1.8)

    def _ax_setup(self, ax, n, title):
        ax.set_xlim(0.5, n + 0.5)
        ax.set_ylim(n + 0.5, 0.5)
        ax.set_xticks(range(1, n + 1))
        ax.set_yticks(range(1, n + 1))
        ax.tick_params(labelsize=7)
        ax.set_title(title, fontweight="bold", fontsize=9)

    # ── Gridworld plot ──────────────────────────────────────────────────

    def _plot_gw(self, agent_state):
        env = self.env
        ax  = self.ax_gw
        n   = env.n
        ax.clear()
        ax.imshow(np.zeros((n, n)), cmap="Blues", alpha=0.08,
                  extent=[0.5, n+0.5, n+0.5, 0.5], aspect="auto")
        self._draw_grid(ax, n)
        self._ax_setup(ax, n, "Gridworld")

        off = 0.33
        for r in range(n):
            for c in range(n):
                # State number
                ax.text(c+1, r+1, str(env.s2i(r, c)+1),
                        ha="center", va="center", fontsize=8, color="#333")
                # Rewards on each edge
                for dr, dc, tx, ty in [(-1,  0,  c+1,    r+1-off),   # N
                                        ( 1,  0,  c+1,    r+1+off),   # S
                                        ( 0,  1,  c+1+off, r+1),      # E
                                        ( 0, -1,  c+1-off, r+1)]:     # W
                    nr, nc = r+dr, c+dc
                    if 0 <= nr < n and 0 <= nc < n:
                        rew = env.R[env.s2i(r, c), env.s2i(nr, nc)]
                    else:
                        rew = env.R[env.s2i(r, c), env.s2i(r, c)]
                    ax.text(tx, ty, f"{int(rew)}", ha="center", va="center",
                            fontsize=5.5, color="#b00020", fontweight="bold")

        ax.plot(agent_state[1]+1, agent_state[0]+1, "o",
                markersize=14, markerfacecolor="red",
                markeredgecolor="darkred", markeredgewidth=1.5)

    # ── Q-table plot ────────────────────────────────────────────────────

    def _plot_qtable(self, agent_state):
        env = self.env
        ax  = self.ax_q
        n   = env.n
        ax.clear()
        ax.imshow(np.zeros((n, n)), cmap="Blues", alpha=0.08,
                  extent=[0.5, n+0.5, n+0.5, 0.5], aspect="auto")
        self._draw_grid(ax, n)
        self._ax_setup(ax, n, "Q-Table")

        off = 0.33
        positions = [(0, 0, -off,  "N"),   # N: above centre
                     (0, 0,  off,  "S"),   # S: below centre
                     (off, 0, 0,   "E"),   # E: right of centre
                     (-off, 0, 0,  "W")]   # W: left of centre
        for r in range(n):
            for c in range(n):
                q = env.qtable[r, c, :]
                # N/S (col fixed, row offset)
                ax.text(c+1,      r+1-off, f"{q[0]:.2g}", ha="center", va="center", fontsize=5.5, fontweight="bold")
                ax.text(c+1,      r+1+off, f"{q[1]:.2g}", ha="center", va="center", fontsize=5.5, fontweight="bold")
                ax.text(c+1+off,  r+1,     f"{q[2]:.2g}", ha="center", va="center", fontsize=5.5, fontweight="bold")
                ax.text(c+1-off,  r+1,     f"{q[3]:.2g}", ha="center", va="center", fontsize=5.5, fontweight="bold")

        ax.plot(agent_state[1]+1, agent_state[0]+1, "o",
                markersize=12, markerfacecolor="red", alpha=0.65,
                markeredgecolor="darkred", markeredgewidth=1.5)

    # ── Policy plot ─────────────────────────────────────────────────────

    def _plot_policy(self):
        env = self.env
        ax  = self.ax_p
        n   = env.n
        ax.clear()
        best = np.argmax(env.qtable, axis=2)   # shape (n, n)
        cmap = plt.cm.get_cmap("RdYlGn", 4)
        ax.imshow(best, cmap=cmap, vmin=-0.5, vmax=3.5, alpha=0.55,
                  extent=[0.5, n+0.5, n+0.5, 0.5], aspect="auto")
        self._draw_grid(ax, n)
        self._ax_setup(ax, n, "Policy  (best action per state)")

        for r in range(n):
            for c in range(n):
                a = best[r, c]
                ax.plot(c+1, r+1, self.ACTION_MARKERS[a],
                        markersize=11, color="black",
                        markeredgecolor="white", markeredgewidth=0.8)

    # ── Reward chart ────────────────────────────────────────────────────

    def _plot_rewards_chart(self):
        ax = self.ax_r
        ax.clear()
        ax.set_title("Training Progress", fontweight="bold", fontsize=9)
        ax.set_xlabel("Steps", fontsize=8)
        ax.set_ylabel("Reward", fontsize=8)
        ax.tick_params(labelsize=7)
        if self._step_hist:
            ax.plot(self._step_hist, self._avg_hist,
                    color="#1f77b4", lw=1, label="Avg reward/step", alpha=0.8)
            ax.plot(self._step_hist, self._mavg_hist,
                    color="#ff7f0e", lw=1.5, label="Moving avg", alpha=0.9)
            ax.legend(fontsize=7)
        ax.grid(True, alpha=0.25)

    # ── Crawler plot ────────────────────────────────────────────────────

    def _plot_crawler(self, state_r, state_c, reward):
        env = self.env
        ax  = self.ax_crawler
        ax.clear()

        phi1 = np.deg2rad(env.phi_gw[state_r, state_c, 0])
        phi2 = np.deg2rad(env.phi_gw[state_r, state_c, 1])

        # Advance position
        self.crawler_pos += (
            np.sign(reward - env.STEP_PENALTY) / (5 * env.n)
        )
        pos = self.crawler_pos

        # Forward kinematics
        arm_y = 0.05
        x0 = pos + self.BASE_L / 2
        y0 = self.W_RAD + arm_y
        l1xe = x0  + np.cos(phi1)        * self.LINK1_L
        l1ye = y0  + np.sin(phi1)        * self.LINK1_L
        l2xe = l1xe + np.cos(phi1 + phi2) * self.LINK2_L
        l2ye = l1ye + np.sin(phi1 + phi2) * self.LINK2_L

        # ── Floor ──────────────────────────────────────────
        x_floor = np.arange(round(pos) - 10, round(pos) + 10, 0.1)
        if self.rocky_ground.get():
            ax.plot(x_floor, np.zeros_like(x_floor), "^-",
                    color="gray", markersize=5, lw=0.8)
        else:
            ax.plot([round(pos)-10, round(pos)+10], [0, 0],
                    "-", color="gray", lw=2)

        # ── Base ───────────────────────────────────────────
        bx = [pos - self.BASE_L/2, pos + self.BASE_L/2,
              pos + self.BASE_L/2, pos - self.BASE_L/2]
        by = [self.W_RAD - 0.005, self.W_RAD - 0.005,
              self.W_RAD + 0.050, self.W_RAD + 0.050]
        ax.fill(bx, by, color="#4caf50", zorder=2)

        # ── Wheels ─────────────────────────────────────────
        for wx in [pos - self.W_OFF, pos + self.W_OFF]:
            circ = plt.Circle((wx, self.W_RAD), self.W_RAD,
                               color="black", zorder=3)
            ax.add_patch(circ)

        # ── Arm links ──────────────────────────────────────
        ax.plot([x0, l1xe], [y0, l1ye], "-", color="#cc4400", lw=5, zorder=4)
        ax.plot([l1xe, l2xe], [l1ye, l2ye], "-", color="#e8b800", lw=5, zorder=4)

        # ── Axis limits ────────────────────────────────────
        ax.set_aspect("equal")
        ax.set_xlim(pos - 1.0, pos + 1.0)
        ax.set_ylim(0, 0.30)
        ax.set_title("Crawler Robot", fontweight="bold", fontsize=9)
        ax.tick_params(labelsize=7)
        ax.grid(True, alpha=0.2)


# ─────────────────────────────────────────────────────────────────────────────

def main():
    root = tk.Tk()
    root.geometry("1440x900")
    app = RLGridworldApp(root)
    root.protocol("WM_DELETE_WINDOW", lambda: (app._cancel_timers(), root.destroy()))
    root.mainloop()


if __name__ == "__main__":
    main()