MuJoCo Warp API documentation.

PiperOrigin-RevId: 810016913
Change-Id: Id4c309919068f7329cd35a2e0b63dc97c42343c3

doc/_static/onthispage_mjwarp.js

@@ -0,0 +1,80 @@
+/**
+ * This script hides Sphinx auto-generated attribute links from the
+ * "On this page" sidebar.
+ * Adapted for the Furo theme.
+ * Uses a MutationObserver to wait for Furo's JavaScript to create the
+ * sidebar before attempting to modify it.
+ */
+function hideAttributesFromTOC() {
+  // Furo's "On this page" sidebar container has the class .toc-tree
+  const tocContainerSelector = '.toc-tree';
+
+  // Step 1: Find all attribute definitions and collect their IDs.
+  const attributeDefs = document.querySelectorAll('dl.py.attribute');
+  if (attributeDefs.length === 0) {
+    return false;  // No attributes on this page, nothing to do.
+  }
+  const attributeIds = new Set();
+  attributeDefs.forEach(def => {
+    const term = def.querySelector('dt');
+    if (term && term.id) {
+      attributeIds.add(term.id);
+    }
+  });
+
+  if (attributeIds.size === 0) {
+    return false;  // No attribute IDs found.
+  }
+
+  // Step 2: Find the sidebar container.
+  const tocContainer = document.querySelector(tocContainerSelector);
+  if (!tocContainer) {
+    // Container not found yet. The observer will try again.
+    return false;
+  }
+
+  // Step 3: Get all links within that sidebar.
+  const tocLinks = tocContainer.querySelectorAll('a');
+  if (tocLinks.length === 0) {
+    return false;  // Container found, but it's empty. Let's wait.
+  }
+
+  // Step 4: Iterate through the links and hide the ones that match.
+  let hiddenCount = 0;
+  tocLinks.forEach(link => {
+    const href = link.getAttribute('href');
+    if (href && href.startsWith('#')) {
+      const linkId = href.substring(1);
+      if (attributeIds.has(linkId)) {
+        // In Furo, the link is inside a list item (<li>) which we need to hide.
+        const listItem = link.closest('li');
+        if (listItem && listItem.style.display !== 'none') {
+          console.log(`Hiding sidebar link for: #${linkId}`);
+          listItem.style.display = 'none';
+          hiddenCount++;
+        }
+      }
+    }
+  });
+
+  // If we successfully hid the links, we can stop observing.
+  if (hiddenCount > 0) {
+    return true;  // Signal success
+  }
+
+  // If we found the container but didn't hide anything, maybe it's not fully
+  // rendered. Let the observer run a few more times. A better approach might be
+  // needed if this fails, but for most cases, this is sufficient.
+  return false;
+}
+
+// Use a MutationObserver to wait for the page to be dynamically built.
+const observer = new MutationObserver((mutations, obs) => {
+  // We only need to run our function once successfully.
+  if (hideAttributesFromTOC()) {
+    obs.disconnect();  // Stop observing once the task is done.
+  }
+});
+
+// Start observing the entire document body for added/removed nodes.
+observer.observe(document.body, {childList: true, subtree: true});

doc/conf.py

@@ -25,6 +25,9 @@ import sys
 sys.path.insert(0, os.path.abspath('../'))
 sys.path.append(os.path.abspath('ext'))
 
+# MuJoCo Warp
+sys.path.append(os.path.abspath('../../mujoco_warp'))
+
 from sphinxcontrib import katex  # pylint: disable=g-import-not-at-top
 from sphinxcontrib import youtube  # pylint: disable=g-import-not-at-top,unused-import
 
@@ -42,6 +45,9 @@ master_doc = 'index'
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.viewcode',
     'sphinxcontrib.bibtex',
     'sphinxcontrib.katex',
     'sphinxcontrib.youtube',
@@ -54,6 +60,13 @@ extensions = [
     'mujoco_include',
 ]
 
+# MuJoCo Warp documentation
+napoleon_google_docstring = True
+autodoc_class_signature = 'separated'
+add_module_names = False
+toc_object_entries_show_parents = 'hide'
+default_role = 'literal'
+
 # GitHub-related options
 github_username = 'google-deepmind'
 github_repository = 'mujoco'
