Address diagram issues in the tutorials (#445)

Co-authored-by: casteryh <[email protected]>

Author: svekars

docs/source/_static/custom.css

@@ -1,3 +1,62 @@
+/* Center all Mermaid diagrams */
+.mermaid {
+    display: block;
+    margin-left: auto;
+    margin-right: auto;
+    text-align: center;
+}
+
+/* Center the pre element that contains mermaid diagrams */
+pre.mermaid {
+    display: flex;
+    justify-content: center;
+}
+
+/* Adjust Mermaid line colors based on theme */
+/* Light mode - darker lines for visibility on white background */
+html[data-theme="light"] .mermaid .edgePath .path,
+html[data-theme="light"] .mermaid .flowchart-link {
+    stroke: #555 !important;
+    stroke-width: 2px !important;
+}
+
+/* Light mode - darker arrow tips */
+html[data-theme="light"] .mermaid .arrowheadPath,
+html[data-theme="light"] .mermaid marker path {
+    fill: #555 !important;
+    stroke: #555 !important;
+}
+
+html[data-theme="dark"] .mermaid .arrowheadPath,
+html[data-theme="dark"] .mermaid marker path {
+    fill: #aaa !important;
+    stroke: #aaa !important;
+}
+
+/* Dark mode - lighter lines for visibility on dark background */
+html[data-theme="dark"] .mermaid .edgePath .path,
+html[data-theme="dark"] .mermaid .flowchart-link {
+    stroke: #aaa !important;
+    stroke-width: 2px !important;
+}
+
+/* Dark mode - lighter arrow tips */
+html[data-theme="dark"] .mermaid .arrowheadPath,
+html[data-theme="dark"] .mermaid marker path {
+    fill: #aaa !important;
+    stroke: #aaa !important;
+}
+
+/* Adjust edge labels background based on theme */
+html[data-theme="light"] .mermaid .edgeLabel {
+    background-color: #fff !important;
+}
+
+html[data-theme="dark"] .mermaid .edgeLabel {
+    background-color: #1a1a1a !important;
+    color: #fff !important;
+}
+
 /* Custom CSS for collapsible parameter lists */
 
 /* Hide parameters in signatures */

docs/source/_static/custom.js

@@ -1,3 +1,103 @@
+// Lightbox functionality for Mermaid diagrams
+document.addEventListener('DOMContentLoaded', function() {
+    // Create lightbox modal for Mermaid diagrams
+    const lightbox = document.createElement('div');
+    lightbox.id = 'mermaid-lightbox';
+    lightbox.style.cssText = `
+        display: none;
+        position: fixed;
+        z-index: 9999;
+        left: 0;
+        top: 0;
+        width: 100%;
+        height: 100%;
+        background-color: rgba(0,0,0,0.9);
+        cursor: zoom-out;
+    `;
+
+    const lightboxContent = document.createElement('div');
+    lightboxContent.style.cssText = `
+        position: absolute;
+        top: 50%;
+        left: 50%;
+        transform: translate(-50%, -50%);
+        max-width: 95%;
+        max-height: 95%;
+        overflow: auto;
+    `;
+
+    const closeBtn = document.createElement('span');
+    closeBtn.innerHTML = '×';
+    closeBtn.style.cssText = `
+        position: absolute;
+        top: 15px;
+        right: 35px;
+        color: #f1f1f1;
+        font-size: 40px;
+        font-weight: bold;
+        cursor: pointer;
+        z-index: 10000;
+    `;
+    closeBtn.onclick = function() {
+        lightbox.style.display = 'none';
+    };
+
+    lightbox.appendChild(closeBtn);
+    lightbox.appendChild(lightboxContent);
+    document.body.appendChild(lightbox);
+
+    // Click outside to close
+    lightbox.onclick = function(e) {
+        if (e.target === lightbox) {
+            lightbox.style.display = 'none';
+        }
+    };
+
+    // ESC key to close
+    document.addEventListener('keydown', function(e) {
+        if (e.key === 'Escape' && lightbox.style.display === 'block') {
+            lightbox.style.display = 'none';
+        }
+    });
+
+    // Make all Mermaid diagrams clickable
+    const mermaidDivs = document.querySelectorAll('.mermaid');
+    mermaidDivs.forEach(function(div) {
+        div.style.cursor = 'zoom-in';
+        div.title = 'Click to enlarge';
+
+        div.addEventListener('click', function() {
+            const clone = div.cloneNode(true);
+
+            // Style the cloned diagram to fill the lightbox
+            clone.style.cssText = `
+                cursor: default;
+                width: 90vw;
+                max-width: 1400px;
+                height: auto;
+                margin: auto;
+            `;
+
+            // Find and resize the SVG inside
+            const svg = clone.querySelector('svg');
+            if (svg) {
+                svg.style.cssText = `
+                    width: 100% !important;
+                    height: auto !important;
+                    max-width: none !important;
+                    max-height: 90vh !important;
+                `;
+                svg.removeAttribute('width');
+                svg.removeAttribute('height');
+            }
+
+            lightboxContent.innerHTML = '';
+            lightboxContent.appendChild(clone);
+            lightbox.style.display = 'block';
+        });
+    });
+});
+
 // Custom JavaScript to make long parameter lists in class signatures collapsible
 document.addEventListener('DOMContentLoaded', function() {
     console.log('Collapsible parameters script loaded');

docs/source/conf.py

@@ -178,6 +178,30 @@ def get_version_path():
 # Configure MyST parser to treat mermaid code blocks as mermaid directives
 myst_fence_as_directive = ["mermaid"]
 
+# Disable D3 zoom (we'll use lightbox instead)
+mermaid_d3_zoom = False
+
+# Global Mermaid theme configuration - applies to all diagrams
+mermaid_init_js = """
+import mermaid from 'https://cdn.jsdelivr.net/npm/[email protected]/dist/mermaid.esm.min.mjs';
+mermaid.initialize({
+    startOnLoad: false,
+    theme: 'base',
+    themeVariables: {
+        primaryColor: '#4CAF50',
+        primaryTextColor: '#000',
+        primaryBorderColor: '#fff',
+        lineColor: '#555',
+        secondaryColor: '#FF9800',
+        tertiaryColor: '#ffffde'
+    },
+    flowchart: {
+        curve: 'basis'
+    },
+    themeCSS: '.edgePath .path { stroke-width: 4px; stroke: #555; }'
+});
+"""
+
 autodoc_default_options = {
     "members": True,
     "undoc-members": True,

docs/source/tutorial_sources/zero-to-forge/1_RL_and_Forge_Fundamentals.md

@@ -11,8 +11,10 @@ graph TD
     subgraph Example["Math Tutoring RL Example"]
         Dataset["Dataset: math problems"]
         Policy["Policy: student AI"]
-        Reward["Reward Model: scores answers"]
-        Reference["Reference Model: baseline"]
+        Reward["Reward Model:
+        scores answers"]
+        Reference["Reference Model:
+        baseline"]
         ReplayBuffer["Replay Buffer: stores experiences"]
         Trainer["Trainer: improves student"]
     end
@@ -25,9 +27,11 @@ graph TD
     ReplayBuffer --> Trainer
     Trainer --> Policy
 
-    style Policy fill:#4CAF50
-    style Reward fill:#FF9800
-    style Trainer fill:#E91E63
+    style Policy fill:#4CAF50,stroke:#fff,stroke-width:2px
+    style Reward fill:#FF9800,stroke:#fff,stroke-width:2px
+    style Trainer fill:#E91E63,stroke:#fff,stroke-width:2px
+
+    linkStyle default stroke:#888,stroke-width:2px
 ```
 
 ### RL Components Defined (TorchForge Names)
@@ -76,7 +80,7 @@ Here's the key insight: **Each RL component becomes a TorchForge service**. The
 ```mermaid
 graph LR
     subgraph Concepts["RL Concepts"]
-        direction TB
+
         C1["Dataset"]
         C2["Policy"]
         C3["Reward Model"]
@@ -86,7 +90,7 @@ graph LR
     end
 
     subgraph Services["TorchForge Services (Real Classes)"]
-        direction TB
+
         S1["DatasetActor"]
         S2["Generator"]
         S3["RewardActor"]
@@ -173,11 +177,20 @@ Our simple RL loop above has complex requirements:
 
 ```mermaid
 graph LR
-    A["Policy: Student AI<br/>'What is 2+2?' → 'The answer is 4'"]
-    B["Reward: Teacher<br/>Scores answer: 0.95"]
-    C["Reference: Original Student<br/>Provides baseline comparison"]
-    D["Replay Buffer: Notebook<br/>Stores: question + answer + score"]
-    E["Trainer: Tutor<br/>Improves student using experiences"]
+    A["Policy: Student AI
+    'What is 2+2?' →
+    'The answer is 4'"]
+    B["Reward: Teacher
+    Scores answer: 0.95"]
+    C["Reference: Original Student
+    Provides baseline comparison"]
+    D["Replay Buffer: Notebook
+    Stores: question
+    + answer
+    + score"]
+    E["Trainer: Tutor
+    Improves student
+    using experiences"]
 
     A --> B
     A --> C

docs/source/tutorial_sources/zero-to-forge/2_Forge_Internals.md

@@ -12,22 +12,35 @@ When you call `await policy_service.generate(question)`, here's what actually ha
 
 ```mermaid
 graph TD
-    Call["Your Code:<br/>await policy_service.generate"]
+    Call["Your Code:
+    await policy_service
+    .generate.route"]
 
     subgraph ServiceLayer["Service Layer"]
-        Proxy["Service Proxy: Load balancing, Health checking"]
-        LB["Load Balancer: Replica selection, Circuit breaker"]
+        Proxy["Service Proxy:
+        Load balancing
+        Health checking"]
+        LB["Load Balancer:
+        Replica selection
+        Circuit breaker"]
     end
 
     subgraph Replicas["Replica Management"]
-        R1["Replica 1: GPU 0, Healthy"]
-        R2["Replica 2: GPU 1, Overloaded"]
-        R3["Replica 3: GPU 2, Failed"]
-        R4["Replica 4: GPU 3, Healthy"]
+        R1["Replica 1:
+        GPU 0, Healthy"]
+        R2["Replica 2:
+        GPU 1, Overloaded"]
+        R3["Replica 3:
+        GPU 2, Failed"]
+        R4["Replica 4:
+        GPU 3, Healthy"]
     end
 
     subgraph Compute["Actual Computation"]
-        Actor["Policy Actor: vLLM engine, Model weights, KV cache"]
+        Actor["Policy Actor:
+        vLLM engine,
+        Model weights,
+        KV cache"]
     end
 
     Call --> Proxy
@@ -126,13 +139,17 @@ responses = await policy.generate.route(prompt=prompt)
 ```mermaid
 graph LR
     subgraph Request["Your Request"]
-        Code["await service.method.ADVERB()"]
+        Code["await service
+        .method.ADVERB()"]
     end
 
     subgraph Patterns["Communication Patterns"]
-        Route[".route()<br/>→ One healthy replica"]
-        CallOne[".call_one()<br/>→ Single actor"]
-        Fanout[".fanout()<br/>→ ALL replicas"]
+        Route[".route()
+        → One healthy replica"]
+        CallOne[".call_one()
+        → Single actor"]
+        Fanout[".fanout()
+        → ALL replicas"]
     end
 
     subgraph Replicas["Replicas/Actors"]