1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128 | from collections.abc import Callable
from typing import Optional
# import matplotlib
# matplotlib.use('TkAgg') # Can change to 'Agg' for non-interactive mode
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
from lit_sb3.wrappers.monitor import load_results
X_TIMESTEPS = "timesteps"
X_EPISODES = "episodes"
X_WALLTIME = "walltime_hrs"
POSSIBLE_X_AXES = [X_TIMESTEPS, X_EPISODES, X_WALLTIME]
EPISODES_WINDOW = 100
def rolling_window(array: np.ndarray, window: int) -> np.ndarray:
"""Apply a rolling window to a np.ndarray.
Args:
array: the input Array
window: length of the rolling window
Returns:
rolling window on the input array
"""
shape = array.shape[:-1] + (array.shape[-1] - window + 1, window) # noqa: RUF005
strides = (*array.strides, array.strides[-1])
return np.lib.stride_tricks.as_strided(array, shape=shape, strides=strides)
def window_func(var_1: np.ndarray, var_2: np.ndarray, window: int, func: Callable) -> tuple[np.ndarray, np.ndarray]:
"""Apply a function to the rolling window of 2 arrays.
Args:
var_1: variable 1
var_2: variable 2
window: length of the rolling window
func: function to apply on the rolling window on variable 2 (such as np.mean)
Returns:
the rolling output with applied function
"""
var_2_window = rolling_window(var_2, window)
function_on_var2 = func(var_2_window, axis=-1)
return var_1[window - 1 :], function_on_var2
def ts2xy(data_frame: pd.DataFrame, x_axis: str) -> tuple[np.ndarray, np.ndarray]:
"""Decompose a data frame variable to x and ys.
(y = episodic return).
Args:
data_frame: the input data
x_axis: the x-axis for the x and y output
(can be X_TIMESTEPS='timesteps', X_EPISODES='episodes' or X_WALLTIME='walltime_hrs')
Returns:
the x and y output
"""
if x_axis == X_TIMESTEPS:
x_var = np.cumsum(data_frame.l.values)
y_var = data_frame.r.values
elif x_axis == X_EPISODES:
x_var = np.arange(len(data_frame))
y_var = data_frame.r.values
elif x_axis == X_WALLTIME:
# Convert to hours
x_var = data_frame.t.values / 3600.0
y_var = data_frame.r.values
else:
raise NotImplementedError(f"Unsupported {x_axis=}, please use one of {POSSIBLE_X_AXES}")
return x_var, y_var
def plot_curves(
xy_list: list[tuple[np.ndarray, np.ndarray]], x_axis: str, title: str, figsize: tuple[int, int] = (8, 2)
) -> None:
"""Plot the curves.
Args:
xy_list: the x and y coordinates to plot
x_axis: the axis for the x and y output
(can be X_TIMESTEPS='timesteps', X_EPISODES='episodes' or X_WALLTIME='walltime_hrs')
title: the title of the plot
figsize: Size of the figure (width, height)
"""
plt.figure(title, figsize=figsize)
max_x = max(xy[0][-1] for xy in xy_list)
min_x = 0
for _, (x, y) in enumerate(xy_list):
plt.scatter(x, y, s=2)
# Do not plot the smoothed curve at all if the timeseries is shorter than window size.
if x.shape[0] >= EPISODES_WINDOW:
# Compute and plot rolling mean with window of size EPISODE_WINDOW
x, y_mean = window_func(x, y, EPISODES_WINDOW, np.mean)
plt.plot(x, y_mean)
plt.xlim(min_x, max_x)
plt.title(title)
plt.xlabel(x_axis)
plt.ylabel("Episode Rewards")
plt.tight_layout()
def plot_results(
dirs: list[str], num_timesteps: int | None, x_axis: str, task_name: str, figsize: tuple[int, int] = (8, 2)
) -> None:
"""Plot the results using csv files from ``Monitor`` wrapper.
Args:
dirs: the save location of the results to plot
num_timesteps: only plot the points below this value
x_axis: the axis for the x and y output
(can be X_TIMESTEPS='timesteps', X_EPISODES='episodes' or X_WALLTIME='walltime_hrs')
task_name: the title of the task to plot
figsize: Size of the figure (width, height)
"""
data_frames = []
for folder in dirs:
data_frame = load_results(folder)
if num_timesteps is not None:
data_frame = data_frame[data_frame.l.cumsum() <= num_timesteps]
data_frames.append(data_frame)
xy_list = [ts2xy(data_frame, x_axis) for data_frame in data_frames]
plot_curves(xy_list, x_axis, task_name, figsize)
|