[Feature] Automatic asset download checks (#456)

* add better tooling for asset management

* support registering asset download ids for robots and auto prompt user to download robots if used

* hierarchical data groups and bug fixes

* bug fixes

* fill in asset download ids

mani_skill/agents/base_agent.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
@@ -15,7 +16,7 @@
     PDJointPosControllerConfig,
 )
 from mani_skill.sensors.base_sensor import BaseSensor, BaseSensorConfig
-from mani_skill.utils import sapien_utils
+from mani_skill.utils import assets, download_asset, sapien_utils
 from mani_skill.utils.structs import Actor, Array, Articulation, Pose
 
 from .controllers.base_controller import (
@@ -157,6 +158,25 @@ def _load_articulation(self):
             sapien_utils.check_urdf_config(urdf_config)
             sapien_utils.apply_urdf_config(loader, urdf_config)
 
+        if not os.path.exists(asset_path):
+            print(f"Robot {self.uid} definition file not found at {asset_path}")
+            if len(assets.DATA_GROUPS[self.uid]) > 0:
+                response = download_asset.prompt_yes_no(
+                    f"Robot {self.uid} has assets available for download. Would you like to download them now?"
+                )
+                if response:
+                    for (
+                        asset_id
+                    ) in assets.expand_data_group_into_individual_data_source_ids(
+                        self.uid
+                    ):
+                        download_asset.download(assets.DATA_SOURCES[asset_id])
+                else:
+                    print(f"Exiting as assets for robot {self.uid} are not downloaded")
+                    exit()
+            else:
+                print(f"Exiting as assets for robot {self.uid} are not found")
+                exit()
         self.robot: Articulation = loader.load(asset_path)
         assert self.robot is not None, f"Fail to load URDF/MJCF from {asset_path}"
 

mani_skill/agents/registration.py

@@ -1,21 +1,21 @@
 from dataclasses import dataclass
-from typing import Dict
+from typing import Dict, List
 
 from mani_skill import logger
 from mani_skill.agents.base_agent import BaseAgent
+from mani_skill.utils import assets
 
 
 @dataclass
 class AgentSpec:
-    """Agent specifications. At the moment it is a simple wrapper around the agent_cls but the dataclass is used in case we may need additional metadata"""
-
     agent_cls: type[BaseAgent]
+    asset_download_ids: List[str]
 
 
 REGISTERED_AGENTS: Dict[str, AgentSpec] = {}
 
 
-def register_agent(override=False):
+def register_agent(asset_download_ids: List[str] = [], override=False):
     """A decorator to register agents into ManiSkill so they can be used easily by string uid.
 
     Args:
@@ -34,7 +34,10 @@ def _register_agent(agent_cls: type[BaseAgent]):
                 )
             return agent_cls
 
-        REGISTERED_AGENTS[agent_cls.uid] = AgentSpec(agent_cls=agent_cls)
+        REGISTERED_AGENTS[agent_cls.uid] = AgentSpec(
+            agent_cls=agent_cls, asset_download_ids=asset_download_ids
+        )
+        assets.DATA_GROUPS[agent_cls.uid] = asset_download_ids
         return agent_cls
 
     return _register_agent

mani_skill/agents/robots/anymal/anymal_c.py

@@ -11,7 +11,7 @@
 from mani_skill.utils.structs.articulation import Articulation
 
 
-@register_agent()
+@register_agent(asset_download_ids=["anymal_c"])
 class ANYmalC(BaseAgent):
     uid = "anymal_c"
     urdf_path = f"{ASSET_DIR}/robots/anymal_c/urdf/anymal.urdf"

mani_skill/agents/robots/googlerobot/googlerobot.py

@@ -9,7 +9,7 @@
 
 
 # TODO (stao) (xuanlin): Add mobile base, model it properly based on real2sim
-@register_agent()
+@register_agent(asset_download_ids=["googlerobot"])
 class GoogleRobot(BaseAgent):
     uid = "googlerobot"
     urdf_path = (

mani_skill/agents/robots/stompy/stompy.py

@@ -9,7 +9,7 @@
 from mani_skill.sensors.camera import CameraConfig
 
 
-@register_agent()  # uncomment this if you want to register the agent so you can instantiate it by ID when creating environments
+@register_agent(asset_download_ids=["stompy"])
 class Stompy(BaseAgent):
     uid = "stompy"
     urdf_path = f"{ASSET_DIR}/robots/stompy/robot.urdf"

mani_skill/agents/robots/unitree_g1/g1.py

@@ -8,7 +8,7 @@
 from mani_skill.sensors.camera import CameraConfig
 
 
-@register_agent()
+@register_agent(asset_download_ids=["unitree_g1"])
 class UnitreeG1(BaseAgent):
     uid = "unitree_g1"
     urdf_path = f"{ASSET_DIR}/robots/unitree_g1/g1.urdf"
@@ -19,10 +19,7 @@ class UnitreeG1(BaseAgent):
     keyframes = dict(
         standing=Keyframe(
             pose=sapien.Pose(p=[0, 0, 0.755]),
-            qpos=np.array(
-                [0.0] * 37
-            )
-            * 1,
+            qpos=np.array([0.0] * 37) * 1,
         )
     )
 
@@ -112,7 +109,7 @@ def is_fallen(self):
         return self.robot.pose.p[:, 2] < 0.3
 
 
-@register_agent()
+@register_agent(asset_download_ids=["unitree_g1"])
 class UnitreeG1Simplified(UnitreeG1):
     uid = "unitree_g1_simplified_legs"
     urdf_path = f"{ASSET_DIR}/robots/unitree_g1/g1_simplified_legs.urdf"

mani_skill/agents/robots/unitree_g1/g1_upper_body.py

@@ -11,7 +11,7 @@
 from mani_skill.utils.structs.actor import Actor
 
 
-@register_agent()
+@register_agent(asset_download_ids=["unitree_g1"])
 class UnitreeG1UpperBody(BaseAgent):
     uid = "unitree_g1_simplified_upper_body"
     urdf_path = f"{ASSET_DIR}/robots/unitree_g1/g1_simplified_upper_body.urdf"
@@ -194,7 +194,7 @@ def right_hand_is_grasping(self, object: Actor, min_force=0.5, max_angle=85):
         return torch.logical_and(lflag, rflag)
 
 
-@register_agent()
+@register_agent(asset_download_ids=["unitree_g1"])
 class UnitreeG1UpperBodyRightArm(UnitreeG1UpperBody):
     uid = "unitree_g1_simplified_upper_body_right_arm"
     urdf_path = f"{ASSET_DIR}/robots/unitree_g1/g1_simplified_upper_body.urdf"

mani_skill/agents/robots/unitree_go/unitree_go2.py

@@ -8,7 +8,9 @@
 from mani_skill.sensors.camera import CameraConfig
 
 
-@register_agent()  # uncomment this if you want to register the agent so you can instantiate it by ID when creating environments
+@register_agent(
+    asset_download_ids=["unitree_go2"]
+)  # uncomment this if you want to register the agent so you can instantiate it by ID when creating environments
 class UnitreeGo2(BaseAgent):
     uid = "unitree_go2"
     urdf_path = f"{ASSET_DIR}/robots/unitree_go2/urdf/go2_description.urdf"  # You can use f"{PACKAGE_ASSET_DIR}" to reference a urdf file in the mani_skill /assets package folder

mani_skill/agents/robots/unitree_h1/h1.py

@@ -8,7 +8,7 @@
 from mani_skill.sensors.camera import CameraConfig
 
 
-@register_agent()
+@register_agent(asset_download_ids=["unitree_h1"])
 class UnitreeH1(BaseAgent):
     uid = "unitree_h1"
     urdf_path = f"{ASSET_DIR}/robots/unitree_h1/urdf/h1.urdf"
@@ -114,7 +114,7 @@ def is_fallen(self):
         return self.robot.pose.p[:, 2] < 0.3
 
 
-@register_agent()
+@register_agent(asset_download_ids=["unitree_h1"])
 class UnitreeH1Simplified(UnitreeH1):
     uid = "unitree_h1_simplified"
     urdf_path = f"{ASSET_DIR}/robots/unitree_h1/urdf/h1_simplified.urdf"

mani_skill/agents/robots/ur_e/ur_10e.py

@@ -7,7 +7,7 @@
 from mani_skill.agents.registration import register_agent
 
 
-@register_agent()
+@register_agent(asset_download_ids=["ur10e"])
 class UR10e(BaseAgent):
     uid = "ur_10e"
     mjcf_path = f"{ASSET_DIR}/robots/ur10e/ur10e.xml"

mani_skill/agents/robots/widowx/widowx.py

@@ -9,7 +9,7 @@
 
 
 # TODO (stao) (xuanlin): model it properly based on real2sim
-@register_agent()
+@register_agent(asset_download_ids=["widowx250s"])
 class WidowX250S(BaseAgent):
     uid = "widowx250s"
     urdf_path = f"{ASSET_DIR}/robots/widowx250s/wx250s.urdf"

mani_skill/agents/robots/xmate3/xmate3.py

@@ -15,7 +15,7 @@
 from mani_skill.utils.structs.actor import Actor
 
 
-@register_agent()
+@register_agent(asset_download_ids=["xmate3_robotiq"])
 class Xmate3Robotiq(BaseAgent):
     uid = "xmate3_robotiq"
     urdf_path = f"{ASSET_DIR}/robots/xmate3_robotiq/xmate3_robotiq.urdf"

mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py

@@ -345,7 +345,11 @@ def __init__(self, *args, **kwargs):
         )
 
 
-@register_env("RotateSingleObjectInHandLevel2-v1", max_episode_steps=300)
+@register_env(
+    "RotateSingleObjectInHandLevel2-v1",
+    max_episode_steps=300,
+    asset_download_ids=["ycb"],
+)
 class RotateSingleObjectInHandLevel2(RotateSingleObjectInHand):
     def __init__(self, *args, **kwargs):
         super().__init__(
@@ -357,7 +361,11 @@ def __init__(self, *args, **kwargs):
         )
 
 
-@register_env("RotateSingleObjectInHandLevel3-v1", max_episode_steps=300)
+@register_env(
+    "RotateSingleObjectInHandLevel3-v1",
+    max_episode_steps=300,
+    asset_download_ids=["ycb"],
+)
 class RotateSingleObjectInHandLevel3(RotateSingleObjectInHand):
     def __init__(self, *args, **kwargs):
         super().__init__(

mani_skill/envs/tasks/mobile_manipulation/open_cabinet_drawer.py

@@ -25,7 +25,11 @@
 
 # TODO (stao): we need to cut the meshes of all the cabinets in this dataset for gpu sim, there may be some wierd physics
 # that may happen although it seems okay for state based RL
-@register_env("OpenCabinetDrawer-v1", max_episode_steps=100)
+@register_env(
+    "OpenCabinetDrawer-v1",
+    asset_download_ids=["partnet_mobility_cabinet"],
+    max_episode_steps=100,
+)
 class OpenCabinetDrawerEnv(BaseEnv):
 
     SUPPORTED_ROBOTS = ["fetch"]
@@ -306,9 +310,9 @@ def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: Dict):
             self.target_qpos - self.handle_link.joint.qpos, self.target_qpos
         )
         open_reward = 2 * (1 - amount_to_open_left)
-        reaching_reward[amount_to_open_left < 0.999] = (
-            2  # if joint opens even a tiny bit, we don't need reach reward anymore
-        )
+        reaching_reward[
+            amount_to_open_left < 0.999
+        ] = 2  # if joint opens even a tiny bit, we don't need reach reward anymore
         # print(open_reward.shape)
         open_reward[info["open_enough"]] = 3  # give max reward here
         reward = reaching_reward + open_reward

mani_skill/envs/tasks/tabletop/assembling_kits.py

@@ -19,7 +19,9 @@
 from mani_skill.utils.structs.types import GPUMemoryConfig, SimConfig
 
 
-@register_env("AssemblingKits-v1", max_episode_steps=200)
+@register_env(
+    "AssemblingKits-v1", asset_download_ids=["assembling_kits"], max_episode_steps=200
+)
 class AssemblingKitsEnv(BaseEnv):
     SUPPORTED_REWARD_MODES = ["sparse", "none"]
     SUPPORTED_ROBOTS = ["panda_wristcam"]
@@ -64,7 +66,9 @@ def __init__(
 
     @property
     def _default_sim_config(self):
-        return SimConfig(gpu_memory_cfg=GPUMemoryConfig(max_rigid_contact_count=2**20))
+        return SimConfig(
+            gpu_memory_cfg=GPUMemoryConfig(max_rigid_contact_count=2**20)
+        )
 
     @property
     def _default_sensor_configs(self):

mani_skill/envs/tasks/tabletop/pick_clutter_ycb.py

@@ -196,7 +196,11 @@ def compute_normalized_dense_reward(
         return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward
 
 
-@register_env("PickClutterYCB-v1", max_episode_steps=100)
+@register_env(
+    "PickClutterYCB-v1",
+    asset_download_ids=["ycb", "pick_clutter_ycb_configs"],
+    max_episode_steps=100,
+)
 class PickClutterYCBEnv(PickClutterEnv):
     DEFAULT_EPISODE_JSON = f"{ASSET_DIR}/tasks/pick_clutter/ycb_train_5k.json.gz"
 

mani_skill/envs/tasks/tabletop/pick_single_ycb.py

@@ -24,7 +24,7 @@
 WARNED_ONCE = False
 
 
-@register_env("PickSingleYCB-v1", max_episode_steps=50)
+@register_env("PickSingleYCB-v1", max_episode_steps=50, asset_download_ids=["ycb"])
 class PickSingleYCBEnv(BaseEnv):
 
     SUPPORTED_ROBOTS = ["panda", "panda_wristcam", "xmate3_robotiq", "fetch"]

mani_skill/utils/assets/README.md

@@ -0,0 +1,4 @@
+# ManiSkill Asset Management
+
+Code for asset management in ManiSkill
+

mani_skill/utils/assets/__init__.py

@@ -0,0 +1,7 @@
+from .data import (
+    DATA_GROUPS,
+    DATA_SOURCES,
+    DataSource,
+    expand_data_group_into_individual_data_source_ids,
+    is_data_source_downloaded,
+)