import matplotlib.pyplot as plt import matplotlib as mpl mpl.use('Agg') import seaborn as sns import numpy as np sns.set_style('darkgrid') import imageio def make_rl_gif(action_logits, action_probs, actions, values, rewards, pre_activations, post_activations, inputs, filename): n_steps = len(pre_activations) pre_activations = pre_activations[:,0,:] post_activations = post_activations[:,0,:] if action_logits.shape[1] == 5: class_labels = ['W', 'U', 'D', 'L', 'R'] elif action_logits.shape[1] == 2: class_labels = ['L', 'R'] else: class_labels = [str(i) for i in range(action_logits.shape[1])] max_target = len(class_labels) figscale = 0.28 frames = [] n_neurons_to_visualise = 15 # Create mosaic layout mosaic = [['img_data', 'img_data', 'img_data', 'img_data', 'action_logits', 'action_logits', 'action_log_probs', 'action_log_probs'] for _ in range(2)] + \ [['img_data', 'img_data', 'img_data', 'img_data', 'action_logits', 'action_logits', 'action_log_probs', 'action_log_probs'] for _ in range(2)] + \ [['value', 'value', 'value', 'value', 'value', 'value', 'value', 'value']] + \ [['reward', 'reward', 'reward', 'reward', 'reward', 'reward', 'reward', 'reward']] + \ [[f'trace_{ti}', f'trace_{ti}', f'trace_{ti}', f'trace_{ti}', f'trace_{ti}', f'trace_{ti}', f'trace_{ti}', f'trace_{ti}'] for ti in range(n_neurons_to_visualise)] # Main plotting loop for stepi in range(n_steps): fig_gif, axes_gif = plt.subplot_mosaic(mosaic=mosaic, figsize=(31*figscale*8/4, 76*figscale)) # Plot action logits these_action_logits = np.array(action_logits)[:, :max_target] colors = ['black' if i == actions[stepi] else ('b' if e >= 0 else 'r') for i, e in enumerate(these_action_logits[stepi])] sort_idxs = np.arange(len(these_action_logits[stepi])) bars = axes_gif['action_logits'].bar(np.arange(len(these_action_logits[stepi][sort_idxs])), these_action_logits[stepi][sort_idxs], color=np.array(colors)[sort_idxs],width=0.9, alpha=0.5) axes_gif['action_logits'].axis('off') for bar, label in zip(bars, class_labels): x = bar.get_x() + bar.get_width() / 2 axes_gif['action_logits'].annotate(label, xy=(x, 0), xytext=(1, 0), textcoords="offset points", ha='center', va='bottom', rotation=90) axes_gif['action_logits'].set_ylim([np.min(these_action_logits), np.max(these_action_logits)]) # Plot action probs these_action_log_probs = np.array(action_probs)[:, :max_target] colors = ['black' if i == actions[stepi] else ('b' if e >= 0 else 'r') for i, e in enumerate(these_action_log_probs[stepi])] sort_idxs = np.arange(len(these_action_log_probs[stepi])) bars = axes_gif['action_log_probs'].bar(np.arange(len(these_action_log_probs[stepi][sort_idxs])), these_action_log_probs[stepi][sort_idxs], color=np.array(colors)[sort_idxs],width=0.9, alpha=0.5) axes_gif['action_log_probs'].axis('off') for bar, label in zip(bars, class_labels): x = bar.get_x() + bar.get_width() / 2 axes_gif['action_log_probs'].annotate(label, xy=(x, 0), xytext=(1, 0), textcoords="offset points", ha='center', va='bottom', rotation=90) axes_gif['action_log_probs'].set_ylim([0,1]) # Plot value trace ax_value = axes_gif['value'] ax_value.plot(np.arange(n_steps), values, 'b-', linewidth=2) ax_value.axvline(x=stepi, color='k', linewidth=2, alpha=0.3) ax_value.set_xticklabels([]) ax_value.set_yticklabels([]) ax_value.grid(False) ax_value.set_xlim([0, n_steps-1]) # Plot reward trace ax_reward = axes_gif['reward'] ax_reward.plot(np.arange(n_steps), rewards, 'g-', linewidth=2) ax_reward.axvline(x=stepi, color='k', linewidth=2, alpha=0.3) ax_reward.set_xticklabels([]) ax_reward.set_yticklabels([]) ax_reward.grid(False) ax_reward.set_xlim([0, n_steps-1]) # Plot neuron traces for neuroni in range(n_neurons_to_visualise): ax = axes_gif[f'trace_{neuroni}'] pre_activation = pre_activations[:, neuroni] post_activation = post_activations[:, neuroni] ax_pre = ax.twinx() pre_min, pre_max = np.min(pre_activation), np.max(pre_activation) post_min, post_max = np.min(post_activation), np.max(post_activation) ax_pre.plot(np.arange(n_steps), pre_activation, color='grey', linestyle='--', linewidth=1, alpha=0.4, label='Pre-activation') color = 'blue' if neuroni % 2 else 'red' ax.plot(np.arange(n_steps), post_activation, color=color, linestyle='-', linewidth=2, alpha=1.0, label='Post-activation') ax.set_xlim([0, n_steps-1]) ax_pre.set_xlim([0, n_steps-1]) ax.set_ylim([post_min, post_max]) ax_pre.set_ylim([pre_min, pre_max]) ax.axvline(x=stepi, color='black', linewidth=1, alpha=0.5) ax.set_xticklabels([]) ax.set_yticklabels([]) ax.grid(False) ax_pre.set_xticklabels([]) ax_pre.set_yticklabels([]) ax_pre.grid(False) ax.set_xlim([0, n_steps-1]) ax.set_xticklabels([]) ax.grid(False) # Show input image this_image = inputs[stepi] axes_gif['img_data'].imshow(this_image, cmap='binary', vmin=0, vmax=1) axes_gif['img_data'].grid(False) axes_gif['img_data'].set_xticks([]) axes_gif['img_data'].set_yticks([]) # Save frames fig_gif.tight_layout(pad=0.1) if stepi == 0: fig_gif.savefig(filename.split('.gif')[0]+'_frame0.png', dpi=100) if stepi == 1: fig_gif.savefig(filename.split('.gif')[0]+'_frame1.png', dpi=100) if stepi == n_steps-1: fig_gif.savefig(filename.split('.gif')[0]+'_frame-1.png', dpi=100) # Convert to frame canvas = fig_gif.canvas canvas.draw() image_numpy = np.frombuffer(canvas.buffer_rgba(), dtype='uint8') image_numpy = image_numpy.reshape(*reversed(canvas.get_width_height()), 4)[:,:,:3] frames.append(image_numpy) plt.close(fig_gif) imageio.mimsave(filename, frames, fps=15, loop=100)