Highlight XML strings using pygments in mjSpec notebook.

PiperOrigin-RevId: 704751058
Change-Id: Iaf7fae7f13ef12d59bac762170a6dfca01494164
This commit is contained in:
Yuval Tassa
2024-12-10 10:05:53 -08:00
committed by Copybara-Service
parent a364308d5a
commit f4593d5d79
+10 -5
View File
@@ -109,7 +109,7 @@
"# More legible printing from numpy.\n", "# More legible printing from numpy.\n",
"np.set_printoptions(precision=3, suppress=True, linewidth=100)\n", "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
"\n", "\n",
"from IPython.display import clear_output\n", "from IPython.display import clear_output, HTML, display\n",
"clear_output()\n", "clear_output()\n",
"\n", "\n",
"# Get MuJoCo's humanoid model and a Franka arm from the MuJoCo Menagerie.\n", "# Get MuJoCo's humanoid model and a Franka arm from the MuJoCo Menagerie.\n",
@@ -121,6 +121,12 @@
"!git clone https://github.com/google-deepmind/mujoco_menagerie\n", "!git clone https://github.com/google-deepmind/mujoco_menagerie\n",
"franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n", "franka_file = 'mujoco_menagerie/franka_fr3/fr3.xml'\n",
"\n", "\n",
"def print_xml(xml_string):\n",
" formatter = pygments.formatters.HtmlFormatter(style='lovelace')\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=250):\n",
" if data is None:\n", " if data is None:\n",
" data = mj.MjData(model)\n", " data = mj.MjData(model)\n",
@@ -172,6 +178,7 @@
"render(model)\n", "render(model)\n",
"\n", "\n",
"# Change the mjSpec, re-compile and re-render\n", "# Change the mjSpec, re-compile and re-render\n",
"spec.modelname = \"edited model\"\n",
"geoms = spec.worldbody.find_all(mj.mjtObj.mjOBJ_GEOM)\n", "geoms = spec.worldbody.find_all(mj.mjtObj.mjOBJ_GEOM)\n",
"geoms[0].name = 'blue_box'\n", "geoms[0].name = 'blue_box'\n",
"geoms[0].rgba = [0, 0, 1, 1]\n", "geoms[0].rgba = [0, 0, 1, 1]\n",
@@ -193,7 +200,7 @@
"id": "Tw_yUwqxKwCI" "id": "Tw_yUwqxKwCI"
}, },
"source": [ "source": [
"`mjSpec` can save XML to string, saving all modifications." "`mjSpec` can save XML to string, with all modifications:"
] ]
}, },
{ {
@@ -204,9 +211,7 @@
}, },
"outputs": [], "outputs": [],
"source": [ "source": [
"#@title Print an XML from an `mjSpec` {vertical-output: true}\n", "print_xml(spec.to_xml())"
"\n",
"print(spec.to_xml())"
] ]
}, },
{ {