From 2c94106e3849a4515c9db481d56502318368471d Mon Sep 17 00:00:00 2001
From: jay23606 <4602056+jay23606@users.noreply.github.com>
Date: Mon, 26 Oct 2026 15:45:52 -0400
Subject: [PATCH] Show the active algorithm in timelapse

---
 qvis/rl/visualization.py                      | 6 +++++-
 tests/test_rl_algorithms_and_visualization.py | 3 ++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/qvis/rl/visualization.py b/qvis/rl/visualization.py
index 690cba6..b9420df 100644
--- a/qvis/rl/visualization.py
+++ b/qvis/rl/visualization.py
@@ -145,7 +145,11 @@ def timelapse(
             gridspec_kw={"width_ratios": [1, 1.4]},
         )
 
-        fig.suptitle("Q-Learning Training Progress")
+        algorithm = type(self.agent).__name__
+        algorithm = {"QLearning": "Q-Learning", "Sarsa": "SARSA"}.get(
+            algorithm, algorithm
+        )
+        fig.suptitle(f"{algorithm} Training Progress")
 
         # ---- Q-table heatmap ----
 
diff --git a/tests/test_rl_algorithms_and_visualization.py b/tests/test_rl_algorithms_and_visualization.py
index 1c68092..3f3af62 100644
--- a/tests/test_rl_algorithms_and_visualization.py
+++ b/tests/test_rl_algorithms_and_visualization.py
@@ -147,11 +147,12 @@ def test_visualizer_timelapse_validates_inputs():
 
 
 def test_visualizer_timelapse_creates_animation_and_stores_reference():
-    agent = _trained_agent(QLearning, episodes=5, snapshot_interval=1)
+    agent = _trained_agent(Sarsa, episodes=5, snapshot_interval=1)
     visualizer = RLVisualizer(agent)
 
     fig, animation = visualizer.timelapse(interval=5, reward_window=2)
 
+    assert "SARSA" in fig._suptitle.get_text()
     assert animation is visualizer.animation
     assert animation is not None
     assert len(fig.axes) >= 2