Merge pull request #3079 from aftersomemath:sysid-pr
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
This commit is contained in:
Copybara-Service
@@ -0,0 +1,216 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Video report section for system identification results."""
|
||||
|
||||
from collections.abc import Callable
|
||||
import os
|
||||
import pathlib
|
||||
from typing import Any
|
||||
|
||||
import mujoco
|
||||
import mujoco.rollout
|
||||
from mujoco.sysid._src import model_modifier
|
||||
from mujoco.sysid._src import parameter
|
||||
from mujoco.sysid._src.plotting import render_rollout
|
||||
from mujoco.sysid._src.trajectory import SystemTrajectory
|
||||
from mujoco.sysid.report.sections.base import ReportSection
|
||||
|
||||
|
||||
def spec_apply(spec, attrs, values):
|
||||
def apply_to_geoms_recursive(body):
|
||||
for g in body.geoms:
|
||||
for attr, value in zip(attrs, values, strict=True):
|
||||
setattr(g, attr, value)
|
||||
for child_body in body.bodies:
|
||||
apply_to_geoms_recursive(child_body)
|
||||
|
||||
for top_body in spec.worldbody.bodies:
|
||||
apply_to_geoms_recursive(top_body)
|
||||
|
||||
|
||||
def generate_video_from_trajectories(
|
||||
initial_params: parameter.ParameterDict,
|
||||
opt_params: parameter.ParameterDict,
|
||||
_build_model: Callable[
|
||||
[parameter.ParameterDict, mujoco.MjSpec], mujoco.MjModel
|
||||
],
|
||||
trajectories: list[SystemTrajectory],
|
||||
model_spec: mujoco.MjSpec,
|
||||
output_filepath: os.PathLike[str],
|
||||
render_initial: bool = True,
|
||||
render_nominal: bool = True,
|
||||
render_opt: bool = True,
|
||||
height: int = 480,
|
||||
width: int = 640,
|
||||
fovy: float = 60,
|
||||
camera: str | int = -1,
|
||||
fps: int = 60,
|
||||
) -> pathlib.Path:
|
||||
"""Render trajectories and concatenate into a single video.
|
||||
|
||||
Each trajectory is rendered with initial/nominal/optimized parameters
|
||||
overlaid, then all frames are concatenated.
|
||||
|
||||
Args:
|
||||
initial_params: Initial parameter values.
|
||||
opt_params: Optimized parameter values.
|
||||
_build_model: Callable to build a model from parameters and spec.
|
||||
trajectories: List of trajectories to render.
|
||||
model_spec: MuJoCo model specification.
|
||||
output_filepath: Path to save the output video.
|
||||
render_initial: Whether to render with initial parameters.
|
||||
render_nominal: Whether to render with nominal parameters.
|
||||
render_opt: Whether to render with optimized parameters.
|
||||
height: Frame height in pixels.
|
||||
width: Frame width in pixels.
|
||||
fovy: Vertical field of view in degrees.
|
||||
camera: Camera index or name.
|
||||
fps: Frames per second.
|
||||
|
||||
Returns:
|
||||
Path to the saved video file.
|
||||
"""
|
||||
import imageio
|
||||
|
||||
all_frames = []
|
||||
for traj in trajectories:
|
||||
# Build models for this trajectory
|
||||
models = []
|
||||
datas = []
|
||||
|
||||
nominal_params = initial_params.copy()
|
||||
nominal_params.reset()
|
||||
|
||||
# initial
|
||||
if render_initial:
|
||||
initial_spec = model_spec.copy()
|
||||
initial_spec = model_modifier.apply_param_modifiers_spec(
|
||||
initial_params, initial_spec
|
||||
)
|
||||
spec_apply(initial_spec, ["rgba"], [[1, 0, 0, 0.5]])
|
||||
initial_model = initial_spec.compile()
|
||||
initial_data = mujoco.MjData(initial_model)
|
||||
models.append(initial_model)
|
||||
datas.append(initial_data)
|
||||
|
||||
# nominal
|
||||
if render_nominal:
|
||||
nominal_spec = model_spec.copy()
|
||||
nominal_spec = model_modifier.apply_param_modifiers_spec(
|
||||
nominal_params, nominal_spec
|
||||
)
|
||||
spec_apply(nominal_spec, ["rgba"], [[0, 1, 0, 0.4]])
|
||||
nominal_model = nominal_spec.compile()
|
||||
nominal_data = mujoco.MjData(nominal_model)
|
||||
models.append(nominal_model)
|
||||
datas.append(nominal_data)
|
||||
|
||||
# pred
|
||||
if render_opt:
|
||||
pred_spec = model_spec.copy()
|
||||
pred_spec = model_modifier.apply_param_modifiers_spec(
|
||||
opt_params, pred_spec
|
||||
)
|
||||
spec_apply(pred_spec, ["rgba"], [[0, 0, 1, 1.0]])
|
||||
pred_model = pred_spec.compile()
|
||||
pred_data = mujoco.MjData(pred_model)
|
||||
models.append(pred_model)
|
||||
datas.append(pred_data)
|
||||
|
||||
control_ts = traj.control.resample(target_dt=models[0].opt.timestep)
|
||||
state, _ = mujoco.rollout.rollout(
|
||||
models, datas, traj.initial_state, control_ts.data
|
||||
)
|
||||
models[0].vis.global_.fovy = fovy
|
||||
models[0].vis.global_.offwidth = width
|
||||
models[0].vis.global_.offheight = height
|
||||
frames = render_rollout(
|
||||
models,
|
||||
datas[0],
|
||||
state,
|
||||
framerate=fps,
|
||||
height=height,
|
||||
width=width,
|
||||
camera=camera,
|
||||
)
|
||||
all_frames.extend(list(frames))
|
||||
|
||||
output_filepath_str = str(output_filepath)
|
||||
writer = imageio.get_writer(output_filepath_str, fps=fps, quality=8)
|
||||
for frame in all_frames:
|
||||
writer.append_data(frame)
|
||||
writer.close()
|
||||
|
||||
return pathlib.Path(output_filepath_str)
|
||||
|
||||
|
||||
class VideoPlayer(ReportSection):
|
||||
"""A report section to embed and display a video file."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
title: str,
|
||||
video_filepath: pathlib.Path,
|
||||
anchor: str = "",
|
||||
width: int | str = 800,
|
||||
height: int | None = 450,
|
||||
autoplay: bool = False,
|
||||
controls: bool = True,
|
||||
muted: bool = False,
|
||||
loop: bool = True,
|
||||
caption: str = "<b>Legend:</b> <span class='color-initial'>Initial</span>, <span class='color-nominal'>Nominal</span>, <span class='color-optimized'>Optimized</span>",
|
||||
collapsible: bool = True,
|
||||
):
|
||||
super().__init__(collapsible=collapsible)
|
||||
self._title = title
|
||||
self._anchor = anchor
|
||||
self._video_filepath = video_filepath
|
||||
self._width = width
|
||||
self._height = height
|
||||
self._autoplay = autoplay
|
||||
self._controls = controls
|
||||
self._muted = muted
|
||||
self._loop = loop
|
||||
self._caption = caption
|
||||
|
||||
@property
|
||||
def title(self) -> str:
|
||||
return self._title
|
||||
|
||||
@property
|
||||
def anchor(self) -> str:
|
||||
return self._anchor
|
||||
|
||||
@property
|
||||
def template_filename(self) -> str:
|
||||
"""Tells the builder to look for 'video.html'."""
|
||||
return "video.html"
|
||||
|
||||
def header_includes(self) -> set[str]:
|
||||
return set()
|
||||
|
||||
def get_context(self) -> dict[str, Any]:
|
||||
"""Returns the data needed to render the video player in the template."""
|
||||
return {
|
||||
"title": self._title,
|
||||
"video_filepath": self._video_filepath.name,
|
||||
"width": self._width,
|
||||
"height": self._height,
|
||||
"autoplay": "autoplay" if self._autoplay else "",
|
||||
"controls": "controls" if self._controls else "",
|
||||
"muted": "muted" if self._muted else "",
|
||||
"loop": "loop" if self._loop else "",
|
||||
"caption": self._caption,
|
||||
}
|
||||
Reference in New Issue
Block a user