Add documentation for MJX API

PiperOrigin-RevId: 860108133
Change-Id: I85469efcfe49d320ef8b89c9d174f33499a22b31

doc/conf.py

@@ -78,6 +78,33 @@ add_module_names = False
 toc_object_entries_show_parents = 'hide'
 default_role = 'literal'
 
+
+# Suppress warnings from docstrings with RST formatting issues
+def setup(app):
+  import logging
+
+  class SphinxWarningFilter(logging.Filter):
+
+    def filter(self, record):
+      msg = record.getMessage()
+      # e.g. qpos0: qpos values at default pose (*, nq)
+      if 'Inline emphasis start-string' in msg:
+        return False
+      # e.g. see `name` for details
+      if 'Inline interpreted text' in msg:
+        return False
+      # e.g. multi-line definition without trailing blank line
+      if 'Definition list ends without a blank line' in msg:
+        return False
+      # e.g. mjsactuator_ reference in C++ docstrings
+      if 'Unknown target name' in msg:
+        return False
+      return True
+
+  for handler in logging.getLogger('sphinx').handlers:
+    handler.addFilter(SphinxWarningFilter())
+
+
 # GitHub-related options
 github_username = 'google-deepmind'
 github_repository = 'mujoco'

doc/mjx.rst

@@ -4,6 +4,11 @@
 MuJoCo XLA (MJX)
 ================
 
+.. toctree::
+    :hidden:
+
+    API <mjx_api.rst>
+
 Starting with version 3.0.0, MuJoCo includes MuJoCo XLA (MJX) under the
 `mjx <https://github.com/google-deepmind/mujoco/tree/main/mjx>`__ directory.  MJX allows MuJoCo to run on compute
 hardware supported by the `XLA <https://www.tensorflow.org/xla>`__ compiler via the

doc/mjx_api.rst

@@ -0,0 +1,44 @@
+MJX API
+=======
+
+.. automodule:: mujoco.mjx
+  :members:
+  :imported-members:
+  :special-members: False
+  :private-members: False
+  :exclude-members:
+      __init__,
+      __format__,
+      __new__,
+      __eq__,
+      __hash__,
+      __repr__,
+      __weakref__,
+      __or__,
+      __and__,
+      __xor__,
+      __ror__,
+      __rand__,
+      __rxor__,
+      __invert__,
+      __delattr__,
+      __setattr__,
+      __getattr__,
+      __getattribute__,
+      __reduce__,
+      __reduce_ex__,
+      __sizeof__,
+      __str__,
+      __module__,
+      __doc__,
+      __dict__,
+      __class__,
+      __le__,
+      __lt__,
+      __ge__,
+      __gt__,
+      __ne__,
+      __bool__,
+      __init_subclass__,
+      replace,
+