Copybara-Service
3ec09f7296
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
429 lines
13 KiB
Python
429 lines
13 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.
|
|
# ==============================================================================
|
|
"""Parameter distribution report section."""
|
|
|
|
import math
|
|
from typing import Any
|
|
|
|
from mujoco.sysid._src import parameter
|
|
from mujoco.sysid.report.sections.base import ReportSection
|
|
from mujoco.sysid.report.utils import plotly_script_tag
|
|
import numpy as np
|
|
from plotly import subplots as plt_subplots
|
|
import plotly.graph_objects as go
|
|
|
|
|
|
class ParameterDistribution(ReportSection):
|
|
"""Displays the identified parameters relative to their bounds and nominal values.
|
|
|
|
Visualization: One horizontal track per parameter with markers for Nominal,
|
|
Identified, and Bounds.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
title: str,
|
|
opt_params: parameter.ParameterDict,
|
|
initial_params: parameter.ParameterDict,
|
|
confidence_intervals: np.ndarray | None = None,
|
|
height_per_param: int = 60,
|
|
anchor: str = "",
|
|
collapsible: bool = True,
|
|
):
|
|
super().__init__(collapsible=collapsible)
|
|
self._title = title
|
|
self._anchor = anchor
|
|
|
|
self._opt_params = opt_params
|
|
self._initial_params = initial_params
|
|
self._nominal_params = initial_params.copy()
|
|
self._nominal_params.reset()
|
|
|
|
self._x_nominal = self._nominal_params.as_vector()
|
|
self._x_hat = self._opt_params.as_vector()
|
|
self._x_initial = self._initial_params.as_vector()
|
|
self._confidence_intervals = confidence_intervals
|
|
self._height_per_param = height_per_param
|
|
|
|
@property
|
|
def template_filename(self) -> str:
|
|
return "plot_generic.html"
|
|
|
|
@property
|
|
def title(self) -> str:
|
|
return self._title
|
|
|
|
@property
|
|
def anchor(self) -> str:
|
|
return self._anchor
|
|
|
|
def header_includes(self) -> set[str]:
|
|
return {plotly_script_tag()}
|
|
|
|
def get_context(self) -> dict[str, Any]:
|
|
fig = self._build_figure()
|
|
|
|
config = {
|
|
"displayModeBar": True,
|
|
"displaylogo": False,
|
|
"responsive": True,
|
|
"toImageButtonOptions": {
|
|
"format": "svg",
|
|
"filename": f"{self._title}_distribution",
|
|
"height": max(400, len(self._x_hat) * 100),
|
|
"width": 1200,
|
|
"scale": 1,
|
|
},
|
|
}
|
|
|
|
return {
|
|
"title": self._title,
|
|
"plot_div": fig.to_html(
|
|
full_html=False, include_plotlyjs=False, config=config
|
|
),
|
|
"caption": (
|
|
"<b>Visualization Guide:</b><br><b>Error Bars ( — )</b>: 95%"
|
|
" Confidence Interval. Smaller is better.<br><span"
|
|
" style='color:green'>◆ High Confidence</span>: Interval is <"
|
|
" 0.5% of the parameter range (error bar may be"
|
|
" invisible).<br><span style='color:blue'>◆ Identified</span>:"
|
|
" Standard confidence interval.<br><span style='color:red'>◆"
|
|
" Unconstrained</span>: Interval is infinite or larger than range"
|
|
" (Red error bar).<br><span style='color:red'>x Nominal</span>,"
|
|
" <span style='color:blue'>- - Bounds</span>: Reference"
|
|
" values.<br>* parameter is frozen."
|
|
),
|
|
}
|
|
|
|
def _build_figure(self) -> go.Figure:
|
|
plot_items = []
|
|
non_frozen_idx = 0
|
|
|
|
for param_name, param in self._opt_params.parameters.items():
|
|
# Get bounds for this specific parameter
|
|
p_min, p_max = param.get_bounds()
|
|
|
|
if param.size == 1:
|
|
flat_value = (
|
|
param.value if np.isscalar(param.value) else param.value.item()
|
|
)
|
|
lb = p_min.item()
|
|
ub = p_max.item()
|
|
|
|
if param.frozen:
|
|
plot_items.append({
|
|
"name": param_name,
|
|
"value": flat_value,
|
|
"is_frozen": True,
|
|
"lb": lb,
|
|
"ub": ub,
|
|
"nominal": flat_value if self._x_nominal is not None else None,
|
|
})
|
|
else:
|
|
conf = (
|
|
self._confidence_intervals[non_frozen_idx]
|
|
if self._confidence_intervals is not None
|
|
else None
|
|
)
|
|
plot_items.append({
|
|
"name": param_name,
|
|
"value": self._x_hat[non_frozen_idx],
|
|
"is_frozen": False,
|
|
"lb": lb,
|
|
"ub": ub,
|
|
"nominal": (
|
|
self._x_nominal[non_frozen_idx]
|
|
if self._x_nominal is not None
|
|
else None
|
|
),
|
|
"conf": conf,
|
|
})
|
|
non_frozen_idx += 1
|
|
else:
|
|
for i in range(param.size):
|
|
if param.shape == (param.size,):
|
|
element_name = f"{param_name}[{i}]"
|
|
val_frozen = param.value[i]
|
|
lb = p_min[i]
|
|
ub = p_max[i]
|
|
else:
|
|
multi_idx = np.unravel_index(i, param.shape)
|
|
idx_str = ",".join(str(x) for x in multi_idx)
|
|
element_name = f"{param_name}[{idx_str}]"
|
|
val_frozen = param.value[multi_idx]
|
|
lb = p_min[i]
|
|
ub = p_max[i]
|
|
|
|
if param.frozen:
|
|
plot_items.append({
|
|
"name": element_name,
|
|
"value": val_frozen,
|
|
"is_frozen": True,
|
|
"lb": lb,
|
|
"ub": ub,
|
|
"nominal": val_frozen if self._x_nominal is not None else None,
|
|
})
|
|
else:
|
|
conf = (
|
|
self._confidence_intervals[non_frozen_idx]
|
|
if self._confidence_intervals is not None
|
|
else None
|
|
)
|
|
plot_items.append({
|
|
"name": element_name,
|
|
"value": self._x_hat[non_frozen_idx],
|
|
"is_frozen": False,
|
|
"lb": lb,
|
|
"ub": ub,
|
|
"nominal": (
|
|
self._x_nominal[non_frozen_idx]
|
|
if self._x_nominal is not None
|
|
else None
|
|
),
|
|
"conf": conf,
|
|
})
|
|
non_frozen_idx += 1
|
|
|
|
n_params = len(plot_items)
|
|
|
|
n_cols = 8 if n_params > 1 else 1
|
|
n_rows = math.ceil(n_params / n_cols)
|
|
|
|
# Prepare subplot titles with conditional formatting
|
|
subplot_titles = []
|
|
for p in plot_items:
|
|
if p["is_frozen"]:
|
|
subplot_titles.append(f"<span style='color: grey;'>{p['name']}*</span>")
|
|
else:
|
|
subplot_titles.append(p["name"])
|
|
|
|
# Create subplots
|
|
fig = plt_subplots.make_subplots(
|
|
rows=n_rows,
|
|
cols=n_cols,
|
|
subplot_titles=subplot_titles,
|
|
vertical_spacing=0.05,
|
|
horizontal_spacing=0.02,
|
|
)
|
|
|
|
# Add dummy trace for Bounds legend
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=[None],
|
|
y=[None],
|
|
mode="lines",
|
|
line=dict(color="blue", width=3, dash="dash"),
|
|
name="Bounds",
|
|
)
|
|
)
|
|
|
|
shown_legends = set()
|
|
|
|
for i, item in enumerate(plot_items):
|
|
row = (i // n_cols) + 1
|
|
col = (i % n_cols) + 1
|
|
|
|
val = item["value"]
|
|
lb, ub = item["lb"], item["ub"]
|
|
|
|
# 1. Bounds lines (For items with valid bounds)
|
|
if lb is not None and ub is not None:
|
|
# Lower Bound (Horizontal line at y=lb)
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=[-1, 1],
|
|
y=[lb, lb],
|
|
mode="lines",
|
|
line=dict(color="blue", width=3, dash="dash"),
|
|
showlegend=False,
|
|
hoverinfo="skip",
|
|
),
|
|
row=row,
|
|
col=col,
|
|
)
|
|
# Upper Bound (Horizontal line at y=ub)
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=[-1, 1],
|
|
y=[ub, ub],
|
|
mode="lines",
|
|
line=dict(color="blue", width=3, dash="dash"),
|
|
showlegend=False,
|
|
hoverinfo="skip",
|
|
),
|
|
row=row,
|
|
col=col,
|
|
)
|
|
|
|
# Calculate margin based on bounds
|
|
dist = ub - lb
|
|
if dist > 0:
|
|
margin = dist * 0.2
|
|
y_range_min = lb
|
|
y_range_max = ub
|
|
else:
|
|
margin = abs(val) * 0.2 if val != 0 else 1.0
|
|
y_range_min = val
|
|
y_range_max = val
|
|
else:
|
|
# Fallback
|
|
margin = abs(val) * 0.2 if val != 0 else 1.0
|
|
y_range_min = val
|
|
y_range_max = val
|
|
|
|
# 2. Add Traces (Frozen vs Optimized)
|
|
if item["is_frozen"]:
|
|
# Frozen Parameter
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=[0],
|
|
y=[val],
|
|
mode="markers",
|
|
marker=dict(
|
|
symbol="circle", size=10, color="gray", opacity=0.7
|
|
),
|
|
name="Frozen",
|
|
showlegend=False,
|
|
hoverinfo="skip",
|
|
),
|
|
row=row,
|
|
col=col,
|
|
)
|
|
else:
|
|
# Optimized Parameter
|
|
|
|
# Determine Confidence Status
|
|
trace_name = "Identified"
|
|
marker_color = "blue"
|
|
legend_group = "identified"
|
|
|
|
error_y = None
|
|
|
|
if item.get("conf") is not None:
|
|
interval = item["conf"]
|
|
min_bound = item["lb"]
|
|
max_bound = item["ub"]
|
|
|
|
error_val = interval
|
|
|
|
rng = 0
|
|
if min_bound is not None and max_bound is not None:
|
|
rng = max_bound - min_bound
|
|
|
|
if not np.isfinite(interval) or (
|
|
rng > 0 and 2.0 * interval > 1.0 * rng
|
|
):
|
|
# Unconstrained
|
|
trace_name = "Unconstrained"
|
|
marker_color = "red"
|
|
legend_group = "unconstrained"
|
|
|
|
if rng > 0:
|
|
error_val = rng
|
|
else:
|
|
error_val = interval # or large value call fallback
|
|
|
|
elif rng > 0 and interval <= 0.005 * rng:
|
|
trace_name = "High Confidence"
|
|
marker_color = "green"
|
|
legend_group = "high_conf"
|
|
|
|
error_bar_color = marker_color
|
|
error_y = dict(
|
|
type="data",
|
|
array=[error_val],
|
|
visible=True,
|
|
thickness=1.5,
|
|
width=3,
|
|
color=error_bar_color,
|
|
)
|
|
|
|
show_leg = trace_name not in shown_legends
|
|
if show_leg:
|
|
shown_legends.add(trace_name)
|
|
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=[0],
|
|
y=[val],
|
|
mode="markers",
|
|
marker=dict(symbol="diamond", size=12, color=marker_color),
|
|
name=trace_name,
|
|
showlegend=show_leg,
|
|
legendgroup=legend_group,
|
|
hovertemplate=(
|
|
f"{trace_name}: %{{y}} ±"
|
|
f" {item.get('conf', 0):.4g}<extra></extra>"
|
|
),
|
|
error_y=error_y,
|
|
),
|
|
row=row,
|
|
col=col,
|
|
)
|
|
|
|
# Nominal Value (if exists)
|
|
if item["nominal"] is not None:
|
|
show_leg_nom = "Nominal" not in shown_legends
|
|
if show_leg_nom:
|
|
shown_legends.add("Nominal")
|
|
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=[0],
|
|
y=[item["nominal"]],
|
|
mode="markers",
|
|
marker=dict(symbol="x", size=10, color="red"),
|
|
name="Nominal",
|
|
showlegend=show_leg_nom,
|
|
legendgroup="nominal",
|
|
hovertemplate="Nominal: %{y}<extra></extra>",
|
|
),
|
|
row=row,
|
|
col=col,
|
|
)
|
|
|
|
fig.update_yaxes(
|
|
range=[y_range_min - margin, y_range_max + margin],
|
|
showgrid=True,
|
|
zeroline=False,
|
|
row=row,
|
|
col=col,
|
|
)
|
|
fig.update_xaxes(
|
|
range=[-1, 1],
|
|
showgrid=False,
|
|
zeroline=False,
|
|
showticklabels=False,
|
|
row=row,
|
|
col=col,
|
|
)
|
|
|
|
# Global Layout
|
|
total_height = max(
|
|
400, n_rows * 180
|
|
) # Increased height per row for better spacing
|
|
fig.update_layout(
|
|
height=total_height,
|
|
template="plotly_white",
|
|
margin=dict(l=60, r=60, t=150, b=60), # Increased side margins
|
|
autosize=True,
|
|
legend=dict(
|
|
orientation="h", yanchor="bottom", y=1.02, xanchor="center", x=0.5
|
|
),
|
|
hovermode="closest",
|
|
)
|
|
|
|
return fig
|