Step Log visualization optimization

[clemens-fricke]

While using the step log visualization tool more and more a few issue came to light.

1. When hiding a trace, by clicking on it in the legend, it would reappear after changing the episode
2. Colors in the top plot repeated for reward and the first objective
3. Only one legend existed and not separate legends for each subplot
4. File sizes get really big if the step log is big
5. Responsiveness of the visualization was sluggish for large files

To combat this the visualization was optimized. The biggest part is the fix for 1, 4, and 5 which were fixed by changing how the slider works. Before all episode traces were plotted separately at the same time but than hidden for all episode but the one that was shown. Now one one trace is plotted for each variable and the data is exchanged for each episode. This solves both the responsiveness, file size and trace hiding problems.

Quantifiable is the easiest for point 4: For a specifically large step log we get a visualization file that is 374Mb in size and for the optimized the file size shrinks to 78Mb.

The new visualization looks like this:

[Copilot]

Pull request overview

This PR optimizes the step log visualization tool to address performance and usability issues. The key improvement is replacing the approach of pre-plotting all episodes (and hiding unselected ones) with dynamically updating trace data via the slider control. This change significantly reduces file sizes (e.g., from 374MB to 78MB in large cases) and improves responsiveness while fixing trace hiding bugs and color repetition issues.

Key changes:

• Refactored slider implementation to update trace data dynamically instead of toggling visibility of pre-rendered traces
• Added separate legends for each subplot instead of a single shared legend
• Improved memory management with explicit cleanup of intermediate dataframes

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The variable names per_episode_x and per_episode_y are not fully descriptive. Consider renaming to episode_x_data and episode_y_data or per_episode_x_values and per_episode_y_values to make it clearer that these are lists of data values for each episode.

[Copilot]

[nitpick] The default value changed from np.iinfo(int).max to 10**12. While 10**12 is simpler, it's a magic number that could theoretically be exceeded (though unlikely in practice). Consider using sys.maxsize or float('inf') which are more idiomatic for representing "no upper limit", or at least add a comment explaining why 10**12 was chosen.

[Copilot]

Incomplete sentence in documentation example. The sentence ends with "from index 3 to 5 (6 is exclusive) from" but doesn't specify what the data is extracted from (e.g., "from the dataframe" or "from each row").

Suggested change

	to 5 (6 is exclusive) from
	to 5 (6 is exclusive) from the dataframe.

[Copilot]

The error message "Empty or malformed step log." is less informative than the original message that included the file path. Consider including the file path to help users identify which file is problematic: f"The provided step log file {step_log_file} is empty or does not follow the current format."

Suggested change

	raise ValueError("Empty or malformed step log.")
	raise ValueError(f"The provided step log file {step_log_file} is empty or does not follow the current format.")

[Copilot]

The episode selection logic no longer handles the case where no episodes match the specified range. The original code had a try-except block to catch IndexError and provide a helpful error message showing the available episode range. Without this error handling, users will get a cryptic IndexError or potentially an empty dataframe that leads to confusing behavior later. Consider adding back error handling:

if not mask.any():
    raise IndexError(
        f"Could not find any episode in the range {episode_start} to {episode_end}. "
        f"The available episodes range from {df['episodes'].min()} to {df['episodes'].max()}"
    )

Suggested change

	mask = (df["episodes"] >= episode_start) & (df["episodes"] <= episode_end)
	mask = (df["episodes"] >= episode_start) & (df["episodes"] <= episode_end)
	if not mask.any():
	raise IndexError(
	f"Could not find any episode in the range {episode_start} to {episode_end}. "
	f"The available episodes range from {df['episodes'].min()} to {df['episodes'].max()}"
	)

[Copilot]

This comment mentions a limitation but it's unclear what the implications are or if this is a known issue that needs to be addressed. Consider clarifying: (1) why this limitation exists, (2) whether it affects the visualization, and (3) if there's a planned solution or workaround. If this is important context, it should be more explicit.

Suggested change

	# Not able to extract the reset values without info field. prev_obs field shows
	# last obs from previous episode not reset value in first step.
	# Limitation: Unable to extract the reset values for each episode because the required
	# "info" field is missing from the step log data. As a result, the "prev_obs" field
	# only contains the last observation from the previous episode, not the reset value at
	# the start of the current episode. This means that the initial state of each episode
	# cannot be visualized, which may affect analyses that rely on episode initialization.
	# Currently, there is no workaround unless the data collection process is updated to
	# include the "info" field with reset values.

[Copilot]

Type annotations were removed from objective_observation and design_variable parameters. This reduces code clarity and type safety. Consider keeping the type annotations: objective_observation: list[tuple[str, tuple[int, int]]] = [("obs", (0, 1))] and design_variable: list[tuple[str, tuple[int, int]]] = [("obs", (1, None))].

Suggested change

	objective_observation=[("obs", (0, 1))],
	design_variable=[("obs", (1, None))],
	objective_observation: list[tuple[str, tuple[int, int]]] = [("obs", (0, 1))],
	design_variable: list[tuple[str, tuple[int, int]]] = [("obs", (1, None))],

[Copilot]

Typo in comment: "everthing" should be "everything".

Suggested change

	# combine everthing into a single dataframe
	# combine everything into a single dataframe

[Copilot]

The logic for creating separate legends uses list(fig.select_yaxes())[::2] which assumes a specific structure that might be fragile if the subplot configuration changes. The slicing [::2] selects every second y-axis, but this implicit assumption isn't documented. Consider adding a comment explaining why every second y-axis is selected, or use a more explicit approach that's self-documenting.

Suggested change

	for i, yaxis in enumerate(list(fig.select_yaxes())[::2], 1):
	legend_name = f"legend{i}"
	fig.update_layout(
	{legend_name: dict(y=yaxis.domain[1], yanchor="top")},
	showlegend=True,
	)
	fig.update_traces(row=i, legend=legend_name)
	# Explicitly place legends for each subplot row.
	# We assume a subplot configuration of 2 rows, 1 column, with secondary y-axes in the first row.
	# This logic is robust to changes in subplot configuration.
	subplot_rows = 2
	subplot_cols = 1
	for row in range(1, subplot_rows + 1):
	# Get the primary y-axis for this subplot
	yaxes = [y for y in fig.select_yaxes() if y.anchor == f"x{row}"]
	if not yaxes:
	continue
	yaxis = yaxes[0]
	legend_name = f"legend{row}"
	fig.update_layout(
	{legend_name: dict(y=yaxis.domain[1], yanchor="top")},
	showlegend=True,
	)
	fig.update_traces(row=row, legend=legend_name)

[Copilot]

The condition not isinstance(figure_size, list) and figure_size == "auto" is incorrect. It should check for both list and tuple since figure_size is typed as Union[tuple[int, int], Literal["auto"]] and the error message mentions tuple. The condition should be: figure_size == "auto" or not isinstance(figure_size, (list, tuple)) and figure_size == "auto". As written, if figure_size is a list, it will fall through to the else clause and raise an error even though lists are handled in line 700.