Matija Kecman
8079ab3d4d
ViewerApp is simplified, interacting with the endpoint and handlers through the Viewer instance. This change makes it possible to write simulation viewers without using the UI/UX provided by ViewerApp (which should be renamed StudioApp) PiperOrigin-RevId: 943642660 Change-Id: Id8eb3d4e6a27ceda93833367fec4699fa6bc838e
50 lines
1.8 KiB
Python
50 lines
1.8 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
|
|
#
|
|
# https://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.
|
|
"""Utility functions for Studio viewer implementations."""
|
|
|
|
import mujoco
|
|
from mujoco.experimental.studio import messages
|
|
from mujoco.experimental.studio import viewer_protocol
|
|
import numpy as np
|
|
|
|
|
|
def apply_perturb(
|
|
viewer: viewer_protocol.Viewer,
|
|
model: mujoco.MjModel | None,
|
|
data: mujoco.MjData | None,
|
|
is_paused: bool = False,
|
|
) -> None:
|
|
"""Apply perturbation to the model and send updated forces to the sim thread."""
|
|
if model is None or data is None:
|
|
return
|
|
perturb = viewer.perturb
|
|
if not is_paused:
|
|
sig = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED)
|
|
size = mujoco.mj_stateSize(model, sig)
|
|
zero_state = np.zeros(size, np.float64)
|
|
mujoco.mj_setState(model, data, zero_state, sig)
|
|
mujoco.mjv_applyPerturbPose(model, data, perturb, 0)
|
|
mujoco.mjv_applyPerturbForce(model, data, perturb)
|
|
else:
|
|
mujoco.mjv_applyPerturbPose(model, data, perturb, 1)
|
|
|
|
xfrc_sig: int = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED)
|
|
xfrc_size: int = mujoco.mj_stateSize(model, xfrc_sig)
|
|
xfrc_state: np.ndarray = np.zeros(xfrc_size, np.float64)
|
|
|
|
mujoco.mj_getState(model, data, xfrc_state, xfrc_sig)
|
|
viewer.send_to_sim(
|
|
messages.PerturbEvent(state=xfrc_state, state_sig=xfrc_sig)
|
|
)
|