From 0040d5049209323b11c3373b697f6f7eb97b027c Mon Sep 17 00:00:00 2001 From: Russ Tedrake Date: Fri, 17 Jan 2025 10:57:05 -0500 Subject: [PATCH] Update robocasa example to use make_drake_compatible_model Adds geometry group re-mapping support to mdcm. Towards #425. --- book/mobile/robocasa.ipynb | 66 +++++++++++++++---- manipulation/make_drake_compatible_model.py | 29 +++++++- manipulation/test/models/test_defaults.xml | 1 + .../test/test_make_drake_compatible_model.py | 2 + 4 files changed, 82 insertions(+), 16 deletions(-) diff --git a/book/mobile/robocasa.ipynb b/book/mobile/robocasa.ipynb index 92e7f2a9..83d43c6c 100644 --- 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 index bda3e376..30c04a29 100644 --- 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 index a181494c..5bfdf96f 100644 --- a/manipulation/test/models/test_defaults.xml +++ b/manipulation/test/models/test_defaults.xml @@ -3,6 +3,7 @@ + diff --git a/manipulation/test/test_make_drake_compatible_model.py b/manipulation/test/test_make_drake_compatible_model.py index f36763dc..2e2eeab0 100644 --- a/manipulation/test/test_make_drake_compatible_model.py +++ b/manipulation/test/test_make_drake_compatible_model.py @@ -128,6 +128,7 @@ 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: @@ -135,6 +136,7 @@ def test_mjcf_defaults(self): 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)