# 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": (
"Visualization Guide:
Error Bars ( — ): 95%"
" Confidence Interval. Smaller is better.
◆ High Confidence: Interval is <"
" 0.5% of the parameter range (error bar may be"
" invisible).
◆ Identified:"
" Standard confidence interval.
◆"
" Unconstrained: Interval is infinite or larger than range"
" (Red error bar).
x Nominal,"
" - - Bounds: Reference"
" values.
* 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"{p['name']}*")
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}"
),
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}",
),
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