Fix pretty printing in mjspec notebook.

PiperOrigin-RevId: 721367994
Change-Id: Ia1c4ee5b64540ff2f9c08a1953295ff366e0c72b

python/mjspec.ipynb

@@ -95,8 +95,6 @@
     "print('Installation successful.')\n",
     "\n",
     "# Other imports and helper functions\n",
-    "import time\n",
-    "import itertools\n",
     "import numpy as np\n",
     "\n",
     "# Graphics and plotting.\n",
@@ -106,8 +104,10 @@
     "import mediapy as media\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "# More legible printing from numpy.\n",
+    "# Printing.\n",
     "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
+    "import pygments\n",
+    "from google.colab import output\n",
     "\n",
     "from IPython.display import clear_output, HTML, display\n",
     "clear_output()\n",
@@ -121,13 +121,17 @@
     "!git clone https://github.com/google-deepmind/mujoco_menagerie\n",
     "franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n",
     "\n",
+    "\n",
+    "is_dark = output.eval_js('document.documentElement.matches(\"[theme=dark]\")')\n",
+    "print_style = 'monokai' if is_dark else 'lovelace'\n",
+    "\n",
     "def print_xml(xml_string):\n",
-    "  formatter = pygments.formatters.HtmlFormatter(style='lovelace')\n",
+    "  formatter = pygments.formatters.HtmlFormatter(style=print_style)\n",
     "  lexer = pygments.lexers.XmlLexer()\n",
     "  highlighted = pygments.highlight(xml_string, lexer, formatter)\n",
     "  display(HTML(f\"<style>{formatter.get_style_defs()}</style>{highlighted}\"))\n",
     "\n",
-    "def render(model, data=None, height=250):\n",
+    "def render(model, data=None, height=300):\n",
     "  if data is None:\n",
     "    data = mj.MjData(model)\n",
     "  with mj.Renderer(model, 480, 640) as renderer:\n",
@@ -254,7 +258,7 @@
     "      pixels = renderer.render()\n",
     "      frames.append(pixels)\n",
     "\n",
-    "media.show_video(frames, fps=framerate)\n"
+    "media.show_video(frames, fps=framerate)"
    ]
   },
   {
@@ -812,9 +816,6 @@
  "metadata": {
   "accelerator": "GPU",
   "colab": {
-   "collapsed_sections": [
-    "sJFuNetilv4m"
-   ],
    "gpuClass": "premium",
    "private_outputs": true
   },