diff --git a/examples/rigid/heterogeneous_simulation.py b/examples/rigid/heterogeneous_simulation.py new file mode 100644 index 000000000..43190e897 --- /dev/null +++ b/examples/rigid/heterogeneous_simulation.py @@ -0,0 +1,111 @@ +import argparse + +import numpy as np +import genesis as gs + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("-v", "--vis", action="store_true", default=False) + parser.add_argument("-n", "--n_envs", type=int, default=4) + args = parser.parse_args() + + ########################## init ########################## + gs.init(backend=gs.gpu, precision="32") + ########################## create a scene ########################## + scene = gs.Scene( + viewer_options=gs.options.ViewerOptions( + camera_pos=(3, -1, 1.5), + camera_lookat=(0.0, 0.0, 0.5), + camera_fov=30, + max_FPS=60, + ), + sim_options=gs.options.SimOptions( + dt=0.01, + ), + show_viewer=args.vis, + ) + + ########################## entities ########################## + plane = scene.add_entity( + gs.morphs.Plane(), + ) + franka = scene.add_entity( + gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"), + ) + + morphs_heterogeneous = [ + gs.morphs.Box( + size=(0.04, 0.04, 0.04), + pos=(0.65, 0.0, 0.02), + ), + gs.morphs.Box( + size=(0.02, 0.02, 0.02), + pos=(0.65, 0.0, 0.02), + ), + gs.morphs.Sphere( + radius=0.015, + pos=(0.65, 0.0, 0.02), + ), + gs.morphs.Sphere( + radius=0.025, + pos=(0.65, 0.0, 0.02), + ), + ] + grasping_object = scene.add_entity( + morph=morphs_heterogeneous, + ) + ########################## build ########################## + scene.build(n_envs=args.n_envs, env_spacing=(1, 1)) + + motors_dof = np.arange(7) + fingers_dof = np.arange(7, 9) + l_qpos = [-1.0124, 1.5559, 1.3662, -1.6878, -1.5799, 1.7757, 1.4602, 0.04, 0.04] + if args.n_envs == 0: + franka.set_qpos(np.array(l_qpos)) + else: + franka.set_qpos(np.array([l_qpos] * args.n_envs)) + scene.step() + + AABB = grasping_object.get_AABB() + mass = grasping_object.get_mass() + print("heterogeneous AABB", AABB) + print("heterogeneous mass", mass) + + end_effector = franka.get_link("hand") + qpos = franka.inverse_kinematics( + link=end_effector, + pos=np.array([[0.65, 0.0, 0.135]] * args.n_envs), + quat=np.array([[0, 1, 0, 0]] * args.n_envs), + ) + print(qpos.shape, motors_dof.shape) + franka.control_dofs_position(qpos[..., :-2], motors_dof) + + # hold + for i in range(100): + print("hold", i) + scene.step() + + # grasp + finder_pos = -0.0 + for i in range(100): + print("grasp", i) + franka.control_dofs_position(qpos[..., :-2], motors_dof) + franka.control_dofs_position(np.array([[finder_pos, finder_pos]] * args.n_envs), fingers_dof) + scene.step() + + # lift + qpos = franka.inverse_kinematics( + link=end_effector, + pos=np.array([[0.65, 0.0, 0.3]] * args.n_envs), + quat=np.array([[0, 1, 0, 0]] * args.n_envs), + ) + for i in range(200): + print("lift", i) + franka.control_dofs_position(qpos[..., :-2], motors_dof) + franka.control_dofs_position(np.array([[finder_pos, finder_pos]] * args.n_envs), fingers_dof) + scene.step() + + +if __name__ == "__main__": + main() diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 534500938..1bbd443fb 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -63,6 +63,7 @@ def __init__( vface_start=0, equality_start=0, visualize_contact: bool = False, + morph_heterogeneous: list[Morph] | None = None, ): super().__init__(idx, scene, morph, solver, material, surface) @@ -89,25 +90,108 @@ def __init__( self._is_free: bool = morph.is_free self._is_built: bool = False + self._morph_heterogeneous = morph_heterogeneous + self._enable_heterogeneous = not (morph_heterogeneous is None or len(morph_heterogeneous) == 0) self._load_model() + def load_morph(self, morph: Morph): + if isinstance(morph, gs.morphs.Mesh): + self._load_mesh(morph, self._surface) + elif isinstance(morph, (gs.morphs.MJCF, gs.morphs.URDF, gs.morphs.Drone)): + self._load_scene(morph, self._surface) + elif isinstance(morph, gs.morphs.Primitive): + self._load_primitive(morph, self._surface) + elif isinstance(morph, gs.morphs.Terrain): + self._load_terrain(morph, self._surface) + else: + gs.raise_exception(f"Unsupported morph: {morph}.") + + def _load_heterogeneous(self): + # TODO + + # handle self._requires_jac_and_IK + # handle self._is_local_collision_mask + # handle contact pair validity + # handle max ndof, nlink, etc + # only support primitive and mesh + # com can be off + + self.list_het_link_start = gs.List() + self.list_het_link_end = gs.List() + self.list_het_n_links = gs.List() + self.list_het_geom_group_start = gs.List() + self.list_het_geom_group_end = gs.List() + self.list_het_vgeom_group_start = gs.List() + self.list_het_vgeom_group_end = gs.List() + + self.list_het_link_start.append(self._link_start) + self.list_het_n_links.append(len(self._links)) + self.list_het_link_end.append(self._link_start + len(self._links)) + self.list_het_geom_group_start.append(self._geom_start) + self.list_het_geom_group_end.append(self._geom_start + len(self.geoms)) + self.list_het_vgeom_group_start.append(self._vgeom_start) + self.list_het_vgeom_group_end.append(self._vgeom_start + len(self.vgeoms)) + + if self._enable_heterogeneous: + for morph in self._morph_heterogeneous: + if isinstance(morph, gs.morphs.Mesh): + g_infos = self._load_mesh(morph, self._surface, load_geom_only_for_heterogeneous=True) + elif isinstance(morph, gs.morphs.Primitive): + g_infos = self._load_primitive(morph, self._surface, load_geom_only_for_heterogeneous=True) + else: + gs.raise_exception( + f"morph_heterogeneous only supports primitive and mesh, Unsupported morph: {morph}." + ) + + if len(self._links) != 1: + gs.raise_exception("morph_heterogeneous only supports single link entity.") + + link = self._links[0] + cg_infos, vg_infos = self._convert_g_infos_to_cg_infos_and_vg_infos(morph, g_infos, False) + + # Add visual geometries + for g_info in vg_infos: + link._add_vgeom( + vmesh=g_info["vmesh"], + init_pos=g_info.get("pos", gu.zero_pos()), + init_quat=g_info.get("quat", gu.identity_quat()), + ) + + # Add collision geometries + for g_info in cg_infos: + friction = self.material.friction + if friction is None: + friction = g_info.get("friction", gu.default_friction()) + link._add_geom( + mesh=g_info["mesh"], + init_pos=g_info.get("pos", gu.zero_pos()), + init_quat=g_info.get("quat", gu.identity_quat()), + type=g_info["type"], + friction=friction, + sol_params=g_info["sol_params"], + data=g_info.get("data"), + needs_coup=self.material.needs_coup, + contype=g_info["contype"], + conaffinity=g_info["conaffinity"], + ) + + self.list_het_link_start.append(self.list_het_link_end[-1]) + self.list_het_link_end.append(self._link_start + len(self._links)) + self.list_het_n_links.append(self.list_het_link_end[-1] - self.list_het_link_start[-1]) + self.list_het_geom_group_start.append(self.list_het_geom_group_end[-1]) + self.list_het_geom_group_end.append(self.list_het_geom_group_end[-1] + len(cg_infos)) + self.list_het_vgeom_group_start.append(self.list_het_vgeom_group_end[-1]) + self.list_het_vgeom_group_end.append(self.list_het_vgeom_group_end[-1] + len(vg_infos)) + def _load_model(self): self._links = gs.List() self._joints = gs.List() self._equalities = gs.List() - if isinstance(self._morph, gs.morphs.Mesh): - self._load_mesh(self._morph, self._surface) - elif isinstance(self._morph, (gs.morphs.MJCF, gs.morphs.URDF, gs.morphs.Drone)): - self._load_scene(self._morph, self._surface) - elif isinstance(self._morph, gs.morphs.Primitive): - self._load_primitive(self._morph, self._surface) - elif isinstance(self._morph, gs.morphs.Terrain): - self._load_terrain(self._morph, self._surface) - else: - gs.raise_exception(f"Unsupported morph: {self._morph}.") + self.load_morph(self._morph) + self._load_heterogeneous() self._requires_jac_and_IK = self._morph.requires_jac_and_IK self._is_local_collision_mask = isinstance(self._morph, gs.morphs.MJCF) @@ -120,7 +204,7 @@ def _update_child_idxs(self): if link.idx not in parent_link.child_idxs: parent_link.child_idxs.append(link.idx) - def _load_primitive(self, morph, surface): + def _load_primitive(self, morph, surface, load_geom_only_for_heterogeneous=False): if morph.fixed: joint_type = gs.JOINT_TYPE.FIXED n_qs = 0 @@ -185,6 +269,9 @@ def _load_primitive(self, morph, surface): ) ) + if load_geom_only_for_heterogeneous: + return g_infos + link, (joint,) = self._add_by_info( l_info=dict( is_robot=False, @@ -209,7 +296,7 @@ def _load_primitive(self, morph, surface): surface=surface, ) - def _load_mesh(self, morph, surface): + def _load_mesh(self, morph, surface, load_geom_only_for_heterogeneous=False): if morph.fixed: joint_type = gs.JOINT_TYPE.FIXED n_qs = 0 @@ -255,8 +342,11 @@ def _load_mesh(self, morph, surface): ) ) - link_name = morph.file.rsplit("/", 1)[-1].replace(".", "_") + if load_geom_only_for_heterogeneous: + return g_infos + link_name = morph.file.rsplit("/", 1)[-1].replace(".", "_") + print("mesh") link, (joint,) = self._add_by_info( l_info=dict( is_robot=False, @@ -715,6 +805,37 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface): ) joints.append(joint) + cg_infos, vg_infos = self._convert_g_infos_to_cg_infos_and_vg_infos(morph, g_infos, l_info["is_robot"]) + + # Add visual geometries + for g_info in vg_infos: + link._add_vgeom( + vmesh=g_info["vmesh"], + init_pos=g_info.get("pos", gu.zero_pos()), + init_quat=g_info.get("quat", gu.identity_quat()), + ) + # Add collision geometries + for g_info in cg_infos: + friction = self.material.friction + if friction is None: + friction = g_info.get("friction", gu.default_friction()) + link._add_geom( + mesh=g_info["mesh"], + init_pos=g_info.get("pos", gu.zero_pos()), + init_quat=g_info.get("quat", gu.identity_quat()), + type=g_info["type"], + friction=friction, + sol_params=g_info["sol_params"], + data=g_info.get("data"), + needs_coup=self.material.needs_coup, + contype=g_info["contype"], + conaffinity=g_info["conaffinity"], + ) + + return link, joints + + @staticmethod + def _convert_g_infos_to_cg_infos_and_vg_infos(morph, g_infos, is_robot): # Separate collision from visual geometry for post-processing cg_infos, vg_infos = [], [] for g_info in g_infos: @@ -740,7 +861,7 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface): # the non-physical part that is added to the original geometries to convexify them are generally inside # the mechanical structure and not interacting directly with the outer world. On top of that, not only # iy increases the memory footprint and compilation time, but also the simulation speed (marginally). - if l_info["is_robot"]: + if is_robot: decompose_error_threshold = morph.decompose_robot_error_threshold else: decompose_error_threshold = morph.decompose_object_error_threshold @@ -760,33 +881,7 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface): mesh = g_info["mesh"] mesh.set_color((*np.random.rand(3), 0.7)) - # Add visual geometries - for g_info in vg_infos: - link._add_vgeom( - vmesh=g_info["vmesh"], - init_pos=g_info.get("pos", gu.zero_pos()), - init_quat=g_info.get("quat", gu.identity_quat()), - ) - - # Add collision geometries - for g_info in cg_infos: - friction = self.material.friction - if friction is None: - friction = g_info.get("friction", gu.default_friction()) - link._add_geom( - mesh=g_info["mesh"], - init_pos=g_info.get("pos", gu.zero_pos()), - init_quat=g_info.get("quat", gu.identity_quat()), - type=g_info["type"], - friction=friction, - sol_params=g_info["sol_params"], - data=g_info.get("data"), - needs_coup=self.material.needs_coup, - contype=g_info["contype"], - conaffinity=g_info["conaffinity"], - ) - - return link, joints + return cg_infos, vg_infos def _add_equality(self, name, type, objs_name, data, sol_params): objs_id = [] @@ -1891,11 +1986,31 @@ def get_AABB(self): if self.n_geoms == 0: gs.raise_exception("Entity has no geoms.") - verts = self.get_verts() - AABB = torch.concatenate( - [verts.min(axis=-2, keepdim=True)[0], verts.max(axis=-2, keepdim=True)[0]], - axis=-2, + gs.engine.solvers.rigid.rigid_solver_decomp.kernel_update_geom_aabbs( + self._solver.geoms_state, + self._solver.geoms_init_AABB, + self._solver._static_rigid_sim_config, + ) + + links_aabb = torch.empty((self._solver._B, self.n_links, 6), dtype=gs.tc_float, device=gs.device) + # links_idx is self.link_start to self.link_end + links_idx = torch.arange(start=self.link_start, end=self.link_end, dtype=gs.tc_int, device=gs.device) + envs_idx = torch.arange(self._solver._B, dtype=gs.tc_int, device=gs.device) + gs.engine.solvers.rigid.rigid_solver_decomp.kernel_get_links_AABB( + links_aabb, + links_idx, + envs_idx, + self._solver.links_info, + self._solver.geoms_state, + self._solver._static_rigid_sim_config, ) + + AABB = torch.empty((self._solver._B, 2, 3), dtype=gs.tc_float, device=gs.device) + AABB[:, 0] = links_aabb[:, :, :3].min(dim=1)[0] + AABB[:, 1] = links_aabb[:, :, 3:].max(dim=1)[0] + if self._solver.n_envs == 0: + AABB = AABB.squeeze(0) + return AABB def _get_idx(self, idx_local, idx_local_max, idx_global_start=0, *, unsafe=False): @@ -2607,10 +2722,10 @@ def get_mass(self): mass : float The total mass of the entity in kg. """ - mass = 0.0 - for link in self.links: - mass += link.get_mass() - return mass + all_links_mass = self._solver.links_info.inertial_mass.to_numpy() + links_idx = np.arange(self.link_start, self.link_end) + entity_mass = all_links_mass[links_idx].sum(0) + return entity_mass # ------------------------------------------------------------------------------------ # ----------------------------------- properties ------------------------------------- diff --git a/genesis/engine/entities/rigid_entity/rigid_link.py b/genesis/engine/entities/rigid_entity/rigid_link.py index 159b59fcc..9a3bc4109 100644 --- a/genesis/engine/entities/rigid_entity/rigid_link.py +++ b/genesis/engine/entities/rigid_entity/rigid_link.py @@ -94,6 +94,45 @@ def __init__( self._geoms: list[RigidGeom] = gs.List() self._vgeoms: list[RigidVisGeom] = gs.List() + @staticmethod + def compute_inertial_info(init_mesh, inertial_quat, rho): + if init_mesh is None: + inertial_mass = 0.0 + else: + if init_mesh.is_watertight: + inertial_mass = init_mesh.volume * rho + else: # TODO: handle non-watertight mesh + inertial_mass = 1.0 + + # inertial_pos + if init_mesh is None: + inertial_pos = gu.zero_pos() + else: + inertial_pos = np.array(init_mesh.center_mass, dtype=gs.np_float) + + # inertial_i + # FIXME: Why coef 0.4 ??? + if init_mesh is None: # use sphere inertia with radius 0.1 + inertial_i = 0.4 * inertial_mass * 0.1**2 * np.eye(3) + + else: + # attempt to fix non-watertight mesh by convexifying + inertia_mesh = init_mesh.copy() + if not inertia_mesh.is_watertight: + inertia_mesh = trimesh.convex.convex_hull(inertia_mesh) + + if inertia_mesh.is_watertight and inertia_mesh.mass > 0: + # TODO: check if this is correct. This is correct if the inertia frame is w.r.t to link frame + T_inertia = gu.trans_quat_to_T(inertial_pos, inertial_quat) + inertial_i = inertia_mesh.moment_inertia_frame(T_inertia) / inertia_mesh.mass * inertial_mass + + else: # approximate with a sphere + inertial_i = ( + 0.4 * inertial_mass * (max(inertia_mesh.bounds[1] - inertia_mesh.bounds[0]) / 2.0) ** 2 * np.eye(3) + ) + + return inertial_mass, inertial_pos, inertial_i + def _build(self): for geom in self._geoms: geom._build() @@ -101,7 +140,7 @@ def _build(self): for vgeom in self._vgeoms: vgeom._build() - self._init_mesh = self._compose_init_mesh() + init_mesh = self._compose_init_mesh(self._geoms, self._vgeoms) # find root link and check if link is fixed solver_links = self._solver.links @@ -115,56 +154,18 @@ def _build(self): self._root_idx = gs.np_int(link.idx) self.is_fixed = is_fixed - # inertial_mass and inertia_i - if self._inertial_mass is None: - if len(self._geoms) == 0 and len(self._vgeoms) == 0: - self._inertial_mass = 0.0 - else: - if self._init_mesh.is_watertight: - self._inertial_mass = self._init_mesh.volume * self.entity.material.rho - else: # TODO: handle non-watertight mesh - self._inertial_mass = 1.0 + mesh_inertial_mass, mesh_inertial_pos, mesh_inertial_i = self.compute_inertial_info( + init_mesh, self._inertial_quat, self.entity.material.rho + ) + self._inertial_mass = mesh_inertial_mass if self._inertial_mass is None else self._inertial_mass + self._inertial_pos = mesh_inertial_pos if self._inertial_pos is None else self._inertial_pos + self._inertial_i = mesh_inertial_i if self._inertial_i is None else self._inertial_i + self._inertial_i = np.asarray(self._inertial_i, dtype=gs.np_float) # Postpone computation of inverse weight if not specified if self._invweight is None: self._invweight = np.full((2,), fill_value=-1.0, dtype=gs.np_float) - # inertial_pos - if self._inertial_pos is None: - if self._init_mesh is None: - self._inertial_pos = gu.zero_pos() - else: - self._inertial_pos = np.array(self._init_mesh.center_mass, dtype=gs.np_float) - - # inertial_i - if self._inertial_i is None: - # FIXME: Why coef 0.4 ??? - if self._init_mesh is None: # use sphere inertia with radius 0.1 - self._inertial_i = 0.4 * self._inertial_mass * 0.1**2 * np.eye(3) - - else: - # attempt to fix non-watertight mesh by convexifying - inertia_mesh = self._init_mesh.copy() - if not inertia_mesh.is_watertight: - inertia_mesh = trimesh.convex.convex_hull(inertia_mesh) - - if inertia_mesh.is_watertight and self._init_mesh.mass > 0: - # TODO: check if this is correct. This is correct if the inertia frame is w.r.t to link frame - T_inertia = gu.trans_quat_to_T(self._inertial_pos, self._inertial_quat) - self._inertial_i = ( - self._init_mesh.moment_inertia_frame(T_inertia) / self._init_mesh.mass * self._inertial_mass - ) - - else: # approximate with a sphere - self._inertial_i = ( - 0.4 - * self._inertial_mass - * (max(self._init_mesh.bounds[1] - self._init_mesh.bounds[0]) / 2.0) ** 2 - * np.eye(3) - ) - - self._inertial_i = np.asarray(self._inertial_i, dtype=gs.np_float) - # override invweight if fixed if is_fixed: self._invweight = np.zeros((2,), dtype=gs.np_float) @@ -173,20 +174,21 @@ def _build(self): self.rsd = rigid_solver_decomp - def _compose_init_mesh(self): - if len(self._geoms) == 0 and len(self._vgeoms) == 0: + @staticmethod + def _compose_init_mesh(geoms, vgeoms): + if len(geoms) == 0 and len(vgeoms) == 0: return None else: init_verts = [] init_faces = [] vert_offset = 0 - if len(self._geoms) > 0: - for geom in self._geoms: + if len(geoms) > 0: + for geom in geoms: init_verts.append(gu.transform_by_trans_quat(geom.init_verts, geom.init_pos, geom.init_quat)) init_faces.append(geom.init_faces + vert_offset) vert_offset += geom.init_verts.shape[0] - elif len(self._vgeoms) > 0: # use vgeom if there's no geom - for vgeom in self._vgeoms: + elif len(vgeoms) > 0: # use vgeom if there's no geom + for vgeom in vgeoms: init_verts.append(gu.transform_by_trans_quat(vgeom.init_vverts, vgeom.init_pos, vgeom.init_quat)) init_faces.append(vgeom.init_vfaces + vert_offset) vert_offset += vgeom.init_vverts.shape[0] diff --git a/genesis/engine/scene.py b/genesis/engine/scene.py index ad3e7c6cc..bdeb8550d 100644 --- a/genesis/engine/scene.py +++ b/genesis/engine/scene.py @@ -273,7 +273,7 @@ def _validate_options( @gs.assert_unbuilt def add_entity( self, - morph: Morph, + morph: Morph | list[Morph], material: Material | None = None, surface: Surface | None = None, visualize_contact: bool = False, @@ -284,8 +284,8 @@ def add_entity( Parameters ---------- - morph : gs.morphs.Morph - The morph of the entity. + morph : gs.morphs.Morph | list[gs.morphs.Morph] + The morph of the entity. If a list of morphs is provided, the entity will be heterogeneous (rigid only). material : gs.materials.Material | None, optional The material of the entity. If None, use ``gs.materials.Rigid()``. surface : gs.surfaces.Surface | None, optional @@ -311,6 +311,9 @@ def add_entity( # small sdf res is sufficient for primitives regardless of size if isinstance(morph, gs.morphs.Primitive): material._sdf_max_res = 32 + else: + if isinstance(morph, list): + gs.raise_exception("Heterogeneous morphs are not supported for non-rigid materials.") # some morph should not smooth surface normal if isinstance(morph, (gs.morphs.Box, gs.morphs.Cylinder, gs.morphs.Terrain)): diff --git a/genesis/engine/simulator.py b/genesis/engine/simulator.py index c69a3a8ab..96cb262be 100644 --- a/genesis/engine/simulator.py +++ b/genesis/engine/simulator.py @@ -155,7 +155,14 @@ def __init__( # sensors self._sensor_manager = SensorManager(self) - def _add_entity(self, morph: Morph, material, surface, visualize_contact=False): + def _add_entity( + self, + morph: Morph | list[Morph], + material, + surface, + visualize_contact=False, + morph_heterogeneous: list[Morph] | None = None, + ): if isinstance(material, gs.materials.Tool): entity = self.tool_solver.add_entity(self.n_entities, material, morph, surface) diff --git a/genesis/engine/solvers/rigid/collider_decomp.py b/genesis/engine/solvers/rigid/collider_decomp.py index 42710ab1f..122af862b 100644 --- a/genesis/engine/solvers/rigid/collider_decomp.py +++ b/genesis/engine/solvers/rigid/collider_decomp.py @@ -201,7 +201,11 @@ def _compute_collision_pair_validity(self): # contype and conaffinity if ( (i_ea == i_eb) - or not (entities_is_local_collision_mask[i_ea] or entities_is_local_collision_mask[i_eb]) + # entities_is_local_collision_mask[i_ea].all() becaues it could be batched. + # Consider using same ndim for batch/non-batch arrays? + or not ( + entities_is_local_collision_mask[i_ea].all() or entities_is_local_collision_mask[i_eb].all() + ) ) and not ( (geoms_contype[i_ga] & geoms_conaffinity[i_gb]) or (geoms_contype[i_gb] & geoms_conaffinity[i_ga]) ): @@ -610,7 +614,6 @@ def func_point_in_geom_aabb( i_g, i_b, geoms_state: array_class.GeomsState, - geoms_info: array_class.GeomsInfo, point: ti.types.vector(3), ): return (point < geoms_state.aabb_max[i_g, i_b]).all() and (point > geoms_state.aabb_min[i_g, i_b]).all() @@ -622,7 +625,6 @@ def func_is_geom_aabbs_overlap( i_gb, i_b, geoms_state: array_class.GeomsState, - geoms_info: array_class.GeomsInfo, ): return not ( (geoms_state.aabb_max[i_ga, i_b] <= geoms_state.aabb_min[i_gb, i_b]).any() @@ -636,7 +638,6 @@ def func_find_intersect_midpoint( i_gb, i_b, geoms_state: array_class.GeomsState, - geoms_info: array_class.GeomsInfo, ): # return the center of the intersecting AABB of AABBs of two geoms intersect_lower = ti.max(geoms_state.aabb_min[i_ga, i_b], geoms_state.aabb_min[i_gb, i_b]) @@ -695,7 +696,7 @@ def func_contact_vertex_sdf( for i_v in range(geoms_info.vert_start[i_ga], geoms_info.vert_end[i_ga]): vertex_pos = gu.ti_transform_by_trans_quat(verts_info.init_pos[i_v], ga_pos, ga_quat) - if func_point_in_geom_aabb(i_gb, i_b, geoms_state, geoms_info, vertex_pos): + if func_point_in_geom_aabb(i_gb, i_b, geoms_state, vertex_pos): new_penetration = -sdf.sdf_func_world(geoms_state, geoms_info, sdf_info, vertex_pos, i_gb, i_b) if new_penetration > penetration: is_col = True @@ -1160,31 +1161,39 @@ def func_broad_phase( """ _B = collider_state.active_buffer.shape[1] n_geoms = collider_state.active_buffer.shape[0] + n_links = links_info.geom_start.shape[0] ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) for i_b in range(_B): axis = 0 - + env_n_geoms = 0 + for i_l in range(n_links): + I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l + env_n_geoms = env_n_geoms + links_info.geom_end[I_l] - links_info.geom_start[I_l] # copy updated geom aabbs to buffer for sorting if collider_state.first_time[i_b]: - for i in range(n_geoms): - collider_state.sort_buffer.value[2 * i, i_b] = geoms_state.aabb_min[i, i_b][axis] - collider_state.sort_buffer.i_g[2 * i, i_b] = i - collider_state.sort_buffer.is_max[2 * i, i_b] = 0 - - collider_state.sort_buffer.value[2 * i + 1, i_b] = geoms_state.aabb_max[i, i_b][axis] - collider_state.sort_buffer.i_g[2 * i + 1, i_b] = i - collider_state.sort_buffer.is_max[2 * i + 1, i_b] = 1 - - geoms_state.min_buffer_idx[i, i_b] = 2 * i - geoms_state.max_buffer_idx[i, i_b] = 2 * i + 1 + i_buffer = 0 + for i_l in range(n_links): + I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l + for i_g in range(links_info.geom_start[I_l], links_info.geom_end[I_l]): + collider_state.sort_buffer.value[2 * i_buffer, i_b] = geoms_state.aabb_min[i_g, i_b][axis] + collider_state.sort_buffer.i_g[2 * i_buffer, i_b] = i_g + collider_state.sort_buffer.is_max[2 * i_buffer, i_b] = 0 + + collider_state.sort_buffer.value[2 * i_buffer + 1, i_b] = geoms_state.aabb_max[i_g, i_b][axis] + collider_state.sort_buffer.i_g[2 * i_buffer + 1, i_b] = i_g + collider_state.sort_buffer.is_max[2 * i_buffer + 1, i_b] = 1 + + geoms_state.min_buffer_idx[i_buffer, i_b] = 2 * i_g + geoms_state.max_buffer_idx[i_buffer, i_b] = 2 * i_g + 1 + i_buffer = i_buffer + 1 collider_state.first_time[i_b] = False else: # warm start. If `use_hibernation=True`, it's already updated in rigid_solver. if ti.static(not static_rigid_sim_config.use_hibernation): - for i in range(n_geoms * 2): + for i in range(env_n_geoms * 2): if collider_state.sort_buffer.is_max[i, i_b]: collider_state.sort_buffer.value[i, i_b] = geoms_state.aabb_max[ collider_state.sort_buffer.i_g[i, i_b], i_b @@ -1195,7 +1204,7 @@ def func_broad_phase( ][axis] # insertion sort, which has complexity near O(n) for nearly sorted array - for i in range(1, 2 * n_geoms): + for i in range(1, 2 * env_n_geoms): key_value = collider_state.sort_buffer.value[i, i_b] key_is_max = collider_state.sort_buffer.is_max[i, i_b] key_i_g = collider_state.sort_buffer.i_g[i, i_b] @@ -1227,7 +1236,7 @@ def func_broad_phase( collider_state.n_broad_pairs[i_b] = 0 if ti.static(not static_rigid_sim_config.use_hibernation): n_active = 0 - for i in range(2 * n_geoms): + for i in range(2 * env_n_geoms): if not collider_state.sort_buffer.is_max[i, i_b]: for j in range(n_active): i_ga = collider_state.active_buffer[j, i_b] @@ -1249,7 +1258,7 @@ def func_broad_phase( ): continue - if not func_is_geom_aabbs_overlap(i_ga, i_gb, i_b, geoms_state, geoms_info): + if not func_is_geom_aabbs_overlap(i_ga, i_gb, i_b, geoms_state): # Clear collision normal cache if not in contact if ti.static(not static_rigid_sim_config.enable_mujoco_compatibility): # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1 @@ -1279,7 +1288,7 @@ def func_broad_phase( if rigid_global_info.n_awake_dofs[i_b] > 0: n_active_awake = 0 n_active_hib = 0 - for i in range(2 * n_geoms): + for i in range(2 * env_n_geoms): is_incoming_geom_hibernated = geoms_state.hibernated[collider_state.sort_buffer.i_g[i, i_b], i_b] if not collider_state.sort_buffer.is_max[i, i_b]: @@ -1304,7 +1313,7 @@ def func_broad_phase( ): continue - if not func_is_geom_aabbs_overlap(i_ga, i_gb, i_b, geoms_state, geoms_info): + if not func_is_geom_aabbs_overlap(i_ga, i_gb, i_b, geoms_state): # Clear collision normal cache if not in contact if ti.static(not static_rigid_sim_config.enable_mujoco_compatibility): # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1 @@ -1339,7 +1348,7 @@ def func_broad_phase( ): continue - if not func_is_geom_aabbs_overlap(i_ga, i_gb, i_b, geoms_state, geoms_info): + if not func_is_geom_aabbs_overlap(i_ga, i_gb, i_b, geoms_state): # Clear collision normal cache if not in contact # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1 collider_state.contact_cache.normal[i_ga, i_gb, i_b] = ti.Vector.zero( diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp.py b/genesis/engine/solvers/rigid/constraint_solver_decomp.py index 9daf94610..afb867ca6 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp.py @@ -1221,8 +1221,9 @@ def func_nt_hessian_direct( constraint_state.nt_H[i_d1, i_d2, i_b] = constraint_state.nt_H[i_d2, i_d1, i_b] for i_e in range(n_entities): - for i_d1 in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): - for i_d2 in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_d1 in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): + for i_d2 in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): constraint_state.nt_H[i_d1, i_d2, i_b] = ( constraint_state.nt_H[i_d1, i_d2, i_b] + rigid_global_info.mass_mat[i_d1, i_d2, i_b] ) @@ -1390,9 +1391,10 @@ def func_ls_init( n_entities = entities_info.dof_start.shape[0] # mv and jv for i_e in range(n_entities): - for i_d1 in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_d1 in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): mv = gs.ti_float(0.0) - for i_d2 in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + for i_d2 in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): mv += rigid_global_info.mass_mat[i_d1, i_d2, i_b] * constraint_state.search[i_d2, i_b] constraint_state.mv[i_d1, i_b] = mv @@ -1974,10 +1976,11 @@ def initialize_Ma( n_entities = entities_info.n_links.shape[0] ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e, i_b in ti.ndrange(n_entities, _B): - for i_d1_ in range(entities_info.n_dofs[i_e]): - i_d1 = entities_info.dof_start[i_e] + i_d1_ + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_d1_ in range(entities_info.n_dofs[I_e]): + i_d1 = entities_info.dof_start[I_e] + i_d1_ Ma_ = gs.ti_float(0.0) - for i_d2 in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + for i_d2 in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): Ma_ += rigid_global_info.mass_mat[i_d1, i_d2, i_b] * qacc[i_d2, i_b] Ma[i_d1, i_b] = Ma_ diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py index a8645be0c..25564634f 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py @@ -239,6 +239,7 @@ def add_collision_constraints_and_wakeup_entities(self, i_island: int, i_b: int) self._solver.geoms_state, self._solver.data_manager.rigid_global_info, self.contact_island, + self._solver._static_rigid_sim_config, ) @ti.func diff --git a/genesis/engine/solvers/rigid/mpr_decomp.py b/genesis/engine/solvers/rigid/mpr_decomp.py index c10dad24f..ab7fb4604 100644 --- a/genesis/engine/solvers/rigid/mpr_decomp.py +++ b/genesis/engine/solvers/rigid/mpr_decomp.py @@ -48,27 +48,6 @@ def clear(mpr_state: ti.template()): mpr_state.simplex_size.fill(0) -@ti.func -def func_point_in_geom_aabb(geoms_state: array_class.GeomsState, point, i_g, i_b): - return (point < geoms_state.aabb_max[i_g, i_b]).all() and (point > geoms_state.aabb_min[i_g, i_b]).all() - - -@ti.func -def func_is_geom_aabbs_overlap(geoms_state: array_class.GeomsState, i_ga, i_gb, i_b): - return not ( - (geoms_state.aabb_max[i_ga, i_b] <= geoms_state.aabb_min[i_gb, i_b]).any() - or (geoms_state.aabb_min[i_ga, i_b] >= geoms_state.aabb_max[i_gb, i_b]).any() - ) - - -@ti.func -def func_find_intersect_midpoint(geoms_state: array_class.GeomsState, i_ga, i_gb, i_b): - # return the center of the intersecting AABB of AABBs of two geoms - intersect_lower = ti.max(geoms_state.aabb_min[i_ga, i_b], geoms_state.aabb_min[i_gb, i_b]) - intersect_upper = ti.min(geoms_state.aabb_max[i_ga, i_b], geoms_state.aabb_max[i_gb, i_b]) - return 0.5 * (intersect_lower + intersect_upper) - - @ti.func def mpr_swap(mpr_state: array_class.MPRState, i, j, i_ga, i_gb, i_b): mpr_state.simplex_support.v1[i, i_b], mpr_state.simplex_support.v1[j, i_b] = ( diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index 95c7b2061..1fc11b730 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -15,6 +15,7 @@ from genesis.engine.solvers.rigid.contact_island import ContactIsland from genesis.engine.states.solvers import RigidSolverState from genesis.options.solvers import RigidOptions +from genesis.options.morphs import Morph from genesis.styles import colors, formats from genesis.utils import linalg as lu from genesis.utils.misc import ti_field_to_torch, DeprecationError, ALLOCATE_TENSOR_WARNING @@ -153,7 +154,7 @@ def __init__(self, scene: "Scene", sim: "Simulator", options: RigidOptions) -> N self._cur_step = -1 - def add_entity(self, idx, material, morph, surface, visualize_contact) -> Entity: + def add_entity(self, idx, material, morph: Morph | list[Morph], surface, visualize_contact) -> Entity: if isinstance(material, gs.materials.Avatar): EntityClass = AvatarEntity if visualize_contact: @@ -163,6 +164,9 @@ def add_entity(self, idx, material, morph, surface, visualize_contact) -> Entity else: EntityClass = RigidEntity + morph_heterogeneous = morph[1:] if isinstance(morph, list) else None + morph = morph[0] if isinstance(morph, list) else morph + if morph.is_free: verts_state_start = self.n_free_verts else: @@ -192,6 +196,7 @@ def add_entity(self, idx, material, morph, surface, visualize_contact) -> Entity vvert_start=self.n_vverts, vface_start=self.n_vfaces, visualize_contact=visualize_contact, + morph_heterogeneous=morph_heterogeneous, ) assert isinstance(entity, RigidEntity) self._entities.append(entity) @@ -207,6 +212,8 @@ def build(self): for entity in self._entities: entity._build() + self._enable_heterogeneous = any(entity._enable_heterogeneous for entity in self._entities) + self._n_qs = self.n_qs self._n_dofs = self.n_dofs self._n_links = self.n_links @@ -258,42 +265,50 @@ def build(self): self.n_equalities_candidate = max(1, self.n_equalities + self._options.max_dynamic_constraints) - # Note optional hibernation_threshold_acc/vel params at the bottom of the initialization list. - # This is caused by this code being also run by AvatarSolver, which inherits from this class - # but does not have all the attributes of the base class. - self._static_rigid_sim_config = self.StaticRigidSimConfig( - para_level=self.sim._para_level, - use_hibernation=getattr(self, "_use_hibernation", False), - use_contact_island=getattr(self, "_use_contact_island", False), - batch_links_info=getattr(self._options, "batch_links_info", False), - batch_dofs_info=getattr(self._options, "batch_dofs_info", False), - batch_joints_info=getattr(self._options, "batch_joints_info", False), - enable_mujoco_compatibility=getattr(self, "_enable_mujoco_compatibility", False), - enable_multi_contact=getattr(self, "_enable_multi_contact", True), - enable_self_collision=getattr(self, "_enable_self_collision", True), - enable_adjacent_collision=getattr(self, "_enable_adjacent_collision", False), - enable_collision=getattr(self, "_enable_collision", False), - box_box_detection=getattr(self, "_box_box_detection", False), - integrator=getattr(self, "_integrator", gs.integrator.implicitfast), - sparse_solve=getattr(self._options, "sparse_solve", False), - solver_type=getattr(self._options, "constraint_solver", gs.constraint_solver.CG), - # dynamic properties - substep_dt=self._substep_dt, - iterations=getattr(self._options, "iterations", 10), - tolerance=getattr(self._options, "tolerance", 1e-6), - ls_iterations=getattr(self._options, "ls_iterations", 10), - ls_tolerance=getattr(self._options, "ls_tolerance", 1e-6), - n_equalities=self._n_equalities, - n_equalities_candidate=self.n_equalities_candidate, - hibernation_thresh_acc=getattr(self, "_hibernation_thresh_acc", 0.0), - hibernation_thresh_vel=getattr(self, "_hibernation_thresh_vel", 0.0), - ) + if self.is_active(): + if hasattr(self._options, "batch_entities_info"): + self._options.batch_entities_info = self._options.batch_entities_info or self._enable_heterogeneous + self._options.batch_links_info = self._options.batch_links_info or self._enable_heterogeneous + # self._options.batch_dofs_info = self._options.batch_dofs_info or self._enable_heterogeneous + # self._options.batch_joints_info = self._options.batch_joints_info or self._enable_heterogeneous + + # Note optional hibernation_threshold_acc/vel params at the bottom of the initialization list. + # This is caused by this code being also run by AvatarSolver, which inherits from this class + # but does not have all the attributes of the base class. + self._static_rigid_sim_config = self.StaticRigidSimConfig( + para_level=self.sim._para_level, + use_hibernation=getattr(self, "_use_hibernation", False), + use_contact_island=getattr(self, "_use_contact_island", False), + batch_entities_info=getattr(self._options, "batch_entities_info", False), + batch_links_info=getattr(self._options, "batch_links_info", False), + batch_dofs_info=getattr(self._options, "batch_dofs_info", False), + batch_joints_info=getattr(self._options, "batch_joints_info", False), + enable_heterogeneous=getattr(self, "_enable_heterogeneous", False), + enable_mujoco_compatibility=getattr(self, "_enable_mujoco_compatibility", False), + enable_multi_contact=getattr(self, "_enable_multi_contact", True), + enable_self_collision=getattr(self, "_enable_self_collision", True), + enable_adjacent_collision=getattr(self, "_enable_adjacent_collision", False), + enable_collision=getattr(self, "_enable_collision", False), + box_box_detection=getattr(self, "_box_box_detection", False), + integrator=getattr(self, "_integrator", gs.integrator.implicitfast), + sparse_solve=getattr(self._options, "sparse_solve", False), + solver_type=getattr(self._options, "constraint_solver", gs.constraint_solver.CG), + # dynamic properties + substep_dt=self._substep_dt, + iterations=getattr(self._options, "iterations", 10), + tolerance=getattr(self._options, "tolerance", 1e-6), + ls_iterations=getattr(self._options, "ls_iterations", 10), + ls_tolerance=getattr(self._options, "ls_tolerance", 1e-6), + n_equalities=self._n_equalities, + n_equalities_candidate=self.n_equalities_candidate, + hibernation_thresh_acc=getattr(self, "_hibernation_thresh_acc", 0.0), + hibernation_thresh_vel=getattr(self, "_hibernation_thresh_vel", 0.0), + ) - # when the migration is finished, we will remove the about two lines - self._func_vel_at_point = func_vel_at_point - self._func_apply_external_force = func_apply_external_force + # when the migration is finished, we will remove the about two lines + self._func_vel_at_point = func_vel_at_point + self._func_apply_external_force = func_apply_external_force - if self.is_active(): self.data_manager = array_class.DataManager(self) self._rigid_global_info = self.data_manager.rigid_global_info @@ -589,11 +604,103 @@ def _init_dof_fields(self): # just in case self.dofs_state.force.fill(0) + def process_heterogeneous_link_info( + self, + links_geom_start, + links_geom_end, + links_vgeom_start, + links_vgeom_end, + links_inertial_pos, + links_inertial_quat, + links_inertial_i, + links_inertial_mass, + ): + for i_l, link in enumerate(self.links): + if link.entity._enable_heterogeneous: + # handle the mapping from env to het morph + np_geom_group_start = np.array(link._entity.list_het_geom_group_start) + np_geom_group_end = np.array(link._entity.list_het_geom_group_end) + np_vgeom_group_start = np.array(link._entity.list_het_vgeom_group_start) + np_vgeom_group_end = np.array(link._entity.list_het_vgeom_group_end) + n_het = len(np_geom_group_start) + base = self._B // n_het + extra = self._B % n_het # first `extra` chunks get one more + # dispatch geoms equally for each heterogeneous environment + sizes = np.r_[np.full(extra, base + 1), np.full(n_het - extra, base)] + geom_idx = np.repeat(np.arange(n_het), sizes) + + links_geom_start[i_l, :] = np_geom_group_start[geom_idx] + links_geom_end[i_l, :] = np_geom_group_end[geom_idx] + links_vgeom_start[i_l, :] = np_vgeom_group_start[geom_idx] + links_vgeom_end[i_l, :] = np_vgeom_group_end[geom_idx] + if self._B < n_het: + gs.raise_exception( + f"{link.name}: Batch size {self._B} must be greater than or equal to " + f"the number of heterogeneous environments {n_het}." + ) + + for geom in link.geoms: + geom.rendered_envs_idx = np.where( + (links_geom_start[i_l] <= geom.idx) & (geom.idx < links_geom_end[i_l]) + )[0] + + for geom in link.vgeoms: + geom.rendered_envs_idx = np.where( + (links_vgeom_start[i_l] <= geom.idx) & (geom.idx < links_vgeom_end[i_l]) + )[0] + + # move the links inertial info for + for i_het in range(n_het): + init_mesh = link._compose_init_mesh( + self.geoms[np_geom_group_start[i_het] : np_geom_group_end[i_het]], + [], + ) + mesh_inertial_mass, mesh_inertial_pos, mesh_inertial_i = link.compute_inertial_info( + init_mesh, link._inertial_quat, link.entity.material.rho + ) + + pos = np.array([link.inertial_pos for link in self.links], dtype=gs.np_float) + links_inertial_pos[i_l, np.where(geom_idx == i_het)[0]] = mesh_inertial_pos + links_inertial_i[i_l, np.where(geom_idx == i_het)[0]] = mesh_inertial_i + links_inertial_mass[i_l, np.where(geom_idx == i_het)[0]] = mesh_inertial_mass + def _init_link_fields(self): self.links_info = self.data_manager.links_info self.links_state = self.data_manager.links_state - links = self.links + + # dispatch geoms for heterogeneous simulation + links_geom_start = np.array([link.geom_start for link in links], dtype=gs.np_int) + links_geom_end = np.array([link.geom_end for link in links], dtype=gs.np_int) + links_vgeom_start = np.array([link.vgeom_start for link in links], dtype=gs.np_int) + links_vgeom_end = np.array([link.vgeom_end for link in links], dtype=gs.np_int) + links_inertial_pos = np.array([link.inertial_pos for link in links], dtype=gs.np_float) + links_inertial_quat = np.array([link.inertial_quat for link in links], dtype=gs.np_float) + links_inertial_i = np.array([link.inertial_i for link in links], dtype=gs.np_float) + links_inertial_mass = np.array([link.inertial_mass for link in links], dtype=gs.np_float) + if self._options.batch_links_info: + # add batch dim, nlinks x self._B + links_geom_start = np.stack([links_geom_start] * self._B, axis=-1) + links_geom_end = np.stack([links_geom_end] * self._B, axis=-1) + links_vgeom_start = np.stack([links_vgeom_start] * self._B, axis=-1) + links_vgeom_end = np.stack([links_vgeom_end] * self._B, axis=-1) + links_inertial_pos = np.stack([links_inertial_pos] * self._B, axis=1) # easier for access in init function + links_inertial_quat = np.stack([links_inertial_quat] * self._B, axis=1) + links_inertial_i = np.stack([links_inertial_i] * self._B, axis=1) + links_inertial_mass = np.stack([links_inertial_mass] * self._B, axis=1) + + if self._enable_heterogeneous: + self.process_heterogeneous_link_info( + links_geom_start, + links_geom_end, + links_vgeom_start, + links_vgeom_end, + links_inertial_pos, + links_inertial_quat, + links_inertial_i, + links_inertial_mass, + ) + kernel_init_link_fields( links_parent_idx=np.array([link.parent_idx for link in links], dtype=gs.np_int), links_root_idx=np.array([link.root_idx for link in links], dtype=gs.np_int), @@ -602,15 +709,19 @@ def _init_link_fields(self): links_joint_start=np.array([link.joint_start for link in links], dtype=gs.np_int), links_q_end=np.array([link.q_end for link in links], dtype=gs.np_int), links_dof_end=np.array([link.dof_end for link in links], dtype=gs.np_int), + links_geom_start=links_geom_start, + links_geom_end=links_geom_end, + links_vgeom_start=links_vgeom_start, + links_vgeom_end=links_vgeom_end, links_joint_end=np.array([link.joint_end for link in links], dtype=gs.np_int), links_invweight=np.array([link.invweight for link in links], dtype=gs.np_float), links_is_fixed=np.array([link.is_fixed for link in links], dtype=gs.np_bool), links_pos=np.array([link.pos for link in links], dtype=gs.np_float), links_quat=np.array([link.quat for link in links], dtype=gs.np_float), - links_inertial_pos=np.array([link.inertial_pos for link in links], dtype=gs.np_float), - links_inertial_quat=np.array([link.inertial_quat for link in links], dtype=gs.np_float), - links_inertial_i=np.array([link.inertial_i for link in links], dtype=gs.np_float), - links_inertial_mass=np.array([link.inertial_mass for link in links], dtype=gs.np_float), + links_inertial_pos=links_inertial_pos, + links_inertial_quat=links_inertial_quat, + links_inertial_i=links_inertial_i, + links_inertial_mass=links_inertial_mass, links_entity_idx=np.array([link._entity_idx_in_solver for link in links], dtype=gs.np_int), # taichi variables links_info=self.links_info, @@ -2124,6 +2235,21 @@ def get_links_root_COM(self, links_idx=None, envs_idx=None, *, unsafe=False): tensor = ti_field_to_torch(self.links_state.COM, envs_idx, links_idx, transpose=True, unsafe=unsafe) return tensor.squeeze(0) if self.n_envs == 0 else tensor + def get_links_AABB(self, links_idx=None, envs_idx=None, *, unsafe=False): + _tensor, links_idx, envs_idx = self._sanitize_2D_io_variables( + None, links_idx, self.n_links, 3, envs_idx, idx_name="links_idx", unsafe=unsafe + ) + tensor = _tensor.unsqueeze(0) if self.n_envs == 0 else _tensor + kernel_get_links_AABB( + tensor, + links_idx, + envs_idx, + self.links_info, + self.geoms_state, + self._static_rigid_sim_config, + ) + return _tensor + def get_links_mass_shift(self, links_idx=None, envs_idx=None, *, unsafe=False): tensor = ti_field_to_torch(self.links_state.mass_shift, envs_idx, links_idx, transpose=True, unsafe=unsafe) return tensor.squeeze(0) if self.n_envs == 0 else tensor @@ -2593,7 +2719,8 @@ def kernel_init_meaninertia( if n_dofs > 0: rigid_global_info.meaninertia[i_b] = 0.0 for i_e in range(n_entities): - for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_d in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): I_d = [i_d, i_b] if ti.static(static_rigid_sim_config.batch_dofs_info) else i_d rigid_global_info.meaninertia[i_b] += rigid_global_info.mass_mat[i_d, i_d, i_b] rigid_global_info.meaninertia[i_b] = rigid_global_info.meaninertia[i_b] / n_dofs @@ -2667,6 +2794,10 @@ def kernel_init_link_fields( links_joint_start: ti.types.ndarray(), links_q_end: ti.types.ndarray(), links_dof_end: ti.types.ndarray(), + links_geom_start: ti.types.ndarray(), + links_geom_end: ti.types.ndarray(), + links_vgeom_start: ti.types.ndarray(), + links_vgeom_end: ti.types.ndarray(), links_joint_end: ti.types.ndarray(), links_invweight: ti.types.ndarray(), links_is_fixed: ti.types.ndarray(), @@ -2705,16 +2836,21 @@ def kernel_init_link_fields( for j in ti.static(range(4)): links_info.quat[I][j] = links_quat[i, j] - links_info.inertial_quat[I][j] = links_inertial_quat[i, j] + links_info.inertial_quat[I][j] = links_inertial_quat[I, j] for j in ti.static(range(3)): links_info.pos[I][j] = links_pos[i, j] - links_info.inertial_pos[I][j] = links_inertial_pos[i, j] + links_info.inertial_pos[I][j] = links_inertial_pos[I, j] - links_info.inertial_mass[I] = links_inertial_mass[i] + links_info.inertial_mass[I] = links_inertial_mass[I] for j1 in ti.static(range(3)): for j2 in ti.static(range(3)): - links_info.inertial_i[I][j1, j2] = links_inertial_i[i, j1, j2] + links_info.inertial_i[I][j1, j2] = links_inertial_i[I, j1, j2] + + links_info.geom_start[I] = links_geom_start[I] + links_info.geom_end[I] = links_geom_end[I] + links_info.vgeom_start[I] = links_vgeom_start[I] + links_info.vgeom_end[I] = links_vgeom_end[I] for i, b in ti.ndrange(n_links, _B): I = [i, b] if ti.static(static_rigid_sim_config.batch_links_info) else i @@ -2950,24 +3086,25 @@ def kernel_init_geom_fields( @ti.kernel def kernel_adjust_link_inertia( link_idx: ti.i32, - ratio: ti.f32, + ratio: ti.f64, # cannot use gs.ti_float() here as there will be a circular import links_info: array_class.LinksInfo, static_rigid_sim_config: ti.template(), ): + ratio_float = gs.ti_float(ratio) if ti.static(static_rigid_sim_config.batch_links_info): _B = links_info.root_idx.shape[1] for i_b in range(_B): for j in ti.static(range(2)): - links_info.invweight[link_idx, i_b][j] /= ratio - links_info.inertial_mass[link_idx, i_b] *= ratio + links_info.invweight[link_idx, i_b][j] /= ratio_float + links_info.inertial_mass[link_idx, i_b] *= ratio_float for j1, j2 in ti.static(ti.ndrange(3, 3)): - links_info.inertial_i[link_idx, i_b][j1, j2] *= ratio + links_info.inertial_i[link_idx, i_b][j1, j2] *= ratio_float else: for j in ti.static(range(2)): - links_info.invweight[link_idx][j] /= ratio - links_info.inertial_mass[link_idx] *= ratio + links_info.invweight[link_idx][j] /= ratio_float + links_info.inertial_mass[link_idx] *= ratio_float for j1, j2 in ti.static(ti.ndrange(3, 3)): - links_info.inertial_i[link_idx][j1, j2] *= ratio + links_info.inertial_i[link_idx][j1, j2] *= ratio_float @ti.kernel @@ -3027,28 +3164,28 @@ def kernel_init_entity_fields( _B = entities_state.hibernated.shape[1] ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) - for i in range(n_entities): - entities_info.dof_start[i] = entities_dof_start[i] - entities_info.dof_end[i] = entities_dof_end[i] - entities_info.n_dofs[i] = entities_dof_end[i] - entities_dof_start[i] + for i_b in range(_B): + for i_e in range(n_entities): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e - entities_info.link_start[i] = entities_link_start[i] - entities_info.link_end[i] = entities_link_end[i] - entities_info.n_links[i] = entities_link_end[i] - entities_link_start[i] + entities_info.dof_start[I_e] = entities_dof_start[i_e] + entities_info.dof_end[I_e] = entities_dof_end[i_e] + entities_info.n_dofs[I_e] = entities_dof_end[i_e] - entities_dof_start[i_e] - entities_info.geom_start[i] = entities_geom_start[i] - entities_info.geom_end[i] = entities_geom_end[i] - entities_info.n_geoms[i] = entities_geom_end[i] - entities_geom_start[i] + entities_info.link_start[I_e] = entities_link_start[i_e] + entities_info.link_end[I_e] = entities_link_end[i_e] + entities_info.n_links[I_e] = entities_link_end[i_e] - entities_link_start[i_e] - entities_info.gravity_compensation[i] = entities_gravity_compensation[i] - entities_info.is_local_collision_mask[i] = entities_is_local_collision_mask[i] + entities_info.geom_start[I_e] = entities_geom_start[i_e] + entities_info.geom_end[I_e] = entities_geom_end[i_e] + entities_info.n_geoms[I_e] = entities_geom_end[i_e] - entities_geom_start[i_e] - if ti.static(static_rigid_sim_config.batch_dofs_info): - for i_d, i_b in ti.ndrange((entities_dof_start[i], entities_dof_end[i]), _B): - dofs_info.dof_start[i_d, i_b] = entities_dof_start[i] - else: - for i_d in range(entities_dof_start[i], entities_dof_end[i]): - dofs_info.dof_start[i_d] = entities_dof_start[i] + entities_info.gravity_compensation[I_e] = entities_gravity_compensation[i_e] + entities_info.is_local_collision_mask[I_e] = entities_is_local_collision_mask[i_e] + + for i_d in range(entities_dof_start[i_e], entities_dof_end[i_e]): + I_d = [i_d, i_b] if ti.static(static_rigid_sim_config.batch_dofs_info) else i_d + dofs_info.dof_start[I_d] = entities_dof_start[i_e] if ti.static(static_rigid_sim_config.use_hibernation): ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) @@ -3179,9 +3316,10 @@ def func_compute_mass_matrix( for i_b in range(_B): for i_e_ in range(rigid_global_info.n_awake_entities[i_b]): i_e = rigid_global_info.awake_entities[i_e_, i_b] - for i in range(entities_info.n_links[i_e]): - i_l = entities_info.link_end[i_e] - 1 - i - I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i in range(entities_info.n_links[I_e]): + i_l = entities_info.link_end[I_e] - 1 - i + l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l i_p = links_info.parent_idx[I_l] if i_p != -1: @@ -3212,20 +3350,21 @@ def func_compute_mass_matrix( for i_b in range(_B): for i_e_ in range(rigid_global_info.n_awake_entities[i_b]): i_e = rigid_global_info.awake_entities[i_e_, i_b] - for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): - for j_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_d in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): + for j_d in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): rigid_global_info.mass_mat[i_d, j_d, i_b] = ( dofs_state.f_ang[i_d, i_b].dot(dofs_state.cdof_ang[j_d, i_b]) + dofs_state.f_vel[i_d, i_b].dot(dofs_state.cdof_vel[j_d, i_b]) ) * rigid_global_info.mass_parent_mask[i_d, j_d] # FIXME: Updating the lower-part of the mass matrix is irrelevant - for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): - for j_d in range(i_d + 1, entities_info.dof_end[i_e]): + for i_d in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): + for j_d in range(i_d + 1, entities_info.dof_end[I_e]): rigid_global_info.mass_mat[i_d, j_d, i_b] = rigid_global_info.mass_mat[j_d, i_d, i_b] # Take into account motor armature - for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + for i_d in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): I_d = [i_d, i_b] if ti.static(static_rigid_sim_config.batch_dofs_info) else i_d rigid_global_info.mass_mat[i_d, i_d, i_b] = ( rigid_global_info.mass_mat[i_d, i_d, i_b] + dofs_info.armature[I_d] @@ -3233,7 +3372,7 @@ def func_compute_mass_matrix( # Take into account first-order correction terms for implicit integration scheme right away if ti.static(implicit_damping): - for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + for i_d in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): I_d = [i_d, i_b] if ti.static(static_rigid_sim_config.batch_dofs_info) else i_d rigid_global_info.mass_mat[i_d, i_d, i_b] += ( dofs_info.damping[I_d] * static_rigid_sim_config.substep_dt @@ -3257,8 +3396,9 @@ def func_compute_mass_matrix( # crb ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) for i_e, i_b in ti.ndrange(n_entities, _B): - for i in range(entities_info.n_links[i_e]): - i_l = entities_info.link_end[i_e] - 1 - i + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i in range(entities_info.n_links[I_e]): + i_l = entities_info.link_end[I_e] - 1 - i I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l i_p = links_info.parent_idx[I_l] @@ -3286,9 +3426,10 @@ def func_compute_mass_matrix( ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e, i_b in ti.ndrange(n_entities, _B): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e for i_d, j_d in ti.ndrange( - (entities_info.dof_start[i_e], entities_info.dof_end[i_e]), - (entities_info.dof_start[i_e], entities_info.dof_end[i_e]), + (entities_info.dof_start[I_e], entities_info.dof_end[I_e]), + (entities_info.dof_start[I_e], entities_info.dof_end[I_e]), ): rigid_global_info.mass_mat[i_d, j_d, i_b] = ( dofs_state.f_ang[i_d, i_b].dot(dofs_state.cdof_ang[j_d, i_b]) @@ -3296,8 +3437,8 @@ def func_compute_mass_matrix( ) * rigid_global_info.mass_parent_mask[i_d, j_d] # FIXME: Updating the lower-part of the mass matrix is irrelevant - for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): - for j_d in range(i_d + 1, entities_info.dof_end[i_e]): + for i_d in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): + for j_d in range(i_d + 1, entities_info.dof_end[I_e]): rigid_global_info.mass_mat[i_d, j_d, i_b] = rigid_global_info.mass_mat[j_d, i_d, i_b] # Take into account motor armature @@ -3341,11 +3482,12 @@ def func_factor_mass( for i_b in range(_B): for i_e_ in range(rigid_global_info.n_awake_entities[i_b]): i_e = rigid_global_info.awake_entities[i_e_, i_b] + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e if rigid_global_info._mass_mat_mask[i_e, i_b] == 1: - entity_dof_start = entities_info.dof_start[i_e] - entity_dof_end = entities_info.dof_end[i_e] - n_dofs = entities_info.n_dofs[i_e] + entity_dof_start = entities_info.dof_start[I_e] + entity_dof_end = entities_info.dof_end[I_e] + n_dofs = entities_info.n_dofs[I_e] for i_d in range(entity_dof_start, entity_dof_end): for j_d in range(entity_dof_start, i_d + 1): @@ -3382,10 +3524,11 @@ def func_factor_mass( else: ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e, i_b in ti.ndrange(n_entities, _B): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e if rigid_global_info._mass_mat_mask[i_e, i_b] == 1: - entity_dof_start = entities_info.dof_start[i_e] - entity_dof_end = entities_info.dof_end[i_e] - n_dofs = entities_info.n_dofs[i_e] + entity_dof_start = entities_info.dof_start[I_e] + entity_dof_end = entities_info.dof_end[I_e] + n_dofs = entities_info.n_dofs[I_e] for i_d in range(entity_dof_start, entity_dof_end): for j_d in range(entity_dof_start, i_d + 1): @@ -3436,11 +3579,12 @@ def func_solve_mass_batched( ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e_ in range(rigid_global_info.n_awake_entities[i_b]): i_e = rigid_global_info.awake_entities[i_e_, i_b] + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e if rigid_global_info._mass_mat_mask[i_e, i_b] == 1: - entity_dof_start = entities_info.dof_start[i_e] - entity_dof_end = entities_info.dof_end[i_e] - n_dofs = entities_info.n_dofs[i_e] + entity_dof_start = entities_info.dof_start[I_e] + entity_dof_end = entities_info.dof_end[I_e] + n_dofs = entities_info.n_dofs[I_e] # Step 1: Solve w st. L^T @ w = y for i_d_ in range(n_dofs): @@ -3460,10 +3604,11 @@ def func_solve_mass_batched( else: ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e in range(n_entities): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e if rigid_global_info._mass_mat_mask[i_e, i_b] == 1: - entity_dof_start = entities_info.dof_start[i_e] - entity_dof_end = entities_info.dof_end[i_e] - n_dofs = entities_info.n_dofs[i_e] + entity_dof_start = entities_info.dof_start[I_e] + entity_dof_end = entities_info.dof_end[I_e] + n_dofs = entities_info.n_dofs[I_e] # Step 1: Solve w st. L^T @ w = y for i_d_ in range(n_dofs): @@ -4032,8 +4177,9 @@ def func_implicit_damping( ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e, i_b in ti.ndrange(n_entities, _B): - entity_dof_start = entities_info.dof_start[i_e] - entity_dof_end = entities_info.dof_end[i_e] + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + entity_dof_start = entities_info.dof_start[I_e] + entity_dof_end = entities_info.dof_end[I_e] for i_d in range(entity_dof_start, entity_dof_end): I_d = [i_d, i_b] if ti.static(static_rigid_sim_config.batch_dofs_info) else i_d if dofs_info.damping[I_d] > gs.EPS: @@ -4648,7 +4794,8 @@ def func_forward_kinematics_entity( rigid_global_info: array_class.RigidGlobalInfo, static_rigid_sim_config: ti.template(), ): - for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_l in range(entities_info.link_start[I_e], entities_info.link_end[I_e]): I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l pos = links_info.pos[I_l] @@ -4760,7 +4907,8 @@ def func_forward_velocity_entity( rigid_global_info: array_class.RigidGlobalInfo, static_rigid_sim_config: ti.template(), ): - for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_l in range(entities_info.link_start[I_e], entities_info.link_end[I_e]): I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l cvel_vel = ti.Vector.zero(gs.ti_float, 3) @@ -5092,20 +5240,21 @@ def func_aggregate_awake_entities( rigid_global_info.n_awake_dofs.fill(0) ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e, i_b in ti.ndrange(n_entities, _B): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e if entities_state.hibernated[i_e, i_b] or entities_info.n_dofs[i_e] == 0: continue next_awake_entity_idx = ti.atomic_add(rigid_global_info.n_awake_entities[i_b], 1) rigid_global_info.awake_entities[next_awake_entity_idx, i_b] = i_e - n_dofs = entities_info.n_dofs[i_e] - entity_dofs_base_idx: ti.int32 = entities_info.dof_start[i_e] + n_dofs = entities_info.n_dofs[I_e] + entity_dofs_base_idx: ti.int32 = entities_info.dof_start[I_e] awake_dofs_base_idx = ti.atomic_add(rigid_global_info.n_awake_dofs[i_b], n_dofs) for i in range(n_dofs): rigid_global_info.awake_dofs[awake_dofs_base_idx + i, i_b] = entity_dofs_base_idx + i - n_links = entities_info.n_links[i_e] - entity_links_base_idx: ti.int32 = entities_info.link_start[i_e] + n_links = entities_info.n_links[I_e] + entity_links_base_idx: ti.int32 = entities_info.link_start[I_e] awake_links_base_idx = ti.atomic_add(rigid_global_info.n_awake_links[i_b], n_links) for i in range(n_links): rigid_global_info.awake_links[awake_links_base_idx + i, i_b] = entity_links_base_idx + i @@ -5128,15 +5277,17 @@ def func_hibernate_entity_and_zero_dof_velocities( """ entities_state.hibernated[i_e, i_b] = True - for i_d in range(entities_info.dof_start[i_e], entities_info.dof_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_d in range(entities_info.dof_start[I_e], entities_info.dof_end[I_e]): dofs_state.hibernated[i_d, i_b] = True dofs_state.vel[i_d, i_b] = 0.0 dofs_state.acc[i_d, i_b] = 0.0 - for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_l in range(entities_info.link_start[I_e], entities_info.link_end[I_e]): links_state.hibernated[i_l, i_b] = True - for i_g in range(entities_info.geom_start[i_e], entities_info.geom_end[i_e]): + for i_g in range(entities_info.geom_start[I_e], entities_info.geom_end[I_e]): geoms_state.hibernated[i_g, i_b] = True @@ -5270,7 +5421,8 @@ def func_torque_and_passive_force( ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) for i_e, i_b in ti.ndrange(n_entities, _B): wakeup = False - for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_l in range(entities_info.link_start[I_e], entities_info.link_end[I_e]): I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l if links_info.n_dofs[I_l] == 0: continue @@ -5345,18 +5497,20 @@ def func_torque_and_passive_force( if ti.abs(force) > gs.EPS: wakeup = True - if ti.static(static_rigid_sim_config.use_hibernation) and entities_state.hibernated[i_e, i_b] and wakeup: - func_wakeup_entity_and_its_temp_island( - i_e, - i_b, - entities_state, - entities_info, - dofs_state, - links_state, - geoms_state, - rigid_global_info, - contact_island_state, - ) + if ti.static(static_rigid_sim_config.use_hibernation): + if entities_state.hibernated[i_e, i_b] and wakeup: + func_wakeup_entity_and_its_temp_island( + i_e, + i_b, + entities_state, + entities_info, + dofs_state, + links_state, + geoms_state, + rigid_global_info, + contact_island_state, + static_rigid_sim_config, + ) if ti.static(static_rigid_sim_config.use_hibernation): ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) @@ -5444,13 +5598,14 @@ def func_update_acc( for i_b in range(_B): for i_e_ in range(rigid_global_info.n_awake_entities[i_b]): i_e = rigid_global_info.awake_entities[i_e_, i_b] - for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_l in range(entities_info.link_start[I_e], entities_info.link_end[I_e]): I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l i_p = links_info.parent_idx[I_l] if i_p == -1: links_state.cdd_vel[i_l, i_b] = -rigid_global_info.gravity[i_b] * ( - 1 - entities_info.gravity_compensation[i_e] + 1 - entities_info.gravity_compensation[I_e] ) links_state.cdd_ang[i_l, i_b] = ti.Vector.zero(gs.ti_float, 3) if ti.static(update_cacc): @@ -5482,13 +5637,14 @@ def func_update_acc( else: ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) for i_e, i_b in ti.ndrange(n_entities, _B): - for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_l in range(entities_info.link_start[I_e], entities_info.link_end[I_e]): I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l i_p = links_info.parent_idx[I_l] if i_p == -1: links_state.cdd_vel[i_l, i_b] = -rigid_global_info.gravity[i_b] * ( - 1 - entities_info.gravity_compensation[i_e] + 1 - entities_info.gravity_compensation[I_e] ) links_state.cdd_ang[i_l, i_b] = ti.Vector.zero(gs.ti_float, 3) if ti.static(update_cacc): @@ -5562,8 +5718,9 @@ def func_update_force( for i_b in range(_B): for i_e_ in range(rigid_global_info.n_awake_entities[i_b]): i_e = rigid_global_info.awake_entities[i_e_, i_b] - for i in range(entities_info.n_links[i_e]): - i_l = entities_info.link_end[i_e] - 1 - i + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i in range(entities_info.n_links[I_e]): + i_l = entities_info.link_end[I_e] - 1 - i I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l i_p = links_info.parent_idx[I_l] if i_p != -1: @@ -5595,8 +5752,9 @@ def func_update_force( ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) for i_e, i_b in ti.ndrange(n_entities, _B): - for i in range(entities_info.n_links[i_e]): - i_l = entities_info.link_end[i_e] - 1 - i + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i in range(entities_info.n_links[I_e]): + i_l = entities_info.link_end[I_e] - 1 - i I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l i_p = links_info.parent_idx[I_l] if i_p != -1: @@ -5703,14 +5861,16 @@ def func_compute_qacc( for i_b in range(_B): for i_e_ in range(rigid_global_info.n_awake_entities[i_b]): i_e = rigid_global_info.awake_entities[i_e_, i_b] - for i_d1_ in range(entities_info.n_dofs[i_e]): - i_d1 = entities_info.dof_start[i_e] + i_d1_ + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_d1_ in range(entities_info.n_dofs[I_e]): + i_d1 = entities_info.dof_start[I_e] + i_d1_ dofs_state.acc[i_d1, i_b] = dofs_state.acc_smooth[i_d1, i_b] else: ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e, i_b in ti.ndrange(n_entities, _B): - for i_d1_ in range(entities_info.n_dofs[i_e]): - i_d1 = entities_info.dof_start[i_e] + i_d1_ + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_d1_ in range(entities_info.n_dofs[I_e]): + i_d1 = entities_info.dof_start[I_e] + i_d1_ dofs_state.acc[i_d1, i_b] = dofs_state.acc_smooth[i_d1, i_b] @@ -5906,7 +6066,8 @@ def func_integrate_dq_entity( rigid_global_info: array_class.RigidGlobalInfo, static_rigid_sim_config: ti.template(), ): - for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_l in range(entities_info.link_start[I_e], entities_info.link_end[I_e]): I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l if links_info.n_dofs[I_l] == 0: continue @@ -5917,7 +6078,7 @@ def func_integrate_dq_entity( q_start = links_info.q_start[I_l] dof_start = links_info.dof_start[I_l] - dq_start = links_info.dof_start[I_l] - entities_info.dof_start[i_e] + dq_start = links_info.dof_start[I_l] - entities_info.dof_start[I_e] if joint_type == gs.JOINT_TYPE.FREE: pos = ti.Vector( @@ -6548,7 +6709,8 @@ def kernel_set_dofs_position( ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) for i_e, i_b_ in ti.ndrange(n_entities, envs_idx.shape[0]): i_b = envs_idx[i_b_] - for i_l in range(entities_info.link_start[i_e], entities_info.link_end[i_e]): + I_e = [i_e, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else i_e + for i_l in range(entities_info.link_start[I_e], entities_info.link_end[I_e]): I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l if links_info.n_dofs[I_l] == 0: continue @@ -6709,6 +6871,35 @@ def kernel_get_links_acc( tensor[i_b_, i_l_, i] = acc_classic_lin[i] +@ti.kernel +def kernel_get_links_AABB( + tensor: ti.types.ndarray(), + links_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + links_info: array_class.LinksInfo, + geoms_state: array_class.GeomsState, + static_rigid_sim_config: ti.template(), +): + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL) + for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): + i_l = links_idx[i_l_] + i_b = envs_idx[i_b_] + + lower = gu.ti_vec3(ti.math.inf) + upper = gu.ti_vec3(-ti.math.inf) + I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l + for i_g in range(links_info.geom_start[I_l], links_info.geom_end[I_l]): + print("i_b", i_b, i_g) + aabb_min = geoms_state.aabb_min[i_g, i_b] + aabb_max = geoms_state.aabb_max[i_g, i_b] + lower = ti.min(lower, aabb_min) + upper = ti.max(upper, aabb_max) + + for i_3 in ti.static(range(3)): + tensor[i_b_, i_l_, i_3] = lower[i_3] + tensor[i_b_, i_l_, i_3 + 3] = upper[i_3] + + @ti.kernel def kernel_get_dofs_control_force( tensor: ti.types.ndarray(), diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp_util.py b/genesis/engine/solvers/rigid/rigid_solver_decomp_util.py index e27398c7c..33f114a95 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp_util.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp_util.py @@ -16,6 +16,7 @@ def func_wakeup_entity_and_its_temp_island( geoms_state: array_class.GeomsState, rigid_global_info: array_class.RigidGlobalInfo, contact_island_state: array_class.ContactIslandState, + static_rigid_sim_config: ti.template(), ): # Note: Original function handled non-hibernated & fixed entities. # Now, we require a properly hibernated entity to be passed in. @@ -36,21 +37,23 @@ def func_wakeup_entity_and_its_temp_island( n_awake_entities = ti.atomic_add(rigid_global_info.n_awake_entities[i_b], 1) rigid_global_info.awake_entities[n_awake_entities, i_b] = entity_idx - n_dofs = entities_info.n_dofs[entity_idx] - base_entity_dof_idx = entities_info.dof_start[entity_idx] + I_e = [entity_idx, i_b] if ti.static(static_rigid_sim_config.batch_entities_info) else entity_idx + + n_dofs = entities_info.n_dofs[I_e] + base_entity_dof_idx = entities_info.dof_start[I_e] base_awake_dof_idx = ti.atomic_add(rigid_global_info.n_awake_dofs[i_b], n_dofs) for i in range(n_dofs): i_d = base_entity_dof_idx + i dofs_state.hibernated[i_d, i_b] = False rigid_global_info.awake_dofs[base_awake_dof_idx + i, i_b] = i_d - n_links = entities_info.n_links[entity_idx] - base_entity_link_idx = entities_info.link_start[entity_idx] + n_links = entities_info.n_links[I_e] + base_entity_link_idx = entities_info.link_start[I_e] base_awake_link_idx = ti.atomic_add(rigid_global_info.n_awake_links[i_b], n_links) for i in range(n_links): i_l = base_entity_link_idx + i links_state.hibernated[i_l, i_b] = False rigid_global_info.awake_links[base_awake_link_idx + i, i_b] = i_l - for i_g in range(entities_info.geom_start[entity_idx], entities_info.geom_end[entity_idx]): + for i_g in range(entities_info.geom_start[I_e], entities_info.geom_end[I_e]): geoms_state.hibernated[i_g, i_b] = False diff --git a/genesis/options/solvers.py b/genesis/options/solvers.py index 790862419..090562c76 100644 --- a/genesis/options/solvers.py +++ b/genesis/options/solvers.py @@ -287,6 +287,7 @@ class RigidOptions(Options): IK_max_targets: int = 6 # batching info + batch_entities_info: Optional[bool] = False batch_links_info: Optional[bool] = False batch_joints_info: Optional[bool] = False batch_dofs_info: Optional[bool] = False diff --git a/genesis/utils/array_class.py b/genesis/utils/array_class.py index c820aeaee..248a14963 100644 --- a/genesis/utils/array_class.py +++ b/genesis/utils/array_class.py @@ -1375,6 +1375,10 @@ class StructLinksInfo: dof_end: V_ANNOTATION joint_end: V_ANNOTATION n_dofs: V_ANNOTATION + geom_start: V_ANNOTATION + geom_end: V_ANNOTATION + vgeom_start: V_ANNOTATION + vgeom_end: V_ANNOTATION pos: V_ANNOTATION quat: V_ANNOTATION invweight: V_ANNOTATION @@ -1398,6 +1402,11 @@ def get_links_info(solver): "dof_end": V(dtype=gs.ti_int, shape=links_info_shape), "joint_end": V(dtype=gs.ti_int, shape=links_info_shape), "n_dofs": V(dtype=gs.ti_int, shape=links_info_shape), + "geom_start": V(dtype=gs.ti_int, shape=links_info_shape), + "geom_end": V(dtype=gs.ti_int, shape=links_info_shape), + "vgeom_start": V(dtype=gs.ti_int, shape=links_info_shape), + "vgeom_end": V(dtype=gs.ti_int, shape=links_info_shape), + "n_geoms": V(dtype=gs.ti_int, shape=links_info_shape), "pos": V(dtype=gs.ti_vec3, shape=links_info_shape), "quat": V(dtype=gs.ti_vec4, shape=links_info_shape), "invweight": V(dtype=gs.ti_vec2, shape=links_info_shape), @@ -1949,7 +1958,7 @@ class StructEntitiesInfo: def get_entities_info(solver): - shape = solver.n_entities_ + shape = solver._batch_shape(solver.n_entities_) if solver._options.batch_entities_info else solver.n_entities_ kwargs = { "dof_start": V(dtype=gs.ti_int, shape=shape), "dof_end": V(dtype=gs.ti_int, shape=shape), diff --git a/genesis/vis/rasterizer_context.py b/genesis/vis/rasterizer_context.py index 5c712c77b..38691f040 100644 --- a/genesis/vis/rasterizer_context.py +++ b/genesis/vis/rasterizer_context.py @@ -303,6 +303,12 @@ def set_reflection_mat(self, geom_T): dtype=np.float32, ) + def get_geom_rendered_envs_idx(self, geom, rendered_envs_idx): + geom_rendered_envs_idx = getattr(geom, "rendered_envs_idx", rendered_envs_idx) + # intersection of geom.rendered_envs_idx and rendered_envs_idx + intersection = np.intersect1d(geom_rendered_envs_idx, rendered_envs_idx) + return intersection + def on_rigid(self): if self.sim.rigid_solver.is_active(): # TODO: support dynamic switching in GUI later @@ -319,7 +325,7 @@ def on_rigid(self): mesh = geom.get_sdf_trimesh() else: mesh = geom.get_trimesh() - geom_T = geoms_T[geom.idx][self.rendered_envs_idx] + geom_T = geoms_T[geom.idx][self.get_geom_rendered_envs_idx(geom, self.rendered_envs_idx)] self.add_rigid_node( geom, pyrender.Mesh.from_trimesh( @@ -347,7 +353,7 @@ def update_rigid(self, buffer_updates): geoms_T = self.sim.rigid_solver._geoms_render_T for geom in geoms: - geom_T = geoms_T[geom.idx][self.rendered_envs_idx] + geom_T = geoms_T[geom.idx][self.get_geom_rendered_envs_idx(geom, self.rendered_envs_idx)] node = self.rigid_nodes[geom.uid] node.mesh._bounds = None for primitive in node.mesh.primitives: diff --git a/tests/test_rigid_physics.py b/tests/test_rigid_physics.py index 36db1fb2a..f611c02df 100644 --- a/tests/test_rigid_physics.py +++ b/tests/test_rigid_physics.py @@ -2553,7 +2553,6 @@ def must_cast(value): (-1, -1, gs_link.get_mass, gs_link.set_mass, None), ): getter, spec = (getter_or_spec, None) if callable(getter_or_spec) else (None, getter_or_spec) - # Check getter and setter without row or column masking if getter is not None: datas = getter() @@ -2833,3 +2832,54 @@ def test_contype_conaffinity(show_viewer, tol): assert_allclose(box1.get_pos(), np.array([0.0, 0.0, 0.25]), atol=1e-3) assert_allclose(box2.get_pos(), np.array([0.0, 0.0, 0.75]), atol=1e-3) assert_allclose(box3.get_pos(), np.array([0.0, 0.0, 0.75]), atol=1e-3) + + +@pytest.mark.parametrize("backend", [gs.cpu]) +def test_heterogeneous_simulation(show_viewer, tol): + scene = gs.Scene( + show_viewer=show_viewer, + ) + + ########################## entities ########################## + plane = scene.add_entity( + gs.morphs.Plane(), + ) + + morphs_heterogeneous = [ + gs.morphs.Box( + size=(0.04, 0.04, 0.04), + pos=(0.65, 0.0, 0.02), + ), + gs.morphs.Box( + size=(0.02, 0.02, 0.02), + pos=(0.65, 0.0, 0.02), + ), + gs.morphs.Sphere( + radius=0.015, + pos=(0.65, 0.0, 0.02), + ), + gs.morphs.Sphere( + radius=0.025, + pos=(0.65, 0.0, 0.02), + ), + ] + grasping_object = scene.add_entity( + morph=morphs_heterogeneous, + ) + ########################## build ########################## + scene.build(n_envs=4, env_spacing=(1, 1)) + + for i in range(300): + qdofs = grasping_object.get_qpos() + scene.step() + print(qdofs[:, 2]) + final_qpos_z = grasping_object.get_qpos()[:, 2] + assert_allclose(final_qpos_z, np.array([0.02, 0.01, 0.015, 0.025]), atol=1e-2) + + AABB = grasping_object.get_AABB() + size_x = AABB[:, 1, 0] - AABB[:, 0, 0] + assert_allclose(size_x, np.array([0.04, 0.02, 0.03, 0.05]), atol=tol) + all_mass = grasping_object.get_mass() + # assert it's not the same + diff_mass = np.abs(np.diff(all_mass)).sum() + assert diff_mass > 0.0