plotme.py

import matplotlib.pyplot as plt


class TrainingPlot:
    def __init__(self):
        plt.ion()  # Turn on interactive mode
        self.fig, self.ax = plt.subplots(figsize=(10, 5))
        self.scores_line, = self.ax.plot([], [], label='Scores', color='tab:blue', linewidth=1.8)
        self.mean_scores_line, = self.ax.plot([], [], label='Mean Scores', color='tab:orange', linewidth=2.2)
        self.ax.set_title('Training...')
        self.ax.set_xlabel('Number of Games')
        self.ax.set_ylabel('Score')
        self.ax.grid(True, alpha=0.3)
        self.ax.legend()
        self.scores = []
        self.mean_scores = []
        self.fig.tight_layout()
        plt.show(block=False)

    def update(self, scores, mean_scores):
        self.scores = list(scores)
        self.mean_scores = list(mean_scores)
        x_data = range(1, len(self.scores) + 1)

        self.scores_line.set_xdata(x_data)
        self.scores_line.set_ydata(self.scores)

        self.mean_scores_line.set_xdata(x_data)
        self.mean_scores_line.set_ydata(self.mean_scores)

        if self.scores:
            upper = max(max(self.scores), max(self.mean_scores), 1)
            self.ax.set_xlim(1, len(self.scores))
            self.ax.set_ylim(0, upper + 1)

        self.fig.canvas.draw_idle()
        self.fig.canvas.flush_events()
        plt.pause(0.01)