Update robocasa example to use make_drake_compatible_model

Adds geometry group re-mapping support to mdcm. Towards #425.
RussTedrake · Jan 17, 2025 · 0040d50 · 0040d50
1 parent da1e295
commit 0040d50
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 16 deletions.
diff --git a/book/mobile/robocasa.ipynb b/book/mobile/robocasa.ipynb
@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "99268d0c",
    "metadata": {
     "colab": {},
@@ -29,7 +29,10 @@
     "from zipfile import ZipFile\n",
     "\n",
     "from pydrake.all import ModelVisualizer, PackageMap, Simulator, StartMeshcat\n",
-    "from tqdm.notebook import tqdm"
+    "from tqdm.notebook import tqdm\n",
+    "\n",
+    "from manipulation.make_drake_compatible_model import MakeDrakeCompatibleModel\n",
+    "from manipulation.utils import running_as_notebook"
    ]
   },
   {
@@ -55,15 +58,15 @@
     "        package_name=\"robocasa\",\n",
     "        params=PackageMap.RemoteParams(\n",
     "            urls=[\n",
-    "                f\"https://github.com/robocasa/robocasa/archive/1370b9e0f747d84fb21ed29bacefb1654865301b.zip\"\n",
+    "                f\"https://github.com/robocasa/robocasa/archive/a7586f67b72e51722340c9dbe97a59e0aff1ff8f.zip\"\n",
     "            ],\n",
-    "            sha256=(\"a7218ed369936f96b19467eee5038870b14c2b7f91a9d6108591394ed074b337\"),\n",
-    "            strip_prefix=\"robocasa-1370b9e0f747d84fb21ed29bacefb1654865301b/robocasa/\",\n",
+    "            sha256=(\"64abe91ac8ca9cbb22aa4d7c9461d7d899ba54e700dccecdfab5a827fcf7f322\"),\n",
+    "            strip_prefix=\"robocasa-a7586f67b72e51722340c9dbe97a59e0aff1ff8f/robocasa/\",\n",
     "        ),\n",
     "    )\n",
     "\n",
     "\n",
-    "def DownloadRobocasaKitchenAssets():\n",
+    "def MaybeDownloadRobocasaKitchenAssets():\n",
     "    package_map = PackageMap()\n",
     "    AddRobocasaRemote(package_map)\n",
     "    # This will force the download if it hasn't been done before.\n",
@@ -73,12 +76,36 @@
     "    # https://github.com/robocasa/robocasa/blob/main/robocasa/scripts/download_kitchen_assets.py\n",
     "    # with robocasa_path updated.\n",
     "    DOWNLOAD_ASSET_REGISTRY = dict(\n",
+    "        textures=dict(\n",
+    "            message=\"Downloading environment textures\",\n",
+    "            url=\"https://utexas.box.com/shared/static/otdsyfjontk17jdp24bkhy2hgalofbh4.zip\",\n",
+    "            folder=os.path.join(robocasa_path, \"models/assets/textures\"),\n",
+    "            check_folder_exists=True,\n",
+    "        ),\n",
     "        fixtures=dict(\n",
     "            message=\"Downloading fixtures\",\n",
-    "            url=\"https://utexas.box.com/shared/static/956d0w2ucqs7d3eors1idsohgum57nli.zip\",\n",
+    "            url=\"https://utexas.box.com/shared/static/pobhbsjyacahg2mx8x4rm5fkz3wlmyzp.zip\",\n",
     "            folder=os.path.join(robocasa_path, \"models/assets/fixtures\"),\n",
-    "            check_folder_exists=False,\n",
+    "            check_folder_exists=True,\n",
     "        ),\n",
+    "        objaverse=dict(\n",
+    "            message=\"Downloading objaverse objects\",\n",
+    "            url=\"https://utexas.box.com/shared/static/ejt1kc2v5vhae1rl4k5697i4xvpbjcox.zip\",\n",
+    "            folder=os.path.join(robocasa_path, \"models/assets/objects/objaverse\"),\n",
+    "            check_folder_exists=True,\n",
+    "        ),\n",
+    "        #  aigen_objs=dict(\n",
+    "        #     message=\"Downloading AI-generated objects\",\n",
+    "        #     url=\"https://utexas.box.com/shared/static/os3hrui06lasnuvwqpmwn0wcrduh6jg3.zip\",\n",
+    "        #     folder=os.path.join(robocasa_path, \"models/assets/objects/aigen_objs\"),\n",
+    "        #     check_folder_exists=False,\n",
+    "        # ),\n",
+    "        # generative_textures=dict(\n",
+    "        #     message=\"Downloading AI-generated environment textures\",\n",
+    "        #     url=\"https://utexas.box.com/shared/static/gf9nkadvfrowkb9lmkcx58jwt4d6c1g3.zip\",\n",
+    "        #     folder=os.path.join(robocasa_path, \"models/assets/generative_textures\"),\n",
+    "        #     check_folder_exists=False,\n",
+    "        # ),\n",
     "    )\n",
     "\n",
     "    def show_progress(block_num, block_size, total_size):\n",
@@ -87,6 +114,9 @@
     "        pbar.update(block_size)\n",
     "\n",
     "    for name, info in DOWNLOAD_ASSET_REGISTRY.items():\n",
+    "        if info[\"check_folder_exists\"] and os.path.exists(info[\"folder\"]):\n",
+    "            print(f\"Skipping {name} - already downloaded\")\n",
+    "            continue\n",
     "        with tqdm(unit=\"B\", unit_scale=True, miniters=1, desc=info[\"message\"]) as pbar:\n",
     "            filename, headers = urllib.request.urlretrieve(\n",
     "                info[\"url\"], reporthook=show_progress\n",
@@ -97,9 +127,8 @@
     "        os.remove(filename)\n",
     "\n",
     "\n",
-    "# You'll only want to run this once.\n",
-    "# TODO(russt): Update this to MaybeDownloadRobocasaKitchenAssets.\n",
-    "# DownloadRobocasaKitchenAssets()"
+    "if running_as_notebook:\n",
+    "    MaybeDownloadRobocasaKitchenAssets()"
    ]
   },
   {
@@ -124,10 +153,19 @@
    "outputs": [],
    "source": [
     "visualizer = ModelVisualizer(meshcat=meshcat)\n",
-    "AddRobocasaRemote(visualizer.parser().package_map())\n",
-    "visualizer.AddModels(\n",
-    "    url=\"package://robocasa/models/assets/fixtures/accessories/knife_blocks/dark_wood/model.xml\"\n",
+    "package_map = visualizer.parser().package_map()\n",
+    "AddRobocasaRemote(package_map)\n",
+    "original_model_path = package_map.ResolveUrl(\n",
+    "    \"package://robocasa/models/assets/fixtures/accessories/knife_blocks/dark_wood/model.xml\"\n",
+    ")\n",
+    "drake_model_path = original_model_path.replace(\".xml\", \".drake.xml\")\n",
+    "MakeDrakeCompatibleModel(\n",
+    "    original_model_path,\n",
+    "    drake_model_path,\n",
+    "    overwrite=True,\n",
+    "    remap_mujoco_geometry_groups={0: 3},\n",
     ")\n",
+    "visualizer.AddModels(drake_model_path)\n",
     "visualizer.Run(loop_once=True)\n",
     "meshcat.DeleteAddedControls()"
    ]

diff --git a/manipulation/make_drake_compatible_model.py b/manipulation/make_drake_compatible_model.py
@@ -262,7 +262,12 @@ def _apply_defaults(
     return element
 
 
-def _convert_mjcf(input_filename: str, output_filename: str, overwrite: bool) -> None:
+def _convert_mjcf(
+    input_filename: str,
+    output_filename: str,
+    overwrite: bool,
+    remap_geometry_groups: dict[int, int] = {},
+) -> None:
     """Convert an MJCF file to be compatible with Drake.
 
     Args:
@@ -495,6 +500,15 @@ def process_defaults(element, parent_class="main"):
                 break
             parent = parent.getparent()
 
+    # Remap geometry groups
+    if remap_geometry_groups:
+        geoms = root.findall(".//geom")
+        for geom in geoms:
+            if "group" in geom.attrib:
+                group = int(geom.attrib["group"])
+                if group in remap_geometry_groups:
+                    geom.attrib["group"] = str(remap_geometry_groups[group])
+
     tree.write(output_filename, pretty_print=True)
     print(f"Converted MJCF file {input_filename} to {output_filename}")
 
@@ -504,6 +518,7 @@ def MakeDrakeCompatibleModel(
     output_filename: str,
     package_map: PackageMap = PackageMap(),
     overwrite: bool = False,
+    remap_mujoco_geometry_groups: dict[int, int] = {},
 ) -> None:
     """Converts a model file (currently .urdf or .xml)to be compatible with the
     Drake multibody parsers.
@@ -534,13 +549,23 @@ def MakeDrakeCompatibleModel(
         package_map (PackageMap, optional): The package map to use. Defaults to None.
         overwrite (bool, optional): Whether to overwrite existing files. Defaults
             to False.
+        remap_mujoco_geometry_groups (dict[int, int], optional): Drake's mujoco
+            parser registers visual geometry for geometry groups < 3 (the
+            mujoco default), which is a common, but not universal, convention.
+            This argument allows you to remap (substituting the value for the
+            key).
     """
     if input_filename.lower().endswith(".urdf"):
         _convert_urdf(input_filename, output_filename, package_map, overwrite=overwrite)
     elif input_filename.lower().endswith(".sdf"):
         _convert_sdf(input_filename, output_filename, package_map, overwrite=overwrite)
     elif input_filename.lower().endswith(".xml"):
-        _convert_mjcf(input_filename, output_filename, overwrite=overwrite)
+        _convert_mjcf(
+            input_filename,
+            output_filename,
+            overwrite=overwrite,
+            remap_geometry_groups=remap_mujoco_geometry_groups,
+        )
     else:
         print(
             f"Warning: The file extension of '{input_filename}' is not "

diff --git a/manipulation/test/models/test_defaults.xml b/manipulation/test/models/test_defaults.xml
@@ -3,6 +3,7 @@
     <default class="scaled">
       <mesh scale="0.001 0.002 0.003"/>
     </default>
+    <geom group="0"/>
   </default>
   <asset>
     <mesh name="test_stl" file="cube.stl" class="scaled"/>

diff --git a/manipulation/test/test_make_drake_compatible_model.py b/manipulation/test/test_make_drake_compatible_model.py
@@ -128,13 +128,15 @@ def test_mjcf_defaults(self):
             input_filename=input_filename,
             output_filename=output_filename,
             package_map=package_map,
+            remap_mujoco_geometry_groups={0: 3},
         )
         self.assertTrue(os.path.exists(output_filename))
         with open(output_filename, "r") as f:
             output_content = f.read()
         self.assertIn(
             'file="cube_from_stl_scaled_0.001_0.002_0.003.obj"', output_content
         )
+        self.assertIn('group="3"', output_content)
         # Clean up the temp file
         os.remove(output_filename)