MuhammadOsamaNusrat/rl-gridworld

1import matplotlib.pyplot as plt
2import numpy as np
3import imageio
4import torch
5from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
6
7def draw_grid_frame(path, grid_size=(4, 4), holes=[(1, 1), (1, 3), (3, 0)], goal=(3, 3), episode=None):
8    fig, ax = plt.subplots(figsize=(4, 4))
9    ax.set_xticks(np.arange(grid_size[1]+1)-0.5, minor=True)
10    ax.set_yticks(np.arange(grid_size[0]+1)-0.5, minor=True)
11    ax.grid(which="minor", color="black", linestyle='-', linewidth=1)
12    ax.tick_params(which="minor", bottom=False, left=False)
13    ax.set_xticks([])
14    ax.set_yticks([])
15
16    for i in range(grid_size[0]):
17        for j in range(grid_size[1]):
18            label = chr(65 + i * grid_size[1] + j)
19            color = "white"
20            if (i, j) in holes:
21                color = "#ffcccc"
22            elif (i, j) == goal:
23                color = "#ccffcc"
24            ax.add_patch(plt.Rectangle((j - 0.5, i - 0.5), 1, 1, facecolor=color))
25            ax.text(j, i, label, ha='center', va='center', fontsize=12)
26
27    for i, (x, y) in enumerate(path):
28        ax.add_patch(plt.Circle((y, x), 0.3, color='blue', alpha=0.3 + 0.5 * (i/len(path))))
29
30    if path:
31        x, y = path[-1]
32        ax.add_patch(plt.Circle((y, x), 0.3, color='red'))
33
34    ax.set_xlim(-0.5, grid_size[1]-0.5)
35    ax.set_ylim(-0.5, grid_size[0]-0.5)
36    ax.invert_yaxis()
37
38    if episode is not None:
39        ax.set_title(f"Episode {episode}", fontsize=14)
40
41    return fig
42
43def save_agent_walk_gif(trajectory, filename="agent_walk.gif", episode=None, loop=True):
44    frames = []
45    for i in range(1, len(trajectory) + 1):
46        fig = draw_grid_frame(trajectory[:i], episode=episode)
47        canvas = FigureCanvas(fig)
48        canvas.draw()
49        image = np.frombuffer(canvas.buffer_rgba(), dtype='uint8')
50        image = image.reshape(fig.canvas.get_width_height()[::-1] + (4,))
51        frames.append(image)
52        plt.close(fig)
53    
54    imageio.mimsave(filename, frames, duration=0.5, loop=0 if loop else 1)
55
56
57
58
59def plot_policy(Q, grid_size=4):
60    fig, ax = plt.subplots(figsize=(6, 6))
61    arrows = {'up': '↑', 'down': '↓', 'left': '←', 'right': '→'}
62    labels = np.array([chr(65 + i * grid_size + j) for i in range(grid_size) for j in range(grid_size)])
63    labels = labels.reshape((grid_size, grid_size))
64    holes = [(1, 1), (1, 3), (3, 0)]
65    goal = (3, 3)
66    start = (0, 0)
67
68    cell_text = []
69    cell_colors = []
70
71    for i in range(grid_size):
72        row_text = []
73        row_color = []
74        for j in range(grid_size):
75            state = (i, j)
76            label = labels[i, j]
77
78            if state in holes:
79                row_text.append(f"❌\n{label}")
80                row_color.append("#ffcccc")
81            elif state == goal:
82                row_text.append(f"✅\n{label}")
83                row_color.append("#ccffcc")
84            elif state == start:
85                row_text.append(f"🟦\n{label}")
86                row_color.append("#cce5ff")
87            elif state in Q:
88                best_action = max(Q[state], key=Q[state].get)
89                row_text.append(f"{arrows[best_action]}\n{label}")
90                row_color.append("white")
91            else:
92                row_text.append(label)
93                row_color.append("white")
94
95        cell_text.append(row_text)
96        cell_colors.append(row_color)
97
98    table = ax.table(cellText=cell_text,
99                     cellColours=cell_colors,
100                     loc='center',
101                     cellLoc='center',
102                     colWidths=[0.2]*grid_size)
103
104    table.scale(1, 2)
105    ax.axis('off')
106    plt.tight_layout()
107    return fig
108
109def plot_heatmap(Q):
110    fig, ax = plt.subplots()
111    values = np.zeros((4, 4))
112    for i in range(4):
113        for j in range(4):
114            state = (i, j)
115            if state in Q:
116                values[i][j] = max(Q[state].values())
117    c = ax.imshow(values, cmap='coolwarm', interpolation='nearest')
118    plt.colorbar(c)
119    ax.set_title("Q-Value Heatmap (Best Actions)")
120    return fig
121
122def plot_visits(visits):
123    fig, ax = plt.subplots()
124    counts = np.zeros((4, 4))
125    for (i, j), count in visits.items():
126        counts[i][j] = count
127    c = ax.imshow(counts, cmap='YlGn', interpolation='nearest')
128    plt.colorbar(c)
129    ax.set_title("State Visit Frequency")
130    return fig
131
132def plot_dqn_qvalues(agent_model, actions):
133    q_grid = np.zeros((4, 4))
134    with torch.no_grad():
135        for i in range(4):
136            for j in range(4):
137                input_tensor = torch.FloatTensor([i, j])
138                q_vals = agent_model(input_tensor)
139                best_action_val = torch.max(q_vals).item()
140                q_grid[i, j] = best_action_val
141
142    fig, ax = plt.subplots()
143    im = ax.imshow(q_grid, cmap="coolwarm")
144    ax.set_title("DQN Q-Value Heatmap (Best Actions)")
145    plt.colorbar(im)
146    return fig
147