From d49a37f9d52057070b291cca3e6e61e7b0e49fbb Mon Sep 17 00:00:00 2001 From: StoneT2000 Date: Tue, 12 Mar 2024 02:03:45 -0700 Subject: [PATCH 1/2] move out scene masks in favor of scene idxs --- mani_skill/envs/tasks/assembling_kits.py | 8 ++-- .../dexterity/rotate_single_object_in_hand.py | 10 ++--- .../envs/tasks/dexterity/rotate_valve.py | 7 ++-- mani_skill/envs/tasks/open_cabinet_drawer.py | 2 +- mani_skill/envs/tasks/peg_insertion_side.py | 7 ++-- mani_skill/utils/building/actor_builder.py | 26 ++----------- .../utils/building/articulation_builder.py | 39 +++++++------------ mani_skill/utils/building/articulations.py | 8 ++-- mani_skill/utils/building/urdf_loader.py | 5 ++- mani_skill/utils/structs/actor.py | 15 +++---- mani_skill/utils/structs/articulation.py | 15 +++---- mani_skill/utils/structs/base.py | 11 +++--- mani_skill/utils/structs/joint.py | 2 +- mani_skill/utils/structs/link.py | 18 +++++---- 14 files changed, 74 insertions(+), 99 deletions(-) diff --git a/mani_skill/envs/tasks/assembling_kits.py b/mani_skill/envs/tasks/assembling_kits.py index e526456cf..193ec0bb5 100644 --- a/mani_skill/envs/tasks/assembling_kits.py +++ b/mani_skill/envs/tasks/assembling_kits.py @@ -92,7 +92,7 @@ def _load_scene(self): self.goal_rot = np.zeros((self.num_envs,)) for i, eps_idx in enumerate(eps_idxs): - scene_mask = [i] + scene_idxs = [i] episode = self._episodes[eps_idx] # get the kit builder and the goal positions/rotations of all other objects @@ -102,7 +102,7 @@ def _load_scene(self): object_goal_rot, ) = self._get_kit_builder_and_goals(episode["kit"]) kit = ( - kit_builder.set_scene_idxs(scene_mask) + kit_builder.set_scene_idxs(scene_idxs) .set_initial_pose(sapien.Pose([0, 0, 0.01])) .build_static(f"kit_{i}") ) @@ -110,7 +110,7 @@ def _load_scene(self): # create the object to place and make it dynamic obj_to_place = ( self._get_object_builder(episode["obj_to_place"]) - .set_scene_idxs(scene_mask) + .set_scene_idxs(scene_idxs) .build(f"obj_{i}") ) self.object_ids.append(episode["obj_to_place"]) @@ -119,7 +119,7 @@ def _load_scene(self): # create all other objects and leave them as static as they do not need to be manipulated other_objs = [ self._get_object_builder(obj_id, static=True) - .set_scene_idxs(scene_mask) + .set_scene_idxs(scene_idxs) .set_initial_pose( sapien.Pose( object_goal_pos[obj_id], diff --git a/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py b/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py index 76b2522fd..60cd8a2f7 100644 --- a/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py +++ b/mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py @@ -110,9 +110,8 @@ def _load_scene(self): base_color=np.array([255, 255, 255, 255]) / 255, ), ) - scene_mask = np.zeros(self.num_envs, dtype=bool) - scene_mask[i] = True - builder.set_scene_mask(scene_mask) + scene_idxs = [i] + builder.set_scene_idxs(scene_idxs) actors.append(builder.build(name=f"cube-{i}")) obj_heights.append(half_size) self.obj = Actor.merge(actors, name="cube") @@ -128,9 +127,8 @@ def _load_scene(self): builder, obj_height = build_actor_ycb( model_id, self._scene, name=model_id, return_builder=True ) - scene_mask = np.zeros(self.num_envs, dtype=bool) - scene_mask[i] = True - builder.set_scene_mask(scene_mask) + scene_idxs = [i] + builder.set_scene_idxs(scene_idxs) actors.append(builder.build(name=f"{model_id}-{i}")) obj_heights.append(obj_height) self.obj = Actor.merge(actors, name="ycb_object") diff --git a/mani_skill/envs/tasks/dexterity/rotate_valve.py b/mani_skill/envs/tasks/dexterity/rotate_valve.py index db158c0b1..7f478cddd 100644 --- a/mani_skill/envs/tasks/dexterity/rotate_valve.py +++ b/mani_skill/envs/tasks/dexterity/rotate_valve.py @@ -110,13 +110,12 @@ def _load_articulations(self): valves: List[Articulation] = [] capsule_lens = [] for i, valve_angles in enumerate(valve_angles_list): - scene_mask = np.zeros(self.num_envs, dtype=bool) - scene_mask[i] = True + scene_idxs = [i] if self.difficulty_level < 3: valve, capsule_len = build_robel_valve( self._scene, valve_angles=valve_angles, - scene_mask=scene_mask, + scene_idxs=scene_idxs, name=f"valve_station_{i}", ) else: @@ -124,7 +123,7 @@ def _load_articulations(self): valve, capsule_len = build_robel_valve( self._scene, valve_angles=valve_angles, - scene_mask=scene_mask, + scene_idxs=scene_idxs, name=f"valve_station_{i}", radius_scale=scales[0], capsule_radius_scale=scales[1], diff --git a/mani_skill/envs/tasks/open_cabinet_drawer.py b/mani_skill/envs/tasks/open_cabinet_drawer.py index 46fa60689..d0d9bd274 100644 --- a/mani_skill/envs/tasks/open_cabinet_drawer.py +++ b/mani_skill/envs/tasks/open_cabinet_drawer.py @@ -83,7 +83,7 @@ def _load_cabinets(self, joint_types: List[str]): scene_mask = np.zeros(self.num_envs, dtype=bool) scene_mask[i] = True cabinet, metadata = build_preprocessed_partnet_mobility_articulation( - self._scene, model_id, name=f"{model_id}-{i}", scene_mask=scene_mask + self._scene, model_id, name=f"{model_id}-{i}", scene_idxs=scene_mask ) # TODO (stao): since we processed the assets we know that the bounds[0, 1] is the actual height to set the object at # but in future we will store a visual origin offset so we can place them by using the actual bbox height / 2 diff --git a/mani_skill/envs/tasks/peg_insertion_side.py b/mani_skill/envs/tasks/peg_insertion_side.py index 9d50dd39d..cfaa29917 100644 --- a/mani_skill/envs/tasks/peg_insertion_side.py +++ b/mani_skill/envs/tasks/peg_insertion_side.py @@ -112,8 +112,7 @@ def _load_scene(self): boxes = [] for i in range(self.num_envs): - scene_mask = np.zeros((self.num_envs), dtype=bool) - scene_mask[i] = True + scene_idxs = [i] length = lengths[i] radius = radii[i] builder = self._scene.create_actor_builder() @@ -140,7 +139,7 @@ def _load_scene(self): half_size=[length / 2, radius, radius], material=mat, ) - builder.set_scene_mask(scene_mask) + builder.set_scene_idxs(scene_idxs) peg = builder.build(f"peg_{i}") # box with hole @@ -153,7 +152,7 @@ def _load_scene(self): builder = _build_box_with_hole( self._scene, inner_radius, outer_radius, depth, center=centers[i] ) - builder.set_scene_mask(scene_mask) + builder.set_scene_idxs(scene_idxs) box = builder.build_kinematic(f"box_with_hole_{i}") pegs.append(peg) diff --git a/mani_skill/utils/building/actor_builder.py b/mani_skill/utils/building/actor_builder.py index 2c965648c..926406ce1 100644 --- a/mani_skill/utils/building/actor_builder.py +++ b/mani_skill/utils/building/actor_builder.py @@ -28,23 +28,10 @@ class ActorBuilder(SAPIENActorBuilder): def __init__(self): super().__init__() self.initial_pose = Pose.create(self.initial_pose) - self.scene_mask = None self.scene_idxs = None self._allow_overlapping_plane_collisions = False self._plane_collision_poses = set() - def set_scene_mask( - self, - scene_mask: Optional[ - Union[List[bool], Sequence[bool], torch.Tensor, np.ndarray] - ] = None, - ): - """ - Set a scene mask so that the actor builder builds the actor only in a subset of the environments - """ - self.scene_mask = scene_mask - return self - def set_scene_idxs( self, scene_idxs: Optional[ @@ -187,16 +174,11 @@ def build(self, name): self.set_name(name) num_actors = self.scene.num_envs - if self.scene_mask is not None: - num_actors = np.sum(num_actors) - self.scene_mask = sapien_utils.to_tensor(self.scene_mask).to(bool) - self.scene_idxs = self.scene_mask.argwhere() - elif self.scene_idxs is not None: - self.scene_mask = torch.zeros((self.scene.num_envs), dtype=bool) - self.scene_mask[self.scene_idxs] = True + if self.scene_idxs is not None: + pass else: - self.scene_mask = torch.ones((self.scene.num_envs), dtype=bool) self.scene_idxs = torch.arange((self.scene.num_envs), dtype=int) + num_actors = len(self.scene_idxs) initial_pose = Pose.create(self.initial_pose) initial_pose_b = initial_pose.raw_pose.shape[0] assert initial_pose_b == 1 or initial_pose_b == num_actors @@ -217,7 +199,7 @@ def build(self, name): sub_scene.add_entity(entity) entities.append(entity) i += 1 - actor = Actor._create_from_entities(entities, self.scene, self.scene_mask) + actor = Actor.create_from_entities(entities, self.scene, self.scene_idxs) # if it is a static body type and this is a GPU sim but we are given a single initial pose, we repeat it for the purposes of observations if ( diff --git a/mani_skill/utils/building/articulation_builder.py b/mani_skill/utils/building/articulation_builder.py index ff32ae9d9..5770d6dac 100644 --- a/mani_skill/utils/building/articulation_builder.py +++ b/mani_skill/utils/building/articulation_builder.py @@ -25,23 +25,24 @@ class ArticulationBuilder(SapienArticulationBuilder): def __init__(self): super().__init__() - self.scene_mask = None self.name = None + self.scene_idxs = None def set_name(self, name: str): self.name = name return self - def set_scene_mask( + def set_scene_idxs( self, - scene_mask: Optional[ - Union[List[bool], Sequence[bool], torch.Tensor, np.ndarray] + scene_idxs: Optional[ + Union[List[int], Sequence[int], torch.Tensor, np.ndarray] ] = None, ): """ - Set a scene mask so that the articulation builder builds the articulation only in a subset of the environments + Set a list of scene indices to build this object in. Cannot be used in conjunction with scene mask """ - self.scene_mask = scene_mask + self.scene_idxs = scene_idxs + return self def create_link_builder(self, parent: LinkBuilder = None): if self.link_builders: @@ -105,25 +106,15 @@ def build(self, name=None, fix_root_link=None, build_mimic_joints=True): if name is not None: self.set_name(name) assert self.name is not None - - num_arts = self.scene.num_envs - if self.scene_mask is not None: - assert ( - len(self.scene_mask) == self.scene.num_envs - ), "Scene mask size is not correct. Must be the same as the number of sub scenes" - num_arts = np.sum(num_arts) - self.scene_mask = sapien_utils.to_tensor(self.scene_mask) + if self.scene_idxs is not None: + pass else: - # if scene mask is none, set it here - self.scene_mask = sapien_utils.to_tensor( - torch.ones((self.scene.num_envs), dtype=bool) - ) + self.scene_idxs = torch.arange((self.scene.num_envs), dtype=int) articulations = [] - for scene_idx, scene in enumerate(self.scene.sub_scenes): - if self.scene_mask is not None and self.scene_mask[scene_idx] == False: - continue + for scene_idx in self.scene_idxs: + sub_scene = self.scene.sub_scenes[scene_idx] links: List[sapien.Entity] = self.build_entities( name_prefix=f"scene-{scene_idx}-{self.name}_" ) @@ -176,13 +167,13 @@ def build(self, name=None, fix_root_link=None, build_mimic_joints=True): ) for l in links: - scene.add_entity(l) + sub_scene.add_entity(l) articulation: physx.PhysxArticulation = l.components[0].articulation articulation.name = f"scene-{scene_idx}_{self.name}" articulations.append(articulation) - articulation = Articulation._create_from_physx_articulations( - articulations, self.scene, self.scene_mask + articulation = Articulation.create_from_physx_articulations( + articulations, self.scene, self.scene_idxs ) self.scene.articulations[self.name] = articulation return articulation diff --git a/mani_skill/utils/building/articulations.py b/mani_skill/utils/building/articulations.py index 1941b24b6..94762c1db 100644 --- a/mani_skill/utils/building/articulations.py +++ b/mani_skill/utils/building/articulations.py @@ -79,7 +79,7 @@ def build_preprocessed_partnet_mobility_articulation( name: str, fix_root_link=True, urdf_config: dict = None, - scene_mask=None, + scene_idxs=None, ): """ Builds a physx.PhysxArticulation object into the scene and returns metadata containing annotations of the object's links and joints. @@ -107,7 +107,7 @@ def build_preprocessed_partnet_mobility_articulation( urdf_path = MODEL_DBS["PartnetMobility"]["model_urdf_paths"][model_id] urdf_config = sapien_utils.parse_urdf_config(urdf_config or {}, scene) sapien_utils.apply_urdf_config(loader, urdf_config) - articulation = loader.load(str(urdf_path), name=name, scene_mask=scene_mask) + articulation = loader.load(str(urdf_path), name=name, scene_idxs=scene_idxs) metadata = ArticulationMetadata( joints=dict(), links=dict(), movable_links=[], bbox=None, scale=loader.scale ) @@ -169,7 +169,7 @@ def build_robel_valve( name: str, radius_scale: float = 1.0, capsule_radius_scale: float = 1.0, - scene_mask=None, + scene_idxs=None, ): # Size and geometry of valve are based on the original setting of Robel benchmark, unit: m # Ref: https://github.com/google-research/robel @@ -182,7 +182,7 @@ def build_robel_valve( bearing_height = 0.032 builder = scene.create_articulation_builder() - builder.set_scene_mask(scene_mask) + builder.set_scene_idxs(scene_idxs) # Mount link mount_builder = builder.create_link_builder(parent=None) diff --git a/mani_skill/utils/building/urdf_loader.py b/mani_skill/utils/building/urdf_loader.py index d34bc1adf..a90c603fe 100644 --- a/mani_skill/utils/building/urdf_loader.py +++ b/mani_skill/utils/building/urdf_loader.py @@ -51,7 +51,7 @@ def load( srdf_file=None, package_dir=None, name=None, - scene_mask=None, + scene_idxs=None, ) -> Articulation: """ Args: @@ -59,6 +59,7 @@ def load( srdf_file: SRDF for urdf_file. If srdf_file is None, it defaults to the ".srdf" file with the same as the urdf file package_dir: base directory used to resolve asset files in the URDF file. If an asset path starts with "package://", "package://" is simply removed from the file name name (str): name of the created articulation + scene_idxs (list[int]): the ids of the scenes to build the objects in Returns: returns a single Articulation loaded from the URDF file. It throws an error if multiple objects exists """ @@ -75,7 +76,7 @@ def load( articulations: List[Articulation] = [] for b in articulation_builders: - b.set_scene_mask(scene_mask) + b.set_scene_idxs(scene_idxs) b.disable_self_collisions = self.disable_self_collisions articulations.append(b.build()) diff --git a/mani_skill/utils/structs/actor.py b/mani_skill/utils/structs/actor.py index a35734c9f..2a5af0d62 100644 --- a/mani_skill/utils/structs/actor.py +++ b/mani_skill/utils/structs/actor.py @@ -34,18 +34,18 @@ class Actor(PhysxRigidDynamicComponentStruct, BaseStruct[sapien.Entity]): # track the initial pose of the actor builder for this actor. Necessary to ensure the actor is reset correctly once # gpu system is initialized - _builder_initial_pose: sapien.Pose = None + _builder_initial_pose: Pose = None name: str = None def __hash__(self): return self._objs[0].__hash__() @classmethod - def _create_from_entities( + def create_from_entities( cls, entities: List[sapien.Entity], scene: ManiSkillScene, - scene_mask: torch.Tensor, + scene_idxs: torch.Tensor, ): shared_name = "_".join(entities[0].name.split("_")[1:]) @@ -69,7 +69,7 @@ def _create_from_entities( return cls( _objs=entities, _scene=scene, - _scene_mask=scene_mask, + _scene_idxs=scene_idxs, px_body_type=px_body_type, _bodies=bodies, _body_data_name="cuda_rigid_body_data" @@ -94,11 +94,11 @@ def merge(cls, actors: List["Actor"], name: str = None): objs = [] scene = actors[0]._scene _builder_initial_poses = [] - merged_scene_mask = actors[0]._scene_mask.clone() + merged_scene_idxs = [] num_objs_per_actor = actors[0]._num_objs for actor in actors: objs += actor._objs - merged_scene_mask[actor._scene_mask] = True + merged_scene_idxs.append(actor._scene_idxs) _builder_initial_poses.append(actor._builder_initial_pose.raw_pose) del scene.actors[actor.name] assert ( @@ -106,7 +106,8 @@ def merge(cls, actors: List["Actor"], name: str = None): ), "Each given actor must have the same number of managed objects" # TODO (stao): Can we support e.g. each Actor having len(actor._objs) > 1? It would mean fetching pose data or any kind of data is highly uintuitive # we definitely cannot permit some actors to have more objs than others, otherwise the data is ragged. - merged_actor = Actor._create_from_entities(objs, scene, merged_scene_mask) + merged_scene_idxs = torch.concat(merged_scene_idxs) + merged_actor = Actor.create_from_entities(objs, scene, merged_scene_idxs) merged_actor.name = name merged_actor._builder_initial_pose = Pose.create( torch.vstack(_builder_initial_poses) diff --git a/mani_skill/utils/structs/articulation.py b/mani_skill/utils/structs/articulation.py index 142eaa6f7..b888ef0b7 100644 --- a/mani_skill/utils/structs/articulation.py +++ b/mani_skill/utils/structs/articulation.py @@ -68,11 +68,11 @@ class Articulation(BaseStruct[physx.PhysxArticulation]): """Maps a tuple of link names to pre-saved net contact force queries""" @classmethod - def _create_from_physx_articulations( + def create_from_physx_articulations( cls, physx_articulations: List[physx.PhysxArticulation], scene: ManiSkillScene, - scene_mask: torch.Tensor, + scene_idxs: torch.Tensor, _merged: bool = False, ): shared_name = "_".join(physx_articulations[0].name.split("_")[1:]) @@ -83,7 +83,7 @@ def _create_from_physx_articulations( self = cls( _objs=physx_articulations, _scene=scene, - _scene_mask=scene_mask, + _scene_idxs=scene_idxs, links=[], link_map=OrderedDict(), root=None, @@ -158,17 +158,18 @@ def _create_from_physx_articulations( def merge(cls, articulations: List["Articulation"], name: str = None): objs = [] scene = articulations[0]._scene - merged_scene_mask = articulations[0]._scene_mask.clone() + merged_scene_idxs = [] num_objs_per_actor = articulations[0]._num_objs for articulation in articulations: objs += articulation._objs - merged_scene_mask[articulation._scene_mask] = True + merged_scene_idxs.append(articulation._scene_idxs) del scene.articulations[articulation.name] assert ( articulation._num_objs == num_objs_per_actor ), "Each given articulation must have the same number of managed objects" - merged_articulation = Articulation._create_from_physx_articulations( - objs, scene, merged_scene_mask, _merged=True + merged_scene_idxs = torch.concat(merged_scene_idxs) + merged_articulation = Articulation.create_from_physx_articulations( + objs, scene, merged_scene_idxs, _merged=True ) merged_articulation.name = name scene.articulations[merged_articulation.name] = merged_articulation diff --git a/mani_skill/utils/structs/base.py b/mani_skill/utils/structs/base.py index dfa8df4f5..7f5a7444a 100644 --- a/mani_skill/utils/structs/base.py +++ b/mani_skill/utils/structs/base.py @@ -24,14 +24,15 @@ class BaseStruct(Generic[T]): _objs: List[T] """list of objects of type T managed by this dataclass""" - _scene_mask: torch.Tensor - """a mask over all sub scenes indicating where the objects are located. - Note that torch.sum(_scene_mask) == len(_objs) and both _scene_mask and _objs are - both sequentially ordered the same: - e.g. the nth True of _scene_mask corresponds with the nth element of _objs""" + _scene_idxs: torch.Tensor + """parallel list with _objs indicating which sub-scene each of those objects are actually in by index""" _scene: ManiSkillScene """The ManiSkillScene object that manages the sub-scenes this dataclasses's objects are in""" + def __post_init__(self): + if not isinstance(self._scene_idxs, torch.Tensor): + self._scene_idxs = sapien_utils.to_tensor(self._scene_idxs) + @property def device(self): if physx.is_gpu_enabled(): diff --git a/mani_skill/utils/structs/joint.py b/mani_skill/utils/structs/joint.py index 119c680cf..dbf21e978 100644 --- a/mani_skill/utils/structs/joint.py +++ b/mani_skill/utils/structs/joint.py @@ -76,7 +76,7 @@ def create( active_index=active_joint_index, _objs=physx_joints, _scene=articulation._scene, - _scene_mask=articulation._scene_mask, + _scene_idxs=articulation._scene_idxs, child_link=child_link, parent_link=parent_link, name=shared_name, diff --git a/mani_skill/utils/structs/link.py b/mani_skill/utils/structs/link.py index 018a786df..64af15104 100644 --- a/mani_skill/utils/structs/link.py +++ b/mani_skill/utils/structs/link.py @@ -44,18 +44,18 @@ def create( cls, physx_links: List[physx.PhysxArticulationLinkComponent], articulation: Articulation = None, - scene_mask: torch.Tensor = None, + scene_idxs: torch.Tensor = None, ): shared_name = "_".join( physx_links[0].name.replace(articulation.name, "", 1).split("_")[1:] ) - if scene_mask is None and articulation is not None: - scene_mask = articulation._scene_mask + if scene_idxs is None and articulation is not None: + scene_idxs = articulation._scene_idxs return cls( articulation=articulation, _objs=physx_links, _scene=articulation._scene, - _scene_mask=scene_mask, + _scene_idxs=scene_idxs, name=shared_name, _body_data_name="cuda_rigid_body_data" if isinstance(articulation.px, physx.PhysxGpuSystem) @@ -66,18 +66,20 @@ def create( @classmethod def merge(cls, links: List["Link"], name: str = None): objs = [] - merged_scene_mask = links[0]._scene_mask.clone() + merged_scene_idxs = [] num_objs_per_actor = links[0]._num_objs for link in links: objs += link._objs - merged_scene_mask[link._scene_mask] = True + merged_scene_idxs.append(link._scene_idxs) assert ( link._num_objs == num_objs_per_actor ), "Each given link must have the same number of managed objects" + merged_scene_idxs = torch.concat(merged_scene_idxs) merged_link = Link.create( - objs, articulation=links[0].articulation, scene_mask=merged_scene_mask + objs, articulation=links[0].articulation, scene_idxs=merged_scene_idxs ) - merged_link.articulation = None # remove articulation reference as it does not make sense and is only used to instantiate some properties like the physx system + # remove articulation reference as it does not make sense and is only used to instantiate some properties like the physx system + merged_link.articulation = None merged_link.name = name return merged_link From eeaeee3486b5e669b0aa37a64cb69e596acdbead Mon Sep 17 00:00:00 2001 From: StoneT2000 Date: Tue, 12 Mar 2024 02:04:45 -0700 Subject: [PATCH 2/2] Update link.py --- mani_skill/utils/structs/link.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mani_skill/utils/structs/link.py b/mani_skill/utils/structs/link.py index 64af15104..f5fd57c16 100644 --- a/mani_skill/utils/structs/link.py +++ b/mani_skill/utils/structs/link.py @@ -160,7 +160,7 @@ def pose(self) -> Pose: ) else: assert len(self._objs) == 1 - return Pose.create(self._objs[0].pose) + return Pose.create(self._objs[0].entity_pose) @pose.setter def pose(self, arg1: Union[Pose, sapien.Pose]) -> None: