Copybara-Service
3ec09f7296
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
472 lines
14 KiB
Python
472 lines
14 KiB
Python
# Copyright 2026 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Optimization trace report section."""
|
|
|
|
from collections.abc import Sequence
|
|
import math
|
|
from typing import Any
|
|
|
|
from mujoco.sysid.report.sections.base import ReportSection
|
|
from mujoco.sysid.report.utils import plotly_script_tag
|
|
import numpy as np
|
|
from numpy import typing as npt
|
|
from plotly import subplots as plt_subplots
|
|
import plotly.graph_objects as go
|
|
|
|
|
|
class OptimizationTrace(ReportSection):
|
|
"""Displays plots summarizing the optimization process."""
|
|
|
|
def __init__(
|
|
self,
|
|
title: str,
|
|
objective: Sequence[float],
|
|
candidate: Sequence[npt.ArrayLike],
|
|
bounds: (
|
|
tuple[Sequence[float] | np.ndarray, Sequence[float] | np.ndarray]
|
|
| None
|
|
) = None,
|
|
param_names: Sequence[str] | None = None,
|
|
log_diff: bool = True,
|
|
bound_eps: float = 1e-3,
|
|
dims_per_page: int = 6, # For candidate plot paging
|
|
anchor: str = "",
|
|
collapsible: bool = True,
|
|
):
|
|
super().__init__(collapsible=collapsible)
|
|
self._title = title
|
|
self._objective = objective
|
|
self._candidate = candidate # List of vectors
|
|
self._bounds = bounds
|
|
self._param_names = param_names
|
|
self._log_diff = log_diff
|
|
self._bound_eps = bound_eps
|
|
self._dims_per_page = dims_per_page
|
|
self._anchor = anchor
|
|
|
|
@property
|
|
def template_filename(self) -> str:
|
|
return "optimization_trace.html"
|
|
|
|
@property
|
|
def title(self) -> str:
|
|
return self._title
|
|
|
|
@property
|
|
def anchor(self) -> str:
|
|
return self._anchor
|
|
|
|
def header_sections(self) -> set[str]:
|
|
return {plotly_script_tag()}
|
|
|
|
def get_context(self) -> dict[str, Any]:
|
|
objective_fig = self._get_objective_figure()
|
|
candidate_figs = self._get_candidate_figures()
|
|
candidate_heatmap_fig = self._get_candidate_heatmap_figure()
|
|
|
|
config = {
|
|
"displayModeBar": True,
|
|
"displaylogo": False,
|
|
"toImageButtonOptions": {
|
|
"format": "svg",
|
|
"filename": f"{self._title}_optimization",
|
|
"height": 800,
|
|
"width": 1200,
|
|
"scale": 1,
|
|
},
|
|
"responsive": True,
|
|
}
|
|
|
|
return {
|
|
"title": self._title,
|
|
"objective_plot_html": (
|
|
objective_fig.to_html(
|
|
full_html=False, include_plotlyjs=False, config=config
|
|
)
|
|
if objective_fig
|
|
else None
|
|
),
|
|
"candidate_plots_html": (
|
|
[
|
|
fig.to_html(
|
|
full_html=False, include_plotlyjs=False, config=config
|
|
)
|
|
for fig in candidate_figs
|
|
]
|
|
if candidate_figs
|
|
else None
|
|
),
|
|
"candidate_heatmap_html": (
|
|
candidate_heatmap_fig.to_html(
|
|
full_html=False, include_plotlyjs=False, config=config
|
|
)
|
|
if candidate_heatmap_fig
|
|
else None
|
|
),
|
|
}
|
|
|
|
def _get_objective_figure(self) -> go.Figure | None:
|
|
"""Generates the objective function plot."""
|
|
if not self._objective:
|
|
return None
|
|
|
|
fig = go.Figure()
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
y=self._objective,
|
|
mode="lines+markers",
|
|
marker=dict(size=4),
|
|
line=dict(width=2),
|
|
name="Objective",
|
|
)
|
|
)
|
|
|
|
final_value = self._objective[-1]
|
|
if abs(final_value) < 1e-3 or abs(final_value) > 1e3:
|
|
final_str = f"{final_value:.2e}"
|
|
else:
|
|
final_str = f"{final_value:.4f}"
|
|
|
|
fig.update_layout(
|
|
title=f"Objective Over Time (Final: {final_str})",
|
|
xaxis_title="Iteration",
|
|
yaxis_title="Objective",
|
|
hovermode="x unified",
|
|
height=400,
|
|
autosize=True,
|
|
margin=dict(l=50, r=50, t=80, b=50),
|
|
template="plotly_white",
|
|
)
|
|
fig.update_xaxes(
|
|
showgrid=True, gridwidth=1, gridcolor="rgba(211, 211, 211, 0.7)"
|
|
)
|
|
fig.update_yaxes(
|
|
showgrid=True, gridwidth=1, gridcolor="rgba(211, 211, 211, 0.7)"
|
|
)
|
|
return fig
|
|
|
|
def _get_candidate_figures(self) -> list[go.Figure]:
|
|
"""Generates the parameter candidate plots (paged)."""
|
|
if not self._candidate:
|
|
return []
|
|
|
|
values = np.array(self._candidate).T # shape: (n_dim, n_iter)
|
|
n_dim, n_iter = values.shape
|
|
if n_iter <= 1:
|
|
return []
|
|
diffs = np.diff(values, axis=1) # shape: (n_dim, n_iter-1)
|
|
|
|
if self._bounds is not None:
|
|
mins = np.array(self._bounds[0])
|
|
maxs = np.array(self._bounds[1])
|
|
if not (mins.shape == (n_dim,) and maxs.shape == (n_dim,)):
|
|
raise ValueError("Bounds dimensions do not match parameter dimensions.")
|
|
else:
|
|
mins = np.full(n_dim, -np.inf)
|
|
maxs = np.full(n_dim, np.inf)
|
|
|
|
param_names = (
|
|
self._param_names
|
|
if self._param_names is not None
|
|
else [f"Dim {i}" for i in range(n_dim)]
|
|
)
|
|
if len(param_names) != n_dim:
|
|
raise ValueError(
|
|
"Number of parameter names does not match parameter dimensions."
|
|
)
|
|
|
|
n_pages = math.ceil(n_dim / self._dims_per_page)
|
|
figures = []
|
|
iterations = np.arange(n_iter)
|
|
iterations_diff = np.arange(1, n_iter)
|
|
|
|
for page in range(n_pages):
|
|
start_dim = page * self._dims_per_page
|
|
end_dim = min((page + 1) * self._dims_per_page, n_dim)
|
|
dims_in_page = end_dim - start_dim
|
|
page_param_names = param_names[start_dim:end_dim]
|
|
|
|
fig = plt_subplots.make_subplots(
|
|
rows=dims_in_page,
|
|
cols=2,
|
|
shared_xaxes=True,
|
|
subplot_titles=[
|
|
title
|
|
for name in page_param_names
|
|
for title in (f"{name} Value", f"{name} Δ")
|
|
],
|
|
vertical_spacing=max(0.02, 0.1 / dims_in_page),
|
|
)
|
|
|
|
for i, dim in enumerate(range(start_dim, end_dim)):
|
|
row_idx = i + 1
|
|
vals = values[dim, :]
|
|
diff_vals = diffs[dim, :]
|
|
lower, upper = mins[dim], maxs[dim]
|
|
|
|
# --- Value Plot (Col 1) ---
|
|
# Check for near-bound points
|
|
near_lower = np.abs(vals - lower) < self._bound_eps
|
|
near_upper = np.abs(vals - upper) < self._bound_eps
|
|
near_bound = near_lower | near_upper
|
|
|
|
# Plot segments with different colors if near bounds
|
|
for t in range(1, n_iter):
|
|
is_near_prev = near_bound[t - 1]
|
|
is_near_curr = near_bound[t]
|
|
color = (
|
|
"#d62728" if (is_near_prev or is_near_curr) else "#1f77b4"
|
|
) # Red if current or prev near bound
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=iterations[t - 1 : t + 1],
|
|
y=vals[t - 1 : t + 1],
|
|
mode="lines",
|
|
line=dict(color=color, width=2),
|
|
showlegend=False,
|
|
),
|
|
row=row_idx,
|
|
col=1,
|
|
)
|
|
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=iterations,
|
|
y=vals,
|
|
mode="markers",
|
|
marker=dict(
|
|
size=5,
|
|
color=["#d62728" if nb else "#1f77b4" for nb in near_bound],
|
|
symbol=[
|
|
"triangle-down"
|
|
if nl
|
|
else "triangle-up"
|
|
if nu
|
|
else "circle"
|
|
for nl, nu in zip(near_lower, near_upper, strict=True)
|
|
], # Triangles for bounds.
|
|
),
|
|
name=f"{param_names[dim]}",
|
|
showlegend=False,
|
|
hoverinfo="x+y+name",
|
|
),
|
|
row=row_idx,
|
|
col=1,
|
|
)
|
|
|
|
# Add bound lines if finite.
|
|
if np.isfinite(lower):
|
|
fig.add_hline(
|
|
y=lower,
|
|
line_dash="dash",
|
|
line_color="gray",
|
|
row=row_idx, # pyright: ignore[reportArgumentType]
|
|
col=1, # pyright: ignore[reportArgumentType]
|
|
opacity=0.5,
|
|
)
|
|
if np.isfinite(upper):
|
|
fig.add_hline(
|
|
y=upper,
|
|
line_dash="dash",
|
|
line_color="gray",
|
|
row=row_idx, # pyright: ignore[reportArgumentType]
|
|
col=1, # pyright: ignore[reportArgumentType]
|
|
opacity=0.5,
|
|
)
|
|
|
|
# Add final value annotation.
|
|
final_val = vals[-1]
|
|
final_str = (
|
|
f"{final_val:.2e}"
|
|
if abs(final_val) < 1e-3 or abs(final_val) > 1e3
|
|
else f"{final_val:.4f}"
|
|
)
|
|
fig.add_annotation(
|
|
x=iterations[-1],
|
|
y=final_val,
|
|
text=final_str,
|
|
showarrow=True,
|
|
arrowhead=1,
|
|
ax=20,
|
|
ay=-30,
|
|
row=row_idx,
|
|
col=1,
|
|
font=dict(color="blue", size=9),
|
|
)
|
|
|
|
# --- Diff Plot (Col 2) ---
|
|
if self._log_diff:
|
|
eps = 1e-12
|
|
plot_diff_vals = np.log10(np.abs(diff_vals) + eps)
|
|
yaxis_title = "log |Δ|"
|
|
else:
|
|
plot_diff_vals = diff_vals
|
|
yaxis_title = "Δ"
|
|
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=iterations_diff,
|
|
y=plot_diff_vals,
|
|
mode="lines+markers",
|
|
marker=dict(symbol="x", size=5, color="orange"),
|
|
line=dict(width=1.5, color="orange"),
|
|
name=f"Δ {param_names[dim]}",
|
|
showlegend=False,
|
|
hoverinfo="x+y+name",
|
|
),
|
|
row=row_idx,
|
|
col=2,
|
|
)
|
|
fig.update_yaxes(
|
|
title_text=yaxis_title, row=row_idx, col=2, title_font_size=10
|
|
)
|
|
|
|
# --- Layout Updates for the Page Figure ---
|
|
fig.update_layout(
|
|
title=(
|
|
f"Candidate Values and Changes (Page {page + 1}/{n_pages}, Dims"
|
|
f" {start_dim}-{end_dim - 1})"
|
|
),
|
|
height=max(400, 200 * dims_in_page), # Adjust height based on dims
|
|
autosize=True,
|
|
margin=dict(l=60, r=30, t=100, b=50),
|
|
hovermode="x unified",
|
|
template="plotly_white",
|
|
)
|
|
fig.update_xaxes(
|
|
showgrid=True,
|
|
gridwidth=1,
|
|
gridcolor="rgba(211, 211, 211, 0.7)",
|
|
zeroline=False,
|
|
)
|
|
fig.update_yaxes(
|
|
showgrid=True,
|
|
gridwidth=1,
|
|
gridcolor="rgba(211, 211, 211, 0.7)",
|
|
zeroline=False,
|
|
)
|
|
|
|
# Add common x-axis label to the bottom row.
|
|
fig.update_xaxes(title_text="Iteration", row=dims_in_page, col=1)
|
|
fig.update_xaxes(title_text="Iteration", row=dims_in_page, col=2)
|
|
|
|
for annotation in fig.layout.annotations:
|
|
annotation.font.size = 10
|
|
|
|
figures.append(fig)
|
|
|
|
return figures
|
|
|
|
def _get_candidate_heatmap_figure(self) -> go.Figure | None:
|
|
"""Generates the parameter candidate heatmap."""
|
|
if not self._candidate:
|
|
return None
|
|
|
|
data = np.array(self._candidate).T # shape: (n_dim, n_iter)
|
|
n_dim, n_iter = data.shape
|
|
|
|
param_names = (
|
|
self._param_names
|
|
if self._param_names is not None
|
|
else [f"Dim {i}" for i in range(n_dim)]
|
|
)
|
|
if len(param_names) != n_dim:
|
|
raise ValueError(
|
|
"Number of parameter names does not match parameter dimensions."
|
|
)
|
|
|
|
# Normalize data for heatmap colors if bounds are provided.
|
|
heatmap_data = data.copy()
|
|
normalize = self._bounds is not None
|
|
if normalize and self._bounds is not None:
|
|
min_bounds, max_bounds = self._bounds
|
|
if not (len(min_bounds) == len(max_bounds) == n_dim):
|
|
raise ValueError("Bounds dimensions do not match parameter dimensions.")
|
|
for i in range(n_dim):
|
|
min_val, max_val = min_bounds[i], max_bounds[i]
|
|
denom = max_val - min_val if max_val > min_val else 1.0
|
|
clipped_vals = np.clip(data[i], min_val, max_val)
|
|
heatmap_data[i] = (
|
|
(clipped_vals - min_val) / denom if denom != 0 else 0.5
|
|
) # Center if range is zero
|
|
else:
|
|
row_mins = np.min(data, axis=1, keepdims=True)
|
|
row_maxs = np.max(data, axis=1, keepdims=True)
|
|
row_ranges = row_maxs - row_mins
|
|
row_ranges[row_ranges == 0] = 1.0 # Avoid division by zero
|
|
heatmap_data = (data - row_mins) / row_ranges
|
|
|
|
fig = go.Figure(
|
|
data=go.Heatmap(
|
|
z=heatmap_data[::-1],
|
|
x=np.arange(n_iter),
|
|
y=list(reversed(param_names)),
|
|
colorscale="RdBu",
|
|
colorbar=dict(
|
|
title="Normalized Value"
|
|
if normalize
|
|
else "Row-Normalized Value"
|
|
),
|
|
hovertemplate=(
|
|
"Iter: %{x}<br>Param: %{y}<br>Value:"
|
|
" %{customdata:.4f}<extra></extra>"
|
|
),
|
|
customdata=data[::-1].tolist(),
|
|
)
|
|
)
|
|
|
|
# Add markers for bound hits if bounds exist.
|
|
bound_markers_x = []
|
|
bound_markers_y = []
|
|
if self._bounds is not None:
|
|
min_bounds, max_bounds = self._bounds
|
|
for dim in range(n_dim):
|
|
min_val, max_val = min_bounds[dim], max_bounds[dim]
|
|
for iter_idx, val in enumerate(data[dim]):
|
|
if (
|
|
abs(val - min_val) < self._bound_eps
|
|
or abs(val - max_val) < self._bound_eps
|
|
):
|
|
bound_markers_x.append(iter_idx)
|
|
bound_markers_y.append(param_names[dim])
|
|
|
|
if bound_markers_x:
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=bound_markers_x,
|
|
y=bound_markers_y,
|
|
mode="markers",
|
|
marker=dict(color="black", size=6, symbol="x"),
|
|
name="At Bound",
|
|
showlegend=False,
|
|
hoverinfo="skip",
|
|
)
|
|
)
|
|
|
|
fig.update_layout(
|
|
title="Candidate Parameter Heatmap",
|
|
xaxis_title="Iteration",
|
|
yaxis_title="Parameter",
|
|
height=max(400, 30 * n_dim),
|
|
autosize=True,
|
|
margin=dict(l=150, r=50, t=80, b=50),
|
|
yaxis=dict(
|
|
tickmode="array", tickvals=param_names, ticktext=param_names
|
|
),
|
|
template="plotly_white",
|
|
)
|
|
|
|
return fig
|