Multiple animations in python with matplotlib

I would like to use matplotlib's animation capabilities to display and save multiple animations. Previously, I was just using pyplot with a short pause at each step to fake the animation, but I did not find a way to save these "animations" as videos, so I'm switching to using the real animations. Here is a dummy version of my code (which will run) when I started:

from matplotlib import pyplot as plt
import numpy as np

class Hallway:
    def __init__(self):
        self.end_pos = 5
        self.cur_pos = 0

    def setup(self):
        self.cur_pos = 0

    def go(self, decision):
        if decision == 0 and self.cur_pos > 0:
            self.cur_pos -= 1
        elif decision == 1:
            self.cur_pos += 1
        done = self.cur_pos >= self.end_pos
        return done
    
    def draw(self, fig):
        fig.clear()
        ax = fig.gca()
        ax.set(xlim=(-0.5, 5.5), ylim=(-0.5, 0.5))
        ax.scatter(self.cur_pos, 0., s=350)
        plt.draw()
        plt.pause(0.01)

sim = Hallway()
for num_sim in range(5):
    print("running simulation {}".format(num_sim))
    sim.setup()
    sim.draw(plt.gcf())
    while True:
        done = sim.go(np.random.randint(0,2))
        sim.draw(plt.gcf())
        if done:
            break
    # Save animation here

Key things to note in here:

1. The next state of the Hallway generated with go
2. The frames are generated with draw
3. The done condition indicates when the simulation should end
4. Once an animation ends, I want to save it, but we're not done! After saving the animation, I want to launch a new one. This will happen 5 times with the outer loop.

So I changed my code around so that I could use an animation object, and this is what it is now:

from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np

class Hallway:
    def __init__(self):
        self.end_pos = 5
        self.cur_pos = 0

    def setup(self):
        self.cur_pos = 0

    def go(self, decision):
        if decision == 0 and self.cur_pos > 0:
            self.cur_pos -= 1
        elif decision == 1:
            self.cur_pos += 1
        done = self.cur_pos >= self.end_pos
        return done
    
    def draw(self, fig):
        fig.clear()
        ax = fig.gca()
        ax.set(xlim=(-0.5, 5.5), ylim=(-0.5, 0.5))
        ax.scatter(self.cur_pos, 0., s=350)
        plt.draw()
        plt.pause(0.01)

sim = Hallway()
for num_sim in range(5):
    print("running simulation {}".format(num_sim))
    sim.setup()
    all_done = False
    fig = plt.figure()

    def gen_frame_until_done(): # Using a generator to give me frames as long as not all_done
        global all_done
        i = 0
        while not all_done:
            i += 1
            yield i
    
    def animate(i): # Animate function takes the place of the while loop
        sim.draw(fig)
        done = sim.go(np.random.randint(0,2))
        if done:
            global all_done
            all_done = True
            sim.draw(fig)
    
    anim = FuncAnimation(fig, animate, frames=gen_frame_until_done, repeat=False)
    plt.show()
    # anim.save(...)

This will run, but it won't quite give me what I want. It will show only one animation and the terminal will show running simulation 0. When all_done is triggered and the simulation is over, the program will wait for me to exit the plot window. Once I exit, the program will continue to the next simulation and repeat.

I don't like that I have to manually exit the window. I got a little hack semi-working by replacing the blocking plt.show() with

plt.show(block=False)
plt.pause(3)
plt.close()

This will allow the program to continue without having to manually exit the window. However, it will only allow 3 seconds of the animation to display before going on to the next one.

What I want:

• I want to be able to display the simulation until it is over. When it is over, I want the window to automatically close.
• I want the next simulation to run with a new animation window right after the previous one.

Again, I'm using the animation objects because I need to be able to save the animations as videos. But if there's another way to do this, I'm definitely open to it.

from Multiple animations in python with matplotlib