Add mjviser web viewer backend (--viewer=viser).

Author: kevinzakka

mujoco_warp/viewer.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The Newton Developers
+# Copyright 2026 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.
@@ -62,6 +62,7 @@ class EngineOptions(enum.IntEnum):
 _KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.")
 _DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
 _REPLAY = flags.DEFINE_string("replay", None, "keyframe sequence to replay, keyframe name must prefix match")
+_VIEWER = flags.DEFINE_enum("viewer", "mujoco", ["mujoco", "viser"], "Viewer backend (mujoco native or mjviser web)")
 
 _VIEWER_GLOBAL_STATE = {"running": True, "step_once": False}
 
@@ -105,6 +106,47 @@ def _compile_step(m, d):
   return capture.graph
 
 
+def _make_warp_step_fn(m, d, graph, ctrls=None):
+  ctrlid = [0]
+
+  def step_fn(mjm, mjd):
+    if ctrls is not None and ctrlid[0] < len(ctrls):
+      mjd.ctrl[:] = ctrls[ctrlid[0]]
+      ctrlid[0] += 1
+
+    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))
+
+    if graph is None:
+      mjw.step(m, d)
+    else:
+      wp.capture_launch(graph)
+      wp.synchronize()
+
+    mjw.get_data_into(mjd, mjm, d)
+
+  return step_fn
+
+
+def _make_c_step_fn(ctrls=None):
+  if ctrls is None:
+    return None
+
+  ctrlid = [0]
+
+  def step_fn(mjm, mjd):
+    if ctrlid[0] < len(ctrls):
+      mjd.ctrl[:] = ctrls[ctrlid[0]]
+      ctrlid[0] += 1
+    mujoco.mj_step(mjm, mjd)
+
+  return step_fn
+
+
 def _main(argv: Sequence[str]) -> None:
   """Runs viewer app."""
   if len(argv) < 2:
@@ -169,45 +211,58 @@ def _main(argv: Sequence[str]) -> None:
     print(f"Data\n  nworld: {d.nworld} nconmax: {int(d.naconmax / d.nworld)} njmax: {d.njmax}\n")
     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:
-    opt = copy.copy(mjm.opt)
-
-    while True:
-      start = time.time()
-
-      if ctrls is not None and ctrlid < len(ctrls):
-        mjd.ctrl[:] = ctrls[ctrlid]
-        ctrlid += 1
-
-      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))
-        # if the user changed an option in the MuJoCo Simulate UI, go ahead and recompile the step
-        # TODO: update memory tied to option max iterations
-        if mjm.opt != opt:
-          opt = copy.copy(mjm.opt)
-          m = mjw.put_model(mjm)
-          graph = _compile_step(m, d) if wp.get_device().is_cuda else None
-        if _VIEWER_GLOBAL_STATE["running"] or _VIEWER_GLOBAL_STATE["step_once"]:
-          _VIEWER_GLOBAL_STATE["step_once"] = False
-          if graph is None:
-            mjw.step(m, d)
-          else:
-            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)
+  if _VIEWER.value == "viser":
+    try:
+      from mjviser import Viewer as MjViserViewer
+    except ImportError:
+      raise SystemExit("mjviser required for --viewer=viser: pip install mjviser")
+
+    if _ENGINE.value == EngineOptions.WARP:
+      step_fn = _make_warp_step_fn(m, d, graph, ctrls)
+    else:
+      step_fn = _make_c_step_fn(ctrls)
+
+    MjViserViewer(mjm, mjd, step_fn=step_fn).run()
+  else:
+    with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer:
+      opt = copy.copy(mjm.opt)
+
+      while True:
+        start = time.time()
+
+        if ctrls is not None and ctrlid < len(ctrls):
+          mjd.ctrl[:] = ctrls[ctrlid]
+          ctrlid += 1
+
+        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))
+          # if the user changed an option in the MuJoCo Simulate UI, go ahead and recompile the step
+          # TODO: update memory tied to option max iterations
+          if mjm.opt != opt:
+            opt = copy.copy(mjm.opt)
+            m = mjw.put_model(mjm)
+            graph = _compile_step(m, d) if wp.get_device().is_cuda else None
+          if _VIEWER_GLOBAL_STATE["running"] or _VIEWER_GLOBAL_STATE["step_once"]:
+            _VIEWER_GLOBAL_STATE["step_once"] = False
+            if graph is None:
+              mjw.step(m, d)
+            else:
+              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():