@@ -163,9 +176,11 @@ html_static_path = [
 ]
 html_css_files = [
     'theme_overrides.css',
+    'theme_overrides_mjwarp.css',
 ]
 html_js_files = [
     'linenumbers.js',
+    'onthispage_mjwarp.js',
 ]
 
 favicons = [

doc/css/theme_overrides_mjwarp.css

@@ -0,0 +1,145 @@
+/*
+ * ===================================================================
+ * Unified Documentation Layout
+ *
+ * Creates a consistent two-column layout for class attributes and
+ * function/method parameters. Includes aggressive overrides to resolve
+ * conflicts with `theme_overrides.css`.
+ * ===================================================================
+ */
+
+
+/*
+ * Part 1: Class-Wide Attribute Alignment Grid
+ * -------------------------------------------
+ * This section creates a single grid for all attributes within a class,
+ * ensuring their names and descriptions are perfectly aligned.
+ */
+
+/* The container for all class members becomes the main grid. */
+dl.py.class > dd {
+  display: grid;
+  grid-template-columns: auto 1fr;
+  column-gap: 0.5ch;
+  align-items: baseline;
+}
+
+/* This makes each attribute's <dl> transparent to the grid, promoting its
+ * <dt> and <dd> children to become direct grid items. This is the key
+ * to achieving the unified alignment.
+ */
+dl.py.attribute {
+  display: contents;
+}
+
+/* Other class members (main docstring, methods) must span both columns
+ * to avoid breaking the attribute grid layout.
+ */
+dl.py.class > dd > p,
+dl.py.class > dd > dl:not(.py.attribute) {
+  grid-column: 1 / -1;
+}
+dl.py.class > dd > p { margin-bottom: 1rem; }
+dl.py.class > dd > dl:not(.py.attribute) { margin-top: 1.5rem; }
+
+
+/*
+ * Part 2: Parameter & Return Value Layout
+ * ---------------------------------------
+ * This section applies a similar two-column grid to the parameters
+ * inside functions and methods for a consistent appearance.
+ */
+
+/* Target the inner list of parameters to avoid grabbing the main heading. */
+dl.py.function .field-list > dd > dl,
+dl.py.method .field-list > dd > dl {
+  display: grid;
+  grid-template-columns: auto 1fr;
+  column-gap: 0.5ch;
+  align-items: baseline;
+  margin-bottom: 0.4rem;
+}
+
+
+/*
+ * Part 3: Common Styling for Grid Content
+ * ---------------------------------------
+ * These rules reset default styling for elements within our new grids.
+ */
+
+/* Reset margins on all description blocks (<dd>) within our custom grids. */
+dl.py.attribute > dd,
+dl.py.function .field-list > dd > dl > dd,
+dl.py.method .field-list > dd > dl > dd {
+  margin: 0;
+  padding: 0;
+}
+/* Add vertical spacing between attribute rows. */
+dl.py.attribute > dd {
+  padding-bottom: 0.4rem;
+}
+
+/* Force the main description paragraph of an attribute to stay on one line. */
+dl.py.attribute > dd > p {
+  display: inline;
+  margin: 0;
+}
+
+
+/*
+ * Part 4: Specific Styling for Attribute Type Highlight
+ * -----------------------------------------------------
+ * This styles the highlighted box for an attribute's type.
+ */
+
+/* The container for the type information. */
+dl.py.attribute > dd > dl.field-list {
+  display: flex;
+  align-items: center;
+  margin-top: 0.4rem;
+}
+/* Reset its children. */
+dl.py.attribute > dd > dl.field-list > dt,
+dl.py.attribute > dd > dl.field-list > dd {
+    margin: 0;
+    padding: 0;
+}
+/* The highlight box itself - reverted to a simple implementation. */
+dl.py.attribute > dd > dl.field-list > dd {
+  font-family: monospace;
+  font-size: 90%;
+}
+
+
+/*
+ * Part 5: Aggressive Theme Overrides
+ * ----------------------------------
+ * These rules use `!important` to forcefully win conflicts with
+ * `theme_overrides.css` and ensure our layout is not broken.
+ */
+
+/* CONFLICT FIX: Forcefully remove left margins that break our grids. */
+dl.py.attribute > dd,
+dl.py.attribute > dd > dl.field-list > dd,
+dl.py.function .field-list > dd > dl > dd,
+dl.py.method .field-list > dd > dl > dd {
+  margin-left: 0 !important;
+}
+
+/* CONFLICT FIX: Remove the double colon rendered by the theme. */
+.field-list > dt::after {
+  content: "" !important;
+}
+
+/* CONFLICT FIX: Hide unwanted labels ("TYPE:", "Return type:"). */
+dl.py.attribute > dd > dl.field-list > dt,
+dl.py.function .field-list > dt:last-of-type,
+dl.py.method .field-list > dt:last-of-type {
+    display: none !important;
+}
+
+/* Remove the underline from all links. */
+a {
+  text-decoration: none !important;
+}
+

doc/mjwarp/api.rst

@@ -0,0 +1,24 @@
+MuJoCo Warp API
+===============
+
+.. automodule:: mujoco_warp
+  :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__,
+

doc/mjwarp/index.rst

@@ -4,6 +4,11 @@
 MuJoCo Warp (MJWarp)
 ====================
 
+.. toctree::
+    :hidden:
+
+    API <api.rst>
+
 MuJoCo Warp (MJWarp) is an implementation of MuJoCo written in `Warp <https://nvidia.github.io/warp/>`__ and optimized
 for `Nvidia <https://nvidia.com>`__ GPUs. MJWarp lives in the
 `google-deepmind/mujoco_warp <https://github.com/google-deepmind/mujoco_warp>`__ GitHub repository and is currently in