Levi Burner
115 lines
3.6 KiB
Python
115 lines
3.6 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.
|
|
# ==============================================================================
|
|
|
|
"""Plotting utilities."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import mujoco
|
|
import numpy as np
|
|
|
|
|
|
def render_rollout(
|
|
model: mujoco.MjModel | Sequence[mujoco.MjModel],
|
|
data: mujoco.MjData,
|
|
state: np.ndarray,
|
|
framerate: int,
|
|
camera: str | int = -1,
|
|
width: int = 640,
|
|
height: int = 480,
|
|
light_pos: Sequence[float] | None = None,
|
|
) -> list[np.ndarray]:
|
|
"""Renders a rollout or batch of rollouts.
|
|
|
|
Args:
|
|
model: Single model or list of models (one per batch).
|
|
data: MjData scratch object.
|
|
state: State array of shape (nbatch, nsteps, nstate).
|
|
framerate: Frames per second to render.
|
|
camera: Camera name or ID.
|
|
width: Image width.
|
|
height: Image height.
|
|
light_pos: Optional light position [x, y, z] to add a spotlight.
|
|
|
|
Returns:
|
|
List of rendered frames (numpy arrays).
|
|
"""
|
|
nbatch = state.shape[0]
|
|
|
|
if isinstance(model, mujoco.MjModel):
|
|
models_list = [model] * nbatch
|
|
else:
|
|
models_list = list(model)
|
|
if len(models_list) == 1:
|
|
models_list = models_list * nbatch
|
|
else:
|
|
assert len(models_list) == nbatch
|
|
|
|
# Visual options
|
|
vopt = mujoco.MjvOption()
|
|
vopt.geomgroup[3] = 1 # Show visualization geoms
|
|
|
|
pert = mujoco.MjvPerturb()
|
|
catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC.value
|
|
|
|
# Simulate and render.
|
|
frames = []
|
|
|
|
with mujoco.Renderer(models_list[0], height=height, width=width) as renderer:
|
|
for i in range(state.shape[1]):
|
|
# Check if we should capture this frame based on framerate
|
|
if len(frames) < i * models_list[0].opt.timestep * framerate:
|
|
for j in range(state.shape[0]):
|
|
# Set state
|
|
mujoco.mj_setState(
|
|
models_list[j],
|
|
data,
|
|
state[j, i, :],
|
|
mujoco.mjtState.mjSTATE_FULLPHYSICS.value,
|
|
)
|
|
mujoco.mj_forward(models_list[j], data)
|
|
|
|
# Use first model to make the scene, add subsequent models
|
|
if j == 0:
|
|
renderer.update_scene(data, camera, scene_option=vopt)
|
|
else:
|
|
mujoco.mjv_addGeoms(
|
|
models_list[j], data, vopt, pert, catmask, renderer.scene
|
|
)
|
|
|
|
# Add light, if requested
|
|
if light_pos is not None:
|
|
if renderer.scene.nlight < 100: # check limit
|
|
light = renderer.scene.lights[renderer.scene.nlight]
|
|
light.ambient = [0, 0, 0]
|
|
light.attenuation = [1, 0, 0]
|
|
light.castshadow = 1
|
|
light.cutoff = 45
|
|
light.diffuse = [0.8, 0.8, 0.8]
|
|
light.dir = [0, 0, -1]
|
|
light.type = mujoco.mjtLightType.mjLIGHT_SPOT
|
|
light.exponent = 10
|
|
light.headlight = 0
|
|
light.specular = [0.3, 0.3, 0.3]
|
|
light.pos = light_pos
|
|
renderer.scene.nlight += 1
|
|
|
|
# Render and add the frame.
|
|
pixels = renderer.render()
|
|
frames.append(pixels)
|
|
return frames
|