#917 - Integrated the PrimaiteSession into all tests.

- Ran a full pre-commit hook and thus encountered tons of fixes required

src/primaite/data_viz/__init__.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class PlotlyTemplate(Enum):
+    """The built-in plotly templates."""
+
+    PLOTLY = "plotly"
+    PLOTLY_WHITE = "plotly_white"
+    PLOTLY_DARK = "plotly_dark"
+    GGPLOT2 = "ggplot2"
+    SEABORN = "seaborn"
+    SIMPLE_WHITE = "simple_white"
+    NONE = "none"

src/primaite/data_viz/session_plots.py

@@ -0,0 +1,73 @@
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+import plotly.graph_objects as go
+import polars as pl
+import yaml
+from plotly.graph_objs import Figure
+
+from primaite import _PLATFORM_DIRS
+
+
+def _get_plotly_config() -> Dict:
+    """Get the plotly config from primaite_config.yaml."""
+    user_config_path = _PLATFORM_DIRS.user_config_path / "primaite_config.yaml"
+    with open(user_config_path, "r") as file:
+        primaite_config = yaml.safe_load(file)
+    return primaite_config["session"]["outputs"]["plots"]
+
+
+def plot_av_reward_per_episode(
+    av_reward_per_episode_csv: Union[str, Path],
+    title: Optional[str] = None,
+    subtitle: Optional[str] = None,
+) -> Figure:
+    """
+    Plot the average reward per episode from a csv session output.
+
+    :param av_reward_per_episode_csv: The average reward per episode csv
+        file path.
+    :param title: The plot title. This is optional.
+    :param subtitle: The plot subtitle. This is optional.
+    :return: The plot as an instance of ``plotly.graph_objs._figure.Figure``.
+    """
+    df = pl.read_csv(av_reward_per_episode_csv)
+
+    if title:
+        if subtitle:
+            title = f"{title} <br>{subtitle}</sup>"
+    else:
+        if subtitle:
+            title = subtitle
+
+    config = _get_plotly_config()
+    layout = go.Layout(
+        autosize=config["size"]["auto_size"],
+        width=config["size"]["width"],
+        height=config["size"]["height"],
+    )
+    # Create the line graph with a colored line
+    fig = go.Figure(layout=layout)
+    fig.update_layout(template=config["template"])
+    fig.add_trace(
+        go.Scatter(
+            x=df["Episode"],
+            y=df["Average Reward"],
+            mode="lines",
+            name="Mean Reward per Episode",
+        )
+    )
+
+    # Set the layout of the graph
+    fig.update_layout(
+        xaxis={
+            "title": "Episode",
+            "type": "linear",
+            "rangeslider": {"visible": config["range_slider"]},
+        },
+        yaxis={"title": "Average Reward"},
+        title=title,
+        showlegend=False,
+    )
+
+    return fig