Copy mjwarp viewer s.t. MJX users can debug models effectively.
PiperOrigin-RevId: 792258628 Change-Id: I3b2b5d10414934866d5c647f30fe88e9cd1a8756
This commit is contained in:
Baruch Tabanpour
committed by
Copybara-Service
Copybara-Service
parent
9e5fdd6983
commit
9b9b7030e7
+227
@@ -0,0 +1,227 @@
|
|||||||
|
# Copyright 2025 The Newton Developers
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
# ==============================================================================
|
||||||
|
|
||||||
|
"""mjwarp-viewer: load and simulate an MJCF with MuJoCo Warp.
|
||||||
|
|
||||||
|
Usage: mjwarp-viewer <mjcf XML path> [flags]
|
||||||
|
|
||||||
|
Example:
|
||||||
|
mjwarp-viewer benchmark/humanoid/humanoid.xml -o "opt.solver=cg"
|
||||||
|
"""
|
||||||
|
|
||||||
|
import ast
|
||||||
|
import enum
|
||||||
|
import logging
|
||||||
|
import pickle
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
from typing import Sequence, Union
|
||||||
|
|
||||||
|
import mujoco
|
||||||
|
import mujoco.viewer
|
||||||
|
import numpy as np
|
||||||
|
import warp as wp
|
||||||
|
from absl import app
|
||||||
|
from absl import flags
|
||||||
|
from etils import epath
|
||||||
|
|
||||||
|
import mujoco_warp as mjw
|
||||||
|
|
||||||
|
|
||||||
|
class EngineOptions(enum.IntEnum):
|
||||||
|
WARP = 0
|
||||||
|
C = 1
|
||||||
|
|
||||||
|
|
||||||
|
_CLEAR_KERNEL_CACHE = flags.DEFINE_bool("clear_kernel_cache", False, "Clear kernel cache (to calculate full JIT time)")
|
||||||
|
_ENGINE = flags.DEFINE_enum_class("engine", EngineOptions.WARP, EngineOptions, "Simulation engine")
|
||||||
|
_NCONMAX = flags.DEFINE_integer("nconmax", None, "Maximum number of contacts.")
|
||||||
|
_NJMAX = flags.DEFINE_integer("njmax", None, "Maximum number of constraints per world.")
|
||||||
|
_OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation: foo.bar = baz)", short_name="o")
|
||||||
|
_DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
|
||||||
|
|
||||||
|
|
||||||
|
_VIEWER_GLOBAL_STATE = {"running": True, "step_once": False}
|
||||||
|
|
||||||
|
|
||||||
|
def key_callback(key: int) -> None:
|
||||||
|
if key == 32: # Space bar
|
||||||
|
_VIEWER_GLOBAL_STATE["running"] = not _VIEWER_GLOBAL_STATE["running"]
|
||||||
|
logging.info("RUNNING = %s", _VIEWER_GLOBAL_STATE["running"])
|
||||||
|
elif key == 46: # period
|
||||||
|
_VIEWER_GLOBAL_STATE["step_once"] = True
|
||||||
|
|
||||||
|
|
||||||
|
def _load_model(path: epath.Path) -> mujoco.MjModel:
|
||||||
|
if not path.exists():
|
||||||
|
resource_path = epath.resource_path("mujoco_warp") / path
|
||||||
|
if not resource_path.exists():
|
||||||
|
raise FileNotFoundError(f"file not found: {path}\nalso tried: {resource_path}")
|
||||||
|
path = resource_path
|
||||||
|
|
||||||
|
print(f"Loading model from: {path}...")
|
||||||
|
if path.suffix == ".mjb":
|
||||||
|
return mujoco.MjModel.from_binary_path(path.as_posix())
|
||||||
|
|
||||||
|
spec = mujoco.MjSpec.from_file(path.as_posix())
|
||||||
|
# check if the file has any mujoco.sdf test plugins
|
||||||
|
if any(p.plugin_name.startswith("mujoco.sdf") for p in spec.plugins):
|
||||||
|
from mujoco.mjx.third_party.mujoco_warp.test_data.collision_sdf.utils import register_sdf_plugins as register_sdf_plugins
|
||||||
|
|
||||||
|
register_sdf_plugins(mjw)
|
||||||
|
return spec.compile()
|
||||||
|
|
||||||
|
|
||||||
|
def _override(model: Union[mjw.Model, mujoco.MjModel]):
|
||||||
|
enum_fields = {
|
||||||
|
"opt.integrator": mjw.IntegratorType,
|
||||||
|
"opt.cone": mjw.ConeType,
|
||||||
|
"opt.solver": mjw.SolverType,
|
||||||
|
"opt.broadphase": mjw.BroadphaseType,
|
||||||
|
"opt.broadphase_filter": mjw.BroadphaseFilter,
|
||||||
|
}
|
||||||
|
for override in _OVERRIDE.value:
|
||||||
|
if "=" not in override:
|
||||||
|
raise app.UsageError(f"Invalid override format: {override}")
|
||||||
|
key, val = override.split("=", 1)
|
||||||
|
key, val = key.strip(), val.strip()
|
||||||
|
|
||||||
|
if key in enum_fields:
|
||||||
|
try:
|
||||||
|
val = str(enum_fields[key][val.upper()])
|
||||||
|
except KeyError:
|
||||||
|
raise app.UsageError(f"Unrecognized enum value: {val}")
|
||||||
|
|
||||||
|
obj, attrs = model, key.split(".")
|
||||||
|
for i, attr in enumerate(attrs):
|
||||||
|
if not hasattr(obj, attr):
|
||||||
|
raise app.UsageError(f"Unrecognized model field: {key}")
|
||||||
|
if i < len(attrs) - 1:
|
||||||
|
obj = getattr(obj, attr)
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
val = type(getattr(obj, attr))(ast.literal_eval(val))
|
||||||
|
except (SyntaxError, ValueError):
|
||||||
|
raise app.UsageError(f"Unrecognized value for field: {key}")
|
||||||
|
|
||||||
|
setattr(obj, attr, val)
|
||||||
|
|
||||||
|
|
||||||
|
def _compile_step(m, d):
|
||||||
|
mjw.step(m, d)
|
||||||
|
# double warmup to work around issues with compilation during graph capture:
|
||||||
|
mjw.step(m, d)
|
||||||
|
# capture the whole step function as a CUDA graph
|
||||||
|
with wp.ScopedCapture() as capture:
|
||||||
|
mjw.step(m, d)
|
||||||
|
return capture.graph
|
||||||
|
|
||||||
|
|
||||||
|
def _main(argv: Sequence[str]) -> None:
|
||||||
|
"""Runs viewer app."""
|
||||||
|
if len(argv) < 2:
|
||||||
|
raise app.UsageError("Missing required input: mjcf path.")
|
||||||
|
elif len(argv) > 2:
|
||||||
|
raise app.UsageError("Too many command-line arguments.")
|
||||||
|
|
||||||
|
mjm = _load_model(epath.Path(argv[1]))
|
||||||
|
mjd = mujoco.MjData(mjm)
|
||||||
|
mujoco.mj_forward(mjm, mjd)
|
||||||
|
|
||||||
|
if _ENGINE.value == EngineOptions.C:
|
||||||
|
_override(mjm)
|
||||||
|
print(
|
||||||
|
f" nbody: {mjm.nbody} nv: {mjm.nv} ngeom: {mjm.ngeom} nu: {mjm.nu}\n"
|
||||||
|
f" solver: {mujoco.mjtSolver(mjm.opt.solver).name} cone: {mujoco.mjtCone(mjm.opt.cone).name}"
|
||||||
|
f" iterations: {mjm.opt.iterations} ls_iterations: {mjm.opt.ls_iterations}\n"
|
||||||
|
f" integrator: {mujoco.mjtIntegrator(mjm.opt.integrator).name}\n"
|
||||||
|
)
|
||||||
|
print(f"MuJoCo C simulating with dt = {mjm.opt.timestep:.3f}...")
|
||||||
|
else:
|
||||||
|
wp.config.quiet = flags.FLAGS["verbosity"].value < 1
|
||||||
|
wp.init()
|
||||||
|
if _CLEAR_KERNEL_CACHE.value:
|
||||||
|
wp.clear_kernel_cache()
|
||||||
|
|
||||||
|
with wp.ScopedDevice(_DEVICE.value):
|
||||||
|
m = mjw.put_model(mjm)
|
||||||
|
_override(m)
|
||||||
|
mjm_hash = pickle.dumps(mjm)
|
||||||
|
broadphase, filter = mjw.BroadphaseType(m.opt.broadphase).name, mjw.BroadphaseFilter(m.opt.broadphase_filter).name
|
||||||
|
solver, cone = mjw.SolverType(m.opt.solver).name, mjw.ConeType(m.opt.cone).name
|
||||||
|
integrator = mjw.IntegratorType(m.opt.integrator).name
|
||||||
|
iterations, ls_iterations, ls_parallel = m.opt.iterations, m.opt.ls_iterations, m.opt.ls_parallel
|
||||||
|
print(
|
||||||
|
f" nbody: {m.nbody} nv: {m.nv} ngeom: {m.ngeom} nu: {m.nu} is_sparse: {m.opt.is_sparse}\n"
|
||||||
|
f" broadphase: {broadphase} broadphase_filter: {filter}\n"
|
||||||
|
f" solver: {solver} cone: {cone} iterations: {iterations} ls_iterations: {ls_iterations} ls_parallel: {ls_parallel}\n"
|
||||||
|
f" integrator: {integrator} graph_conditional: {m.opt.graph_conditional}"
|
||||||
|
)
|
||||||
|
d = mjw.put_data(mjm, mjd, nconmax=_NCONMAX.value, njmax=_NJMAX.value)
|
||||||
|
print(f"Data\n nworld: {d.nworld} nconmax: {d.nconmax} njmax: {d.njmax}\n")
|
||||||
|
print("Compiling physics step...", end="")
|
||||||
|
start = time.time()
|
||||||
|
graph = _compile_step(m, d)
|
||||||
|
elapsed = time.time() - start
|
||||||
|
print(f"done ({elapsed:0.2}s).")
|
||||||
|
print(f"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...")
|
||||||
|
|
||||||
|
with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer:
|
||||||
|
while True:
|
||||||
|
start = time.time()
|
||||||
|
|
||||||
|
if _ENGINE.value == EngineOptions.C:
|
||||||
|
mujoco.mj_step(mjm, mjd)
|
||||||
|
else: # mjwarp
|
||||||
|
wp.copy(d.ctrl, wp.array([mjd.ctrl.astype(np.float32)]))
|
||||||
|
wp.copy(d.act, wp.array([mjd.act.astype(np.float32)]))
|
||||||
|
wp.copy(d.xfrc_applied, wp.array([mjd.xfrc_applied.astype(np.float32)]))
|
||||||
|
wp.copy(d.qpos, wp.array([mjd.qpos.astype(np.float32)]))
|
||||||
|
wp.copy(d.qvel, wp.array([mjd.qvel.astype(np.float32)]))
|
||||||
|
wp.copy(d.time, wp.array([mjd.time], dtype=wp.float32))
|
||||||
|
|
||||||
|
hash = pickle.dumps(mjm)
|
||||||
|
if hash != mjm_hash:
|
||||||
|
mjm_hash = hash
|
||||||
|
m = mjw.put_model(mjm)
|
||||||
|
graph = _compile_step(m, d)
|
||||||
|
|
||||||
|
if _VIEWER_GLOBAL_STATE["running"]:
|
||||||
|
wp.capture_launch(graph)
|
||||||
|
wp.synchronize()
|
||||||
|
elif _VIEWER_GLOBAL_STATE["step_once"]:
|
||||||
|
_VIEWER_GLOBAL_STATE["step_once"] = False
|
||||||
|
wp.capture_launch(graph)
|
||||||
|
wp.synchronize()
|
||||||
|
|
||||||
|
mjw.get_data_into(mjd, mjm, d)
|
||||||
|
|
||||||
|
viewer.sync()
|
||||||
|
|
||||||
|
elapsed = time.time() - start
|
||||||
|
if elapsed < mjm.opt.timestep:
|
||||||
|
time.sleep(mjm.opt.timestep - elapsed)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
# absl flags assumes __main__ is the main running module for printing usage documentation
|
||||||
|
# pyproject bin scripts break this assumption, so manually set argv and docstring
|
||||||
|
sys.argv[0] = "mujoco_warp.viewer"
|
||||||
|
sys.modules["__main__"].__doc__ = __doc__
|
||||||
|
app.run(_main)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user