switching to matplotlib.animation

main.py

@@ -2,28 +2,37 @@ from sim import Environment, CAMSReverseAndSidestepAgent, Direction
 import time
 import matplotlib.pyplot as plt
 import numpy as np
+from matplotlib.animation import FuncAnimation
 
 STEPS = 1000
-FRAME_DELAY = 1/60
 
-
-# Fixing random state for reproducibility
 np.random.seed(12345)
-
+fig, ax = plt.subplots()
 env = Environment((40, 30))
+
 first_agent = CAMSReverseAndSidestepAgent(
     environment=env, 
     position=(20,15), 
     initial_direction=Direction.NORTH
 )
 
-fig, ax = plt.subplots()
-im = ax.imshow(env.render(), aspect="equal", origin="lower")
-fig.show()
+second_agent = CAMSReverseAndSidestepAgent(
+    environment=env, 
+    position=(11,13), 
+    initial_direction=Direction.EAST
+)
 
-for i in range(1, STEPS+1):
-    print(i)
-    plt.pause(FRAME_DELAY)
+im = ax.imshow(env.render(), aspect="equal", origin="lower")
+
+def update(frame_number):
     env.step()
     im.set_data(env.render())
-    fig.canvas.draw()
+    return im,
+
+ani = FuncAnimation(
+    fig=fig, 
+    func=update,
+    frames=STEPS,
+    blit=True)
+
+plt.show()