Add graphics modes list to viewer protocol

PiperOrigin-RevId: 929157105
Change-Id: I75586c37bc94666800f7822654225bad0717e51c

python/mujoco/experimental/studio/native_viewer.py

@@ -21,17 +21,18 @@ to use these classes.
 """
 
 import mujoco
-
 from mujoco.experimental.studio import native_viewer_cc as _viewer
 from mujoco.experimental.studio import ux
+from mujoco.experimental.studio import viewer_protocol
 
 
-class NativeViewer:
+class NativeViewer(viewer_protocol.Viewer):
   """Simulation-agnostic native viewer for MuJoCo models."""
 
   def __init__(
       self,
       model: mujoco.MjModel,
+      *,
       camera: mujoco.MjvCamera | None = None,
       vis_options: mujoco.MjvOption | None = None,
       perturb: mujoco.MjvPerturb | None = None,
@@ -39,7 +40,7 @@ class NativeViewer:
       title: str = '',
       width: int = 1200,
       height: int = 800,
-      gfx: str = '',
+      gfx: str | None = None,
   ) -> None:
     """Initializes the NativeViewer.
 
@@ -55,12 +56,12 @@ class NativeViewer:
       title: Title of the viewer window.
       width: Initial width of the viewer window.
       height: Initial height of the viewer window.
-      gfx: Graphics mode.
+      gfx: Graphics mode. If None, uses the platform default.
     """
     self.camera = camera or mujoco.MjvCamera()
     self.perturb = perturb or mujoco.MjvPerturb()
     self.vis_options = vis_options or mujoco.MjvOption()
-    self._viewer = _viewer.Viewer(title, width, height, gfx)
+    self._viewer = _viewer.Viewer(title, width, height, gfx or '')
     self._viewer.InitRenderer(model)
     # This class does not own the model but we need to know if the model being
     # rendered has changed, so we store the unique python object id here so we

python/mujoco/experimental/studio/sample/async.py

@@ -33,11 +33,14 @@ from mujoco.experimental.studio import native_viewer as _viewer
 from mujoco.experimental.studio import sim as _sim
 from mujoco.experimental.studio import studio_app
 from mujoco.experimental.studio import ux
+from mujoco.experimental.studio import viewer_protocol
 import numpy as np
 
 from mujoco.experimental.dear_imgui import dear_imgui as imgui
 
-_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
+_GFX = absl_flags.DEFINE_enum(
+    'gfx', None, viewer_protocol.GFX_MODES, 'Rendering graphics mode.'
+)
 _WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
 _HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
 

python/mujoco/experimental/studio/sample/implot.py

@@ -29,12 +29,15 @@ from absl import flags as absl_flags
 import mujoco
 from mujoco.experimental.studio import native_viewer as _viewer
 from mujoco.experimental.studio import studio_app
+from mujoco.experimental.studio import viewer_protocol
 import numpy as np
 
 from mujoco.experimental.dear_imgui import dear_imgui as imgui
 from mujoco.experimental.implot import implot
 
-_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
+_GFX = absl_flags.DEFINE_enum(
+    'gfx', None, viewer_protocol.GFX_MODES, 'Rendering graphics mode.'
+)
 _WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
 _HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
 

python/mujoco/experimental/studio/sample/render.py

@@ -21,11 +21,14 @@ from absl import flags
 import mujoco
 from mujoco.experimental.studio import parser
 from mujoco.experimental.studio import renderer
+from mujoco.experimental.studio import viewer_protocol
 from PIL import Image
 
 _MODEL = flags.DEFINE_string('model', '', 'Model file to load.')
 _OUTPUT = flags.DEFINE_string('output', '', 'Output file to save.')
-_GFX = flags.DEFINE_string('gfx', '', 'Renderer to use.')
+_GFX = flags.DEFINE_enum(
+    'gfx', None, viewer_protocol.GFX_MODES, 'Rendering graphics mode.'
+)
 _WIDTH = flags.DEFINE_integer('width', 320, 'Width of the output image.')
 _HEIGHT = flags.DEFINE_integer('height', 240, 'Height of the output image.')
 _STEPS = flags.DEFINE_integer('steps', 1, 'Number of steps before render.')
@@ -50,7 +53,7 @@ def main(argv):
     mujoco.mj_step(model, data)
 
   try:
-    r = renderer.Renderer(_GFX.value)
+    r = renderer.Renderer(_GFX.value or '')
     r.Init(model)
     pixels = r.Render(
         model, data, None, None, None, _WIDTH.value, _HEIGHT.value

python/mujoco/experimental/studio/studio.py

@@ -17,8 +17,11 @@ from absl import app as absl_app
 from absl import flags as absl_flags
 from mujoco.experimental.studio import native_viewer
 from mujoco.experimental.studio import studio_app
+from mujoco.experimental.studio import viewer_protocol
 
-_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
+_GFX = absl_flags.DEFINE_enum(
+    'gfx', None, viewer_protocol.GFX_MODES, 'Rendering graphics mode.'
+)
 _WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
 _HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
 

python/mujoco/experimental/studio/viewer_protocol.py

@@ -21,6 +21,17 @@ from typing import Protocol
 import mujoco
 from mujoco.experimental.studio import ux
 
+GFX_MODES = (
+    'classic',
+    'classic_headless',
+    'opengl',
+    'opengl_headless',
+    'opengl_software',
+    'vulkan',
+    'vulkan_software',
+    'webgl',
+)
+
 
 class Viewer(Protocol):
   """Structural interface for any viewer."""