Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update robocasa example to use make_drake_compatible_model #427

Merged
merged 1 commit into from
Jan 17, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 52 additions & 14 deletions book/mobile/robocasa.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 1,
"id": "99268d0c",
"metadata": {
"colab": {},
Expand All @@ -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"
]
},
{
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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()"
]
},
{
Expand All @@ -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()"
]
Expand Down
29 changes: 27 additions & 2 deletions manipulation/make_drake_compatible_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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}")

Expand All @@ -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.
Expand Down Expand Up @@ -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 "
Expand Down
1 change: 1 addition & 0 deletions manipulation/test/models/test_defaults.xml
Original file line number Diff line number Diff line change
Expand Up @@ -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"/>
Expand Down
2 changes: 2 additions & 0 deletions manipulation/test/test_make_drake_compatible_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
Loading