diff --git a/genesis/engine/solvers/avatar_solver.py b/genesis/engine/solvers/avatar_solver.py index bb8788cbff..eb3782acad 100644 --- a/genesis/engine/solvers/avatar_solver.py +++ b/genesis/engine/solvers/avatar_solver.py @@ -76,10 +76,10 @@ def get_state(self, f): def print_contact_data(self): batch_idx = 0 - n_contacts = self.collider.n_contacts[batch_idx] + n_contacts = self.collider._collider_state.n_contacts[batch_idx] print("collision_pairs:") if n_contacts > 0: - contact_data = self.collider.contact_data.to_numpy() + contact_data = self.collider._collider_state.contact_data.to_numpy() links_a = contact_data["link_a"][:n_contacts, batch_idx] links_b = contact_data["link_b"][:n_contacts, batch_idx] link_pairs = np.vstack([links_a, links_b]).T diff --git a/genesis/engine/solvers/rigid/array_class.py b/genesis/engine/solvers/rigid/array_class.py index d719360879..c8e8cc67ac 100644 --- a/genesis/engine/solvers/rigid/array_class.py +++ b/genesis/engine/solvers/rigid/array_class.py @@ -4,14 +4,140 @@ import taichi as ti import genesis as gs +import numpy as np # we will use struct for DofsState and DofsInfo after Hugh adds array_struct feature to taichi DofsState = ti.template() DofsInfo = ti.template() +GeomsState = ti.template() +GeomsInfo = ti.template() +GeomsInitAABB = ti.template() +LinksState = ti.template() +LinksInfo = ti.template() +VertsInfo = ti.template() +EdgesInfo = ti.template() @ti.data_oriented class RigidGlobalInfo: - def __init__(self, n_dofs: int, n_entities: int, n_geoms: int, f_batch: Callable): + def __init__(self, n_dofs: int, n_entities: int, n_geoms: int, _B: int, f_batch: Callable): self.n_awake_dofs = ti.field(dtype=gs.ti_int, shape=f_batch()) self.awake_dofs = ti.field(dtype=gs.ti_int, shape=f_batch(n_dofs)) + + +# =========================================== Collider =========================================== + + +@ti.data_oriented +class ColliderState: + """ + Class to store the mutable collider data, all of which type is [ti.fields]. + """ + + def __init__(self, solver, n_possible_pairs, n_vert_neighbors, collider_info): + """ + Parameters: + ---------- + n_possible_pairs: int + Maximum number of possible collision pairs based on geom configurations. For instance, when adjacent + collision is disabled, adjacent geoms are not considered in counting possible pairs. + n_vert_neighbors: int + Size of the vertex neighbors array. + """ + _B = solver._B + f_batch = solver._batch_shape + n_geoms = solver.n_geoms_ + n_verts = solver.n_verts_ + max_collision_pairs = min(solver._max_collision_pairs, n_possible_pairs) + max_contact_pairs = max_collision_pairs * collider_info.n_contacts_per_pair + use_hibernation = solver._static_rigid_sim_config.use_hibernation + box_box_detection = solver._static_rigid_sim_config.box_box_detection + + ############## vertex connectivity ############## + self.vert_neighbors = ti.field(dtype=gs.ti_int, shape=max(1, n_vert_neighbors)) + self.vert_neighbor_start = ti.field(dtype=gs.ti_int, shape=n_verts) + self.vert_n_neighbors = ti.field(dtype=gs.ti_int, shape=n_verts) + + ############## broad phase SAP ############## + # This buffer stores the AABBs along the search axis of all geoms + struct_sort_buffer = ti.types.struct(value=gs.ti_float, i_g=gs.ti_int, is_max=gs.ti_int) + self.sort_buffer = struct_sort_buffer.field(shape=f_batch(2 * n_geoms), layout=ti.Layout.SOA) + + # This buffer stores indexes of active geoms during SAP search + if use_hibernation: + self.active_buffer_awake = ti.field(dtype=gs.ti_int, shape=f_batch(n_geoms)) + self.active_buffer_hib = ti.field(dtype=gs.ti_int, shape=f_batch(n_geoms)) + self.active_buffer = ti.field(dtype=gs.ti_int, shape=f_batch(n_geoms)) + + # Stores the validity of the collision pairs + self.collision_pair_validity = ti.field(dtype=gs.ti_int, shape=(n_geoms, n_geoms)) + + # Whether or not this is the first time to run the broad phase for each batch + self.first_time = ti.field(gs.ti_int, shape=_B) + + # Number of possible pairs of collision, store them in a field to avoid recompilation + self._max_possible_pairs = ti.field(dtype=gs.ti_int, shape=()) + self._max_collision_pairs = ti.field(dtype=gs.ti_int, shape=()) + self._max_contact_pairs = ti.field(dtype=gs.ti_int, shape=()) + + # Final results of the broad phase + self.n_broad_pairs = ti.field(dtype=gs.ti_int, shape=_B) + self.broad_collision_pairs = ti.Vector.field(2, dtype=gs.ti_int, shape=f_batch(max(1, max_collision_pairs))) + + ############## narrow phase ############## + struct_contact_data = ti.types.struct( + geom_a=gs.ti_int, + geom_b=gs.ti_int, + penetration=gs.ti_float, + normal=gs.ti_vec3, + pos=gs.ti_vec3, + friction=gs.ti_float, + sol_params=gs.ti_vec7, + force=gs.ti_vec3, + link_a=gs.ti_int, + link_b=gs.ti_int, + ) + self.contact_data = struct_contact_data.field( + shape=f_batch(max(1, max_contact_pairs)), + layout=ti.Layout.SOA, + ) + # total number of contacts, including hibernated contacts + self.n_contacts = ti.field(gs.ti_int, shape=_B) + self.n_contacts_hibernated = ti.field(gs.ti_int, shape=_B) + + # contact caching for warmstart collision detection + struct_contact_cache = ti.types.struct( + # i_va_ws=gs.ti_int, + # penetration=gs.ti_float, + normal=gs.ti_vec3, + ) + self.contact_cache = struct_contact_cache.field(shape=f_batch((n_geoms, n_geoms)), layout=ti.Layout.SOA) + + ########## Box-box contact detection ########## + if box_box_detection: + # With the existing Box-Box collision detection algorithm, it is not clear where the contact points are + # located depending of the pose and size of each box. In practice, up to 11 contact points have been + # observed. The theoretical worst case scenario would be 2 cubes roughly the same size and same center, + # with transform RPY = (45, 45, 45), resulting in 3 contact points per faces for a total of 16 points. + self.box_depth = ti.field(dtype=gs.ti_float, shape=f_batch(collider_info.box_MAXCONPAIR)) + self.box_points = ti.field(gs.ti_vec3, shape=f_batch(collider_info.box_MAXCONPAIR)) + self.box_pts = ti.field(gs.ti_vec3, shape=f_batch(6)) + self.box_lines = ti.field(gs.ti_vec6, shape=f_batch(4)) + self.box_linesu = ti.field(gs.ti_vec6, shape=f_batch(4)) + self.box_axi = ti.field(gs.ti_vec3, shape=f_batch(3)) + self.box_ppts2 = ti.field(dtype=gs.ti_float, shape=f_batch((4, 2))) + self.box_pu = ti.field(gs.ti_vec3, shape=f_batch(4)) + + ########## Terrain contact detection ########## + if collider_info.has_terrain: + links_idx = solver.geoms_info.link_idx.to_numpy()[solver.geoms_info.type.to_numpy() == gs.GEOM_TYPE.TERRAIN] + entity = solver._entities[solver.links_info.entity_idx.to_numpy()[links_idx[0]]] + + self.terrain_hf = ti.field(dtype=gs.ti_float, shape=entity.terrain_hf.shape) + self.terrain_rc = ti.field(dtype=gs.ti_int, shape=2) + self.terrain_scale = ti.field(dtype=gs.ti_float, shape=2) + self.terrain_xyz_maxmin = ti.field(dtype=gs.ti_float, shape=6) + + # for faster compilation + self.xyz_max_min = ti.field(dtype=gs.ti_float, shape=f_batch(6)) + self.prism = ti.field(dtype=gs.ti_vec3, shape=f_batch(6)) diff --git a/genesis/engine/solvers/rigid/collider_decomp.py b/genesis/engine/solvers/rigid/collider_decomp.py index 2a9cbeb528..88db9201d3 100644 --- a/genesis/engine/solvers/rigid/collider_decomp.py +++ b/genesis/engine/solvers/rigid/collider_decomp.py @@ -1,5 +1,6 @@ import sys from typing import TYPE_CHECKING +from dataclasses import dataclass import numpy as np import numpy.typing as npt @@ -11,6 +12,7 @@ import genesis.utils.geom as gu from genesis.styles import colors, formats from genesis.utils.misc import ti_field_to_torch +import genesis.engine.solvers.rigid.array_class as array_class from .mpr_decomp import MPR from .gjk_decomp import GJK @@ -52,80 +54,122 @@ def rotmatx(matin, i0, i1, i2, f0, f1, f2): @ti.data_oriented class Collider: + @dataclass(frozen=True) + class ColliderInfo: + """ + Class to store the immutable collider data. + """ + + # store static information here + has_nonconvex_nonterrain: bool = False + has_terrain: bool = False + # multi contact perturbation and tolerance + mc_perturbation: float = 0.0 + mc_tolerance: float = 0.0 + mpr_to_sdf_overlap_ratio: float = 0.0 + # maximum number of contact pairs per collision pair + n_contacts_per_pair: int = 5 + # maximum number of contact points for box-box collision detection + box_MAXCONPAIR: int = 16 + # ccd algorithm + ccd_algorithm: CCD_ALGORITHM_CODE = CCD_ALGORITHM_CODE.MPR + def __init__(self, rigid_solver: "RigidSolver"): self._solver = rigid_solver - self._init_verts_connectivity() - self._init_collision_fields() + self._init_info() + self._init_state() + # self._init_collision_fields() + + # FIXME: MPR is necessary because it is used for terrain collision detection + self._mpr = MPR(rigid_solver) + self._gjk = GJK(rigid_solver) if self._solver._options.use_gjk_collision else None + + def _init_info(self) -> None: # Identify the convex collision detection (ccd) algorithm if self._solver._options.use_gjk_collision: if self._solver._enable_mujoco_compatibility: - self.ccd_algorithm = CCD_ALGORITHM_CODE.MJ_GJK + ccd_algorithm = CCD_ALGORITHM_CODE.MJ_GJK else: - self.ccd_algorithm = CCD_ALGORITHM_CODE.GJK + ccd_algorithm = CCD_ALGORITHM_CODE.GJK else: if self._solver._enable_mujoco_compatibility: - self.ccd_algorithm = CCD_ALGORITHM_CODE.MJ_MPR + ccd_algorithm = CCD_ALGORITHM_CODE.MJ_MPR else: - self.ccd_algorithm = CCD_ALGORITHM_CODE.MPR + ccd_algorithm = CCD_ALGORITHM_CODE.MPR + + # Initialize the static info + self._collider_info = Collider.ColliderInfo( + has_nonconvex_nonterrain=np.logical_and( + self._solver.geoms_info.is_convex.to_numpy() == 0, + self._solver.geoms_info.type.to_numpy() != gs.GEOM_TYPE.TERRAIN, + ).any(), + has_terrain=(self._solver.geoms_info.type.to_numpy() == gs.GEOM_TYPE.TERRAIN).any(), + mc_perturbation=1e-3 if self._solver._enable_mujoco_compatibility else 1e-2, + mc_tolerance=1e-3 if self._solver._enable_mujoco_compatibility else 1e-2, + mpr_to_sdf_overlap_ratio=0.4, + n_contacts_per_pair=5, + box_MAXCONPAIR=16, + ccd_algorithm=ccd_algorithm, + ) - # FIXME: MPR is necessary because it is used for terrain collision detection - self._mpr = MPR(rigid_solver) - self._gjk = GJK(rigid_solver) if self._solver._options.use_gjk_collision else None + def _init_state(self) -> None: + # Compute number of possible collision pairs, as it is needed to initialize the collider state + n_possible_pairs, collision_pair_validity = self._compute_collision_pair_validity() + vert_neighbors, vert_neighbor_start, vert_n_neighbors = self._compute_verts_connectivity() + n_vert_neighbors = len(vert_neighbors) - # multi contact perturbation and tolerance - if self._solver._enable_mujoco_compatibility: - self._mc_perturbation = 1e-3 - self._mc_tolerance = 1e-3 - else: - self._mc_perturbation = 1e-2 - self._mc_tolerance = 1e-2 - self._mpr_to_sdf_overlap_ratio = 0.4 + # Initialize [state], which stores every mutable collision data. + self._collider_state = array_class.ColliderState( + self._solver, n_possible_pairs, n_vert_neighbors, self._collider_info + ) + self._init_collision_pair_validity(collision_pair_validity) + self._init_verts_connectivity(vert_neighbors, vert_neighbor_start, vert_n_neighbors) + self._init_max_contact_pairs(n_possible_pairs) + self._init_terrain_state() - def _init_verts_connectivity(self) -> None: - vert_neighbors = [] - vert_neighbor_start = [] - vert_n_neighbors = [] - offset = 0 - for geom in self._solver.geoms: - vert_neighbors.append(geom.vert_neighbors + geom.vert_start) - vert_neighbor_start.append(geom.vert_neighbor_start + offset) - vert_n_neighbors.append(geom.vert_n_neighbors) - offset += len(geom.vert_neighbors) + # [contacts_info_cache] is not used in Taichi kernels, so keep it outside of the collider state. + self._contacts_info_cache = {} - if self._solver.n_verts > 0: - vert_neighbors = np.concatenate(vert_neighbors, dtype=gs.np_int) - vert_neighbor_start = np.concatenate(vert_neighbor_start, dtype=gs.np_int) - vert_n_neighbors = np.concatenate(vert_n_neighbors, dtype=gs.np_int) + # FIXME: 'ti.static_print' cannot be used as it will be printed systematically, completely ignoring guard + # condition, while 'print' is slowing down the kernel even if every called in practice... + self._warn_msg_max_collision_pairs = ( + f"{colors.YELLOW}[Genesis] [00:00:00] [WARNING] Ignoring contact pair to avoid exceeding max " + f"({self._collider_state._max_contact_pairs[None]}). Please increase the value of RigidSolver's option " + f"'max_collision_pairs'.{formats.RESET}" + ) - self.vert_neighbors = ti.field(dtype=gs.ti_int, shape=max(1, len(vert_neighbors))) - self.vert_neighbor_start = ti.field(dtype=gs.ti_int, shape=self._solver.n_verts_) - self.vert_n_neighbors = ti.field(dtype=gs.ti_int, shape=self._solver.n_verts_) + self.reset() - if self._solver.n_verts > 0: - self.vert_neighbors.from_numpy(vert_neighbors) - self.vert_neighbor_start.from_numpy(vert_neighbor_start) - self.vert_n_neighbors.from_numpy(vert_n_neighbors) - - def _init_collision_fields(self) -> None: - # Compute collision pair validity - self.collision_pair_validity = ti.field(gs.ti_int, shape=(self._solver.n_geoms_, self._solver.n_geoms_)) - - geoms_link_idx = self._solver.geoms_info.link_idx.to_numpy() - geoms_contype = self._solver.geoms_info.contype.to_numpy() - geoms_conaffinity = self._solver.geoms_info.conaffinity.to_numpy() - links_entity_idx = self._solver.links_info.entity_idx.to_numpy() - links_root_idx = self._solver.links_info.root_idx.to_numpy() - links_parent_idx = self._solver.links_info.parent_idx.to_numpy() - links_is_fixed = self._solver.links_info.is_fixed.to_numpy() - if self._solver._options.batch_links_info: + def _compute_collision_pair_validity(self): + """ + Compute the collision pair validity matrix. + + For each pair of geoms, determine if they can collide based on their properties and the solver configuration. + """ + solver = self._solver + n_geoms = solver.n_geoms_ + enable_self_collision = solver._static_rigid_sim_config.enable_self_collision + enable_adjacent_collision = solver._static_rigid_sim_config.enable_adjacent_collision + batch_links_info = solver._static_rigid_sim_config.batch_links_info + + geoms_link_idx = solver.geoms_info.link_idx.to_numpy() + geoms_contype = solver.geoms_info.contype.to_numpy() + geoms_conaffinity = solver.geoms_info.conaffinity.to_numpy() + links_entity_idx = solver.links_info.entity_idx.to_numpy() + links_root_idx = solver.links_info.root_idx.to_numpy() + links_parent_idx = solver.links_info.parent_idx.to_numpy() + links_is_fixed = solver.links_info.is_fixed.to_numpy() + if batch_links_info: links_entity_idx = links_entity_idx[:, 0] links_root_idx = links_root_idx[:, 0] links_parent_idx = links_parent_idx[:, 0] links_is_fixed = links_is_fixed[:, 0] + n_possible_pairs = 0 - for i_ga in range(self._solver.n_geoms): - for i_gb in range(i_ga + 1, self._solver.n_geoms): + collision_pair_validity = np.zeros((n_geoms, n_geoms), dtype=gs.np_int) + for i_ga in range(n_geoms): + for i_gb in range(i_ga + 1, n_geoms): i_la = geoms_link_idx[i_ga] i_lb = geoms_link_idx[i_gb] @@ -135,11 +179,11 @@ def _init_collision_fields(self) -> None: # self collision if links_root_idx[i_la] == links_root_idx[i_lb]: - if not self._solver._enable_self_collision: + if not enable_self_collision: continue # adjacent links - if not self._solver._enable_adjacent_collision and ( + if not enable_adjacent_collision and ( links_parent_idx[i_la] == i_lb or links_parent_idx[i_lb] == i_la ): continue @@ -154,231 +198,294 @@ def _init_collision_fields(self) -> None: if links_is_fixed[i_la] and links_is_fixed[i_lb]: continue - self.collision_pair_validity[i_ga, i_gb] = 1 + collision_pair_validity[i_ga, i_gb] = 1 n_possible_pairs += 1 - self._n_contacts_per_pair = 5 # CONSTANT. CANNOT NOT BE CHANGED. - self._max_possible_pairs = n_possible_pairs - self._max_collision_pairs = min(n_possible_pairs, self._solver._max_collision_pairs) - self._max_contact_pairs = self._max_collision_pairs * self._n_contacts_per_pair + return n_possible_pairs, collision_pair_validity - # FIXME: 'ti.static_print' cannot be used as it will be printed systematically, completely ignoring guard - # condition, while 'print' is slowing down the kernel even if every called in practice... - self._warn_msg_max_collision_pairs = ( - f"{colors.YELLOW}[Genesis] [00:00:00] [WARNING] Ignoring contact pair to avoid exceeding max " - f"({self._max_contact_pairs}). Please increase the value of RigidSolver's option " - f"'max_collision_pairs'.{formats.RESET}" - ) + def _compute_verts_connectivity(self): + """ + Compute the vertex connectivity. + """ + vert_neighbors = [] + vert_neighbor_start = [] + vert_n_neighbors = [] + offset = 0 + for geom in self._solver.geoms: + vert_neighbors.append(geom.vert_neighbors + geom.vert_start) + vert_neighbor_start.append(geom.vert_neighbor_start + offset) + vert_n_neighbors.append(geom.vert_n_neighbors) + offset += len(geom.vert_neighbors) - ############## broad phase SAP ############## - # This buffer stores the AABBs along the search axis of all geoms - struct_sort_buffer = ti.types.struct(value=gs.ti_float, i_g=gs.ti_int, is_max=gs.ti_int) - self.sort_buffer = struct_sort_buffer.field( - shape=self._solver._batch_shape(2 * self._solver.n_geoms_), - layout=ti.Layout.SOA, - ) - # This buffer stores indexes of active geoms during SAP search - if self._solver._use_hibernation: - self.active_buffer_awake = ti.field(dtype=gs.ti_int, shape=self._solver._batch_shape(self._solver.n_geoms_)) - self.active_buffer_hib = ti.field(dtype=gs.ti_int, shape=self._solver._batch_shape(self._solver.n_geoms_)) - self.active_buffer = ti.field(dtype=gs.ti_int, shape=self._solver._batch_shape(self._solver.n_geoms_)) - - self.n_broad_pairs = ti.field(dtype=gs.ti_int, shape=self._solver._B) - self.broad_collision_pairs = ti.Vector.field( - 2, dtype=gs.ti_int, shape=self._solver._batch_shape(max(1, self._max_collision_pairs)) - ) + if self._solver.n_verts > 0: + vert_neighbors = np.concatenate(vert_neighbors, dtype=gs.np_int) + vert_neighbor_start = np.concatenate(vert_neighbor_start, dtype=gs.np_int) + vert_n_neighbors = np.concatenate(vert_n_neighbors, dtype=gs.np_int) - self.first_time = ti.field(gs.ti_int, shape=self._solver._B) - - ############## narrow phase ############## - struct_contact_data = ti.types.struct( - geom_a=gs.ti_int, - geom_b=gs.ti_int, - penetration=gs.ti_float, - normal=gs.ti_vec3, - pos=gs.ti_vec3, - friction=gs.ti_float, - sol_params=gs.ti_vec7, - force=gs.ti_vec3, - link_a=gs.ti_int, - link_b=gs.ti_int, - ) - self.contact_data = struct_contact_data.field( - shape=self._solver._batch_shape(max(1, self._max_contact_pairs)), - layout=ti.Layout.SOA, - ) - self.n_contacts = ti.field( - gs.ti_int, shape=self._solver._B - ) # total number of contacts, including hibernated contacts - self.n_contacts_hibernated = ti.field(gs.ti_int, shape=self._solver._B) - self._contacts_info_cache = {} + return vert_neighbors, vert_neighbor_start, vert_n_neighbors - # contact caching for warmstart collision detection - struct_contact_cache = ti.types.struct( - # i_va_ws=gs.ti_int, - # penetration=gs.ti_float, - normal=gs.ti_vec3, - ) - self.contact_cache = struct_contact_cache.field( - shape=self._solver._batch_shape((self._solver.n_geoms_, self._solver.n_geoms_)), - layout=ti.Layout.SOA, - ) + def _init_collision_pair_validity(self, collision_pair_validity): + self._collider_state.collision_pair_validity.from_numpy(collision_pair_validity) - # for faster compilation - self._has_nonconvex_nonterrain = np.logical_and( - self._solver.geoms_info.is_convex.to_numpy() == 0, - self._solver.geoms_info.type.to_numpy() != gs.GEOM_TYPE.TERRAIN, - ).any() - self._has_terrain = (self._solver.geoms_info.type.to_numpy() == gs.GEOM_TYPE.TERRAIN).any() - - if self._has_terrain: - self.xyz_max_min = ti.field(dtype=gs.ti_float, shape=self._solver._batch_shape(6)) - self.prism = ti.field(dtype=gs.ti_vec3, shape=self._solver._batch_shape(6)) - - ##---------------- box box - if self._solver._box_box_detection: - # With the existing Box-Box collision detection algorithm, it is not clear where the contact points are - # located depending of the pose and size of each box. In practice, up to 11 contact points have been - # observed. The theoretical worst case scenario would be 2 cubes roughly the same size and same center, - # with transform RPY = (45, 45, 45), resulting in 3 contact points per faces for a total of 16 points. - self.box_MAXCONPAIR = 16 - self.box_depth = ti.field(dtype=gs.ti_float, shape=self._solver._batch_shape(self.box_MAXCONPAIR)) - self.box_points = ti.field(gs.ti_vec3, shape=self._solver._batch_shape(self.box_MAXCONPAIR)) - self.box_pts = ti.field(gs.ti_vec3, shape=self._solver._batch_shape(6)) - self.box_lines = ti.field(gs.ti_vec6, shape=self._solver._batch_shape(4)) - self.box_linesu = ti.field(gs.ti_vec6, shape=self._solver._batch_shape(4)) - self.box_axi = ti.field(gs.ti_vec3, shape=self._solver._batch_shape(3)) - self.box_ppts2 = ti.field(dtype=gs.ti_float, shape=self._solver._batch_shape((4, 2))) - self.box_pu = ti.field(gs.ti_vec3, shape=self._solver._batch_shape(4)) - ##---------------- box box + def _init_verts_connectivity(self, vert_neighbors, vert_neighbor_start, vert_n_neighbors): + if self._solver.n_verts > 0: + self._collider_state.vert_neighbors.from_numpy(vert_neighbors) + self._collider_state.vert_neighbor_start.from_numpy(vert_neighbor_start) + self._collider_state.vert_n_neighbors.from_numpy(vert_n_neighbors) + + def _init_max_contact_pairs(self, n_possible_pairs): + max_collision_pairs = min(self._solver._max_collision_pairs, n_possible_pairs) + max_contact_pairs = max_collision_pairs * self._collider_info.n_contacts_per_pair + + self._collider_state._max_possible_pairs[None] = n_possible_pairs + self._collider_state._max_collision_pairs[None] = max_collision_pairs + self._collider_state._max_contact_pairs[None] = max_contact_pairs + + def _init_terrain_state(self): + if self._collider_info.has_terrain: + solver = self._solver + links_idx = solver.geoms_info.link_idx.to_numpy()[solver.geoms_info.type.to_numpy() == gs.GEOM_TYPE.TERRAIN] + entity = solver._entities[solver.links_info.entity_idx.to_numpy()[links_idx[0]]] + + scale = entity.terrain_scale.astype(gs.np_float) + rc = np.array(entity.terrain_hf.shape, dtype=gs.np_int) + hf = entity.terrain_hf.astype(gs.np_float) * scale[1] + xyz_maxmin = np.array( + [rc[0] * scale[0], rc[1] * scale[0], hf.max(), 0, 0, hf.min() - 1.0], + dtype=gs.np_float, + ) - self.reset() + self._collider_state.terrain_hf.from_numpy(hf) + self._collider_state.terrain_rc.from_numpy(rc) + self._collider_state.terrain_scale.from_numpy(scale) + self._collider_state.terrain_xyz_maxmin.from_numpy(xyz_maxmin) def reset(self, envs_idx: npt.NDArray[np.int32] | None = None) -> None: if envs_idx is None: envs_idx = self._solver._scene._envs_idx - self._kernel_reset(envs_idx) + self._kernel_reset( + self._solver._rigid_global_info, self._solver._static_rigid_sim_config, self._collider_state, envs_idx + ) self._contacts_info_cache = {} @ti.kernel def _kernel_reset( - self, + self_unused, + rigid_global_info: ti.template(), + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), envs_idx: ti.types.ndarray(), ): - ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL) + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) for i_b_ in range(envs_idx.shape[0]): i_b = envs_idx[i_b_] - self.first_time[i_b] = 1 - for i_ga in range(self._solver.n_geoms): - for i_gb in range(self._solver.n_geoms): + collider_state.first_time[i_b] = 1 + n_geoms = collider_state.active_buffer.shape[0] + for i_ga in range(n_geoms): + for i_gb in range(n_geoms): # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1 # self.contact_cache[i_ga, i_gb, i_b].penetration = 0.0 - self.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) + collider_state.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) def clear(self, envs_idx=None): if envs_idx is None: envs_idx = self._solver._scene._envs_idx - self._kernel_clear(envs_idx) + self._kernel_clear( + self._solver.links_state, + self._solver.links_info, + self._solver._static_rigid_sim_config, + self._collider_state, + envs_idx, + ) @ti.kernel def _kernel_clear( - self, + self_unused, + links_state: array_class.LinksState, + links_info: array_class.LinksInfo, + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), envs_idx: ti.types.ndarray(), ): - ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL) + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) for i_b_ in range(envs_idx.shape[0]): i_b = envs_idx[i_b_] - if ti.static(self._solver._use_hibernation): - self.n_contacts_hibernated[i_b] = 0 + if ti.static(static_rigid_sim_config.use_hibernation): + collider_state.n_contacts_hibernated[i_b] = 0 # advect hibernated contacts - for i_c in range(self.n_contacts[i_b]): - i_la = self.contact_data[i_c, i_b].link_a - i_lb = self.contact_data[i_c, i_b].link_b + for i_c in range(collider_state.n_contacts[i_b]): + i_la = collider_state.contact_data[i_c, i_b].link_a + i_lb = collider_state.contact_data[i_c, i_b].link_b - I_la = [i_la, i_b] if ti.static(self._solver._options.batch_links_info) else i_la - I_lb = [i_lb, i_b] if ti.static(self._solver._options.batch_links_info) else i_lb + I_la = [i_la, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_la + I_lb = [i_lb, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_lb # pair of hibernated-fixed links -> hibernated contact # TODO: we should also include hibernated-hibernated links and wake up the whole contact island # once a new collision is detected - if (self._solver.links_state[i_la, i_b].hibernated and self._solver.links_info[I_lb].is_fixed) or ( - self._solver.links_state[i_lb, i_b].hibernated and self._solver.links_info[I_la].is_fixed + if (links_state[i_la, i_b].hibernated and links_info[I_lb].is_fixed) or ( + links_state[i_lb, i_b].hibernated and links_info[I_la].is_fixed ): - i_c_hibernated = self.n_contacts_hibernated[i_b] + i_c_hibernated = collider_state.n_contacts_hibernated[i_b] if i_c != i_c_hibernated: - self.contact_data[i_c_hibernated, i_b] = self.contact_data[i_c, i_b] - self.n_contacts_hibernated[i_b] = i_c_hibernated + 1 + collider_state.contact_data[i_c_hibernated, i_b] = collider_state.contact_data[i_c, i_b] + collider_state.n_contacts_hibernated[i_b] = i_c_hibernated + 1 - self.n_contacts[i_b] = self.n_contacts_hibernated[i_b] + collider_state.n_contacts[i_b] = collider_state.n_contacts_hibernated[i_b] else: - self.n_contacts[i_b] = 0 + collider_state.n_contacts[i_b] = 0 def detection(self) -> None: # from genesis.utils.tools import create_timer self._contacts_info_cache = {} # timer = create_timer(name="69477ab0-5e75-47cb-a4a5-d4eebd9336ca", level=3, ti_sync=True, skip_first_call=True) - self._func_update_aabbs() + self._func_update_aabbs(self._solver) # timer.stamp("func_update_aabbs") - self._func_broad_phase() + self._func_broad_phase( + self._solver.links_state, + self._solver.links_info, + self._solver.geoms_state, + self._solver.geoms_info, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + self._collider_state, + ) # timer.stamp("func_broad_phase") - self._func_narrow_phase_convex_vs_convex() - self._func_narrow_phase_convex_specializations() + self._func_narrow_phase_convex_vs_convex( + self._solver.links_state, + self._solver.links_info, + self._solver.geoms_state, + self._solver.geoms_info, + self._solver.geoms_init_AABB, + self._solver.verts_info, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + self._collider_state, + self._collider_info, + self._mpr, + self._gjk, + self._solver.sdf, + ) + self._func_narrow_phase_convex_specializations( + self._solver.geoms_state, + self._solver.geoms_info, + self._solver.geoms_init_AABB, + self._solver.verts_info, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + self._collider_state, + self._collider_info, + self._mpr, + ) # timer.stamp("func_narrow_phase") - if self._has_terrain: - self._func_narrow_phase_any_vs_terrain() + if self._collider_info.has_terrain: + self._func_narrow_phase_any_vs_terrain( + self._solver.geoms_state, + self._solver.geoms_info, + self._solver.geoms_init_AABB, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + self._collider_state, + self._collider_info, + self._mpr, + ) # timer.stamp("_func_narrow_phase_any_vs_terrain") - if self._has_nonconvex_nonterrain: - self._func_narrow_phase_nonconvex_vs_nonterrain() + if self._collider_info.has_nonconvex_nonterrain: + self._func_narrow_phase_nonconvex_vs_nonterrain( + self._solver.links_state, + self._solver.links_info, + self._solver.geoms_state, + self._solver.geoms_info, + self._solver.geoms_init_AABB, + self._solver.verts_info, + self._solver.edges_info, + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + self._collider_state, + self._collider_info, + self._solver.sdf, + ) # timer.stamp("_func_narrow_phase_nonconvex_vs_nonterrain") @ti.func - def _func_point_in_geom_aabb(self, point, i_g, i_b): - return (point < self._solver.geoms_state[i_g, i_b].aabb_max).all() and ( - point > self._solver.geoms_state[i_g, i_b].aabb_min - ).all() + def _func_point_in_geom_aabb( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + point: ti.types.vector(3), + i_g, + i_b, + ): + return (point < geoms_state[i_g, i_b].aabb_max).all() and (point > geoms_state[i_g, i_b].aabb_min).all() @ti.func - def _func_is_geom_aabbs_overlap(self, i_ga, i_gb, i_b): + def _func_is_geom_aabbs_overlap( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + i_ga, + i_gb, + i_b, + ): return not ( - (self._solver.geoms_state[i_ga, i_b].aabb_max <= self._solver.geoms_state[i_gb, i_b].aabb_min).any() - or (self._solver.geoms_state[i_ga, i_b].aabb_min >= self._solver.geoms_state[i_gb, i_b].aabb_max).any() + (geoms_state[i_ga, i_b].aabb_max <= geoms_state[i_gb, i_b].aabb_min).any() + or (geoms_state[i_ga, i_b].aabb_min >= geoms_state[i_gb, i_b].aabb_max).any() ) @ti.func - def _func_find_intersect_midpoint(self, i_ga, i_gb): + def _func_find_intersect_midpoint( + self_unused, geoms_state: array_class.GeomsState, geoms_info: array_class.GeomsInfo, i_ga, i_gb + ): # return the center of the intersecting AABB of AABBs of two geoms - intersect_lower = ti.max(self._solver.geoms_state[i_ga].aabb_min, self._solver.geoms_state[i_gb].aabb_min) - intersect_upper = ti.min(self._solver.geoms_state[i_ga].aabb_max, self._solver.geoms_state[i_gb].aabb_max) + intersect_lower = ti.max(geoms_state[i_ga].aabb_min, geoms_state[i_gb].aabb_min) + intersect_upper = ti.min(geoms_state[i_ga].aabb_max, geoms_state[i_gb].aabb_max) return 0.5 * (intersect_lower + intersect_upper) @ti.func - def _func_contact_sphere_sdf(self, i_ga, i_gb, i_b): - ga_info = self._solver.geoms_info[i_ga] + def _func_contact_sphere_sdf( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + sdf: ti.template(), + i_ga, + i_gb, + i_b, + ): + ga_info = geoms_info[i_ga] is_col = False penetration = gs.ti_float(0.0) normal = ti.Vector.zero(gs.ti_float, 3) contact_pos = ti.Vector.zero(gs.ti_float, 3) - sphere_center = self._solver.geoms_state[i_ga, i_b].pos + sphere_center = geoms_state[i_ga, i_b].pos sphere_radius = ga_info.data[0] - center_to_b_dist = self._solver.sdf.sdf_world(sphere_center, i_gb, i_b) + center_to_b_dist = sdf.sdf_world(sphere_center, i_gb, i_b) if center_to_b_dist < sphere_radius: is_col = True - normal = self._solver.sdf.sdf_normal_world(sphere_center, i_gb, i_b) + normal = sdf.sdf_normal_world(sphere_center, i_gb, i_b) penetration = sphere_radius - center_to_b_dist contact_pos = sphere_center - (sphere_radius - 0.5 * penetration) * normal return is_col, normal, penetration, contact_pos @ti.func - def _func_contact_vertex_sdf(self, i_ga, i_gb, i_b): - ga_info = self._solver.geoms_info[i_ga] - ga_pos = self._solver.geoms_state[i_ga, i_b].pos - ga_quat = self._solver.geoms_state[i_ga, i_b].quat + def _func_contact_vertex_sdf( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + verts_info: array_class.VertsInfo, + sdf: ti.template(), + i_ga, + i_gb, + i_b, + ): + ga_info = geoms_info[i_ga] + ga_pos = geoms_state[i_ga, i_b].pos + ga_quat = geoms_state[i_ga, i_b].quat is_col = False penetration = gs.ti_float(0.0) @@ -386,9 +493,9 @@ def _func_contact_vertex_sdf(self, i_ga, i_gb, i_b): contact_pos = ti.Vector.zero(gs.ti_float, 3) for i_v in range(ga_info.vert_start, ga_info.vert_end): - vertex_pos = gu.ti_transform_by_trans_quat(self._solver.verts_info[i_v].init_pos, ga_pos, ga_quat) - if self._func_point_in_geom_aabb(vertex_pos, i_gb, i_b): - new_penetration = -self._solver.sdf.sdf_world(vertex_pos, i_gb, i_b) + vertex_pos = gu.ti_transform_by_trans_quat(verts_info[i_v].init_pos, ga_pos, ga_quat) + if self_unused._func_point_in_geom_aabb(geoms_state, geoms_info, vertex_pos, i_gb, i_b): + new_penetration = -sdf.sdf_world(vertex_pos, i_gb, i_b) if new_penetration > penetration: is_col = True contact_pos = vertex_pos @@ -396,7 +503,7 @@ def _func_contact_vertex_sdf(self, i_ga, i_gb, i_b): # Compute contact normal only once, and only in case of contact if is_col: - normal = self._solver.sdf.sdf_normal_world(contact_pos, i_gb, i_b) + normal = sdf.sdf_normal_world(contact_pos, i_gb, i_b) # The contact point must be offsetted by half the penetration depth contact_pos += 0.5 * penetration * normal @@ -404,34 +511,44 @@ def _func_contact_vertex_sdf(self, i_ga, i_gb, i_b): return is_col, normal, penetration, contact_pos @ti.func - def _func_contact_edge_sdf(self, i_ga, i_gb, i_b): - ga_info = self._solver.geoms_info[i_ga] - ga_state = self._solver.geoms_state[i_ga, i_b] + def _func_contact_edge_sdf( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + verts_info: array_class.VertsInfo, + edges_info: array_class.EdgesInfo, + sdf: ti.template(), + i_ga, + i_gb, + i_b, + ): + ga_info = geoms_info[i_ga] + ga_state = geoms_state[i_ga, i_b] is_col = False penetration = gs.ti_float(0.0) normal = ti.Vector.zero(gs.ti_float, 3) contact_pos = ti.Vector.zero(gs.ti_float, 3) - ga_sdf_cell_size = self._solver.sdf.geoms_info[i_ga].sdf_cell_size + ga_sdf_cell_size = sdf.geoms_info[i_ga].sdf_cell_size for i_e in range(ga_info.edge_start, ga_info.edge_end): - cur_length = self._solver.edges_info[i_e].length + cur_length = edges_info[i_e].length if cur_length > ga_sdf_cell_size: - i_v0 = self._solver.edges_info[i_e].v0 - i_v1 = self._solver.edges_info[i_e].v1 + i_v0 = edges_info[i_e].v0 + i_v1 = edges_info[i_e].v1 - p_0 = gu.ti_transform_by_trans_quat(self._solver.verts_info[i_v0].init_pos, ga_state.pos, ga_state.quat) - p_1 = gu.ti_transform_by_trans_quat(self._solver.verts_info[i_v1].init_pos, ga_state.pos, ga_state.quat) + p_0 = gu.ti_transform_by_trans_quat(verts_info[i_v0].init_pos, ga_state.pos, ga_state.quat) + p_1 = gu.ti_transform_by_trans_quat(verts_info[i_v1].init_pos, ga_state.pos, ga_state.quat) vec_01 = gu.ti_normalize(p_1 - p_0) - sdf_grad_0_b = self._solver.sdf.sdf_grad_world(p_0, i_gb, i_b) - sdf_grad_1_b = self._solver.sdf.sdf_grad_world(p_1, i_gb, i_b) + sdf_grad_0_b = sdf.sdf_grad_world(p_0, i_gb, i_b) + sdf_grad_1_b = sdf.sdf_grad_world(p_1, i_gb, i_b) # check if the edge on a is facing towards mesh b (I am not 100% sure about this, subject to removal) - sdf_grad_0_a = self._solver.sdf.sdf_grad_world(p_0, i_ga, i_b) - sdf_grad_1_a = self._solver.sdf.sdf_grad_world(p_1, i_ga, i_b) + sdf_grad_0_a = sdf.sdf_grad_world(p_0, i_ga, i_b) + sdf_grad_1_a = sdf.sdf_grad_world(p_1, i_ga, i_b) normal_edge_0 = sdf_grad_0_a - sdf_grad_0_a.dot(vec_01) * vec_01 normal_edge_1 = sdf_grad_1_a - sdf_grad_1_a.dot(vec_01) * vec_01 @@ -442,18 +559,18 @@ def _func_contact_edge_sdf(self, i_ga, i_gb, i_b): while cur_length > ga_sdf_cell_size: p_mid = 0.5 * (p_0 + p_1) - if self._solver.sdf.sdf_grad_world(p_mid, i_gb, i_b).dot(vec_01) < 0: + if sdf.sdf_grad_world(p_mid, i_gb, i_b).dot(vec_01) < 0: p_0 = p_mid else: p_1 = p_mid cur_length = 0.5 * cur_length p = 0.5 * (p_0 + p_1) - new_penetration = -self._solver.sdf.sdf_world(p, i_gb, i_b) + new_penetration = -sdf.sdf_world(p, i_gb, i_b) if new_penetration > penetration: is_col = True - normal = self._solver.sdf.sdf_normal_world(p, i_gb, i_b) + normal = sdf.sdf_normal_world(p, i_gb, i_b) contact_pos = p penetration = new_penetration @@ -463,12 +580,23 @@ def _func_contact_edge_sdf(self, i_ga, i_gb, i_b): return is_col, normal, penetration, contact_pos @ti.func - def _func_contact_convex_convex_sdf(self, i_ga, i_gb, i_b, i_va_ws): - gb_vert_start = self._solver.geoms_info[i_gb].vert_start - ga_pos = self._solver.geoms_state[i_ga, i_b].pos - ga_quat = self._solver.geoms_state[i_ga, i_b].quat - gb_pos = self._solver.geoms_state[i_gb, i_b].pos - gb_quat = self._solver.geoms_state[i_gb, i_b].quat + def _func_contact_convex_convex_sdf( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + verts_info: array_class.VertsInfo, + collider_state: ti.template(), + sdf: ti.template(), + i_ga, + i_gb, + i_b, + i_va_ws, + ): + gb_vert_start = geoms_info[i_gb].vert_start + ga_pos = geoms_state[i_ga, i_b].pos + ga_quat = geoms_state[i_ga, i_b].quat + gb_pos = geoms_state[i_gb, i_b].pos + gb_quat = geoms_state[i_gb, i_b].quat is_col = False penetration = gs.ti_float(0.0) @@ -478,21 +606,20 @@ def _func_contact_convex_convex_sdf(self, i_ga, i_gb, i_b, i_va_ws): i_va = i_va_ws if i_va == -1: # start traversing on the vertex graph with a smart initial vertex - pos_vb = gu.ti_transform_by_trans_quat(self._solver.verts_info[gb_vert_start].init_pos, gb_pos, gb_quat) - i_va = self._solver.sdf._func_find_closest_vert(pos_vb, i_ga, i_b) + pos_vb = gu.ti_transform_by_trans_quat(verts_info[gb_vert_start].init_pos, gb_pos, gb_quat) + i_va = sdf._func_find_closest_vert(pos_vb, i_ga, i_b) i_v_closest = i_va - pos_v_closest = gu.ti_transform_by_trans_quat(self._solver.verts_info[i_v_closest].init_pos, ga_pos, ga_quat) - sd_v_closest = self._solver.sdf.sdf_world(pos_v_closest, i_gb, i_b) + pos_v_closest = gu.ti_transform_by_trans_quat(verts_info[i_v_closest].init_pos, ga_pos, ga_quat) + sd_v_closest = sdf.sdf_world(pos_v_closest, i_gb, i_b) while True: for i_neighbor_ in range( - self.vert_neighbor_start[i_va], self.vert_neighbor_start[i_va] + self.vert_n_neighbors[i_va] + collider_state.vert_neighbor_start[i_va], + collider_state.vert_neighbor_start[i_va] + collider_state.vert_n_neighbors[i_va], ): - i_neighbor = self.vert_neighbors[i_neighbor_] - pos_neighbor = gu.ti_transform_by_trans_quat( - self._solver.verts_info[i_neighbor].init_pos, ga_pos, ga_quat - ) - sd_neighbor = self._solver.sdf.sdf_world(pos_neighbor, i_gb, i_b) + i_neighbor = collider_state.vert_neighbors[i_neighbor_] + pos_neighbor = gu.ti_transform_by_trans_quat(verts_info[i_neighbor].init_pos, ga_pos, ga_quat) + sd_neighbor = sdf.sdf_world(pos_neighbor, i_gb, i_b) if ( sd_neighbor < sd_v_closest - 1e-5 ): # 1e-5 (0.01mm) to avoid endless loop due to numerical instability @@ -509,26 +636,27 @@ def _func_contact_convex_convex_sdf(self, i_ga, i_gb, i_b, i_va_ws): pos_a = pos_v_closest if sd_v_closest < 0: is_col = True - normal = self._solver.sdf.sdf_normal_world(pos_a, i_gb, i_b) + normal = sdf.sdf_normal_world(pos_a, i_gb, i_b) contact_pos = pos_a penetration = -sd_v_closest else: # check edge surrounding it for i_neighbor_ in range( - self.vert_neighbor_start[i_va], self.vert_neighbor_start[i_va] + self.vert_n_neighbors[i_va] + collider_state.vert_neighbor_start[i_va], + collider_state.vert_neighbor_start[i_va] + collider_state.vert_n_neighbors[i_va], ): - i_neighbor = self.vert_neighbors[i_neighbor_] + i_neighbor = collider_state.vert_neighbors[i_neighbor_] p_0 = pos_v_closest - p_1 = gu.ti_transform_by_trans_quat(self._solver.verts_info[i_neighbor].init_pos, ga_pos, ga_quat) + p_1 = gu.ti_transform_by_trans_quat(verts_info[i_neighbor].init_pos, ga_pos, ga_quat) vec_01 = gu.ti_normalize(p_1 - p_0) - sdf_grad_0_b = self._solver.sdf.sdf_grad_world(p_0, i_gb, i_b) - sdf_grad_1_b = self._solver.sdf.sdf_grad_world(p_1, i_gb, i_b) + sdf_grad_0_b = sdf.sdf_grad_world(p_0, i_gb, i_b) + sdf_grad_1_b = sdf.sdf_grad_world(p_1, i_gb, i_b) # check if the edge on a is facing towards mesh b (I am not 100% sure about this, subject to removal) - sdf_grad_0_a = self._solver.sdf.sdf_grad_world(p_0, i_ga, i_b) - sdf_grad_1_a = self._solver.sdf.sdf_grad_world(p_1, i_ga, i_b) + sdf_grad_0_a = sdf.sdf_grad_world(p_0, i_ga, i_b) + sdf_grad_1_a = sdf.sdf_grad_world(p_1, i_ga, i_b) normal_edge_0 = sdf_grad_0_a - sdf_grad_0_a.dot(vec_01) * vec_01 normal_edge_1 = sdf_grad_1_a - sdf_grad_1_a.dot(vec_01) * vec_01 @@ -536,10 +664,10 @@ def _func_contact_convex_convex_sdf(self, i_ga, i_gb, i_b, i_va_ws): # check if closest point is between the two points if sdf_grad_0_b.dot(vec_01) < 0 and sdf_grad_1_b.dot(vec_01) > 0: cur_length = (p_1 - p_0).norm() - ga_sdf_cell_size = self._solver.sdf.geoms_info[i_ga].sdf_cell_size + ga_sdf_cell_size = sdf.geoms_info[i_ga].sdf_cell_size while cur_length > ga_sdf_cell_size: p_mid = 0.5 * (p_0 + p_1) - if self._solver.sdf.sdf_grad_world(p_mid, i_gb, i_b).dot(vec_01) < 0: + if sdf.sdf_grad_world(p_mid, i_gb, i_b).dot(vec_01) < 0: p_0 = p_mid else: p_1 = p_mid @@ -548,11 +676,11 @@ def _func_contact_convex_convex_sdf(self, i_ga, i_gb, i_b, i_va_ws): p = 0.5 * (p_0 + p_1) - new_penetration = -self._solver.sdf.sdf_world(p, i_gb, i_b) + new_penetration = -sdf.sdf_world(p, i_gb, i_b) if new_penetration > 0: is_col = True - normal = self._solver.sdf.sdf_normal_world(p, i_gb, i_b) + normal = sdf.sdf_normal_world(p, i_gb, i_b) contact_pos = p penetration = new_penetration break @@ -560,13 +688,24 @@ def _func_contact_convex_convex_sdf(self, i_ga, i_gb, i_b, i_va_ws): return is_col, normal, penetration, contact_pos, i_va @ti.func - def _func_contact_mpr_terrain(self, i_ga, i_gb, i_b): - ga_pos, ga_quat = self._solver.geoms_state[i_ga, i_b].pos, self._solver.geoms_state[i_ga, i_b].quat - gb_pos, gb_quat = self._solver.geoms_state[i_gb, i_b].pos, self._solver.geoms_state[i_gb, i_b].quat + def _func_contact_mpr_terrain( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + collider_state: ti.template(), + collider_info: ti.template(), + mpr: ti.template(), + i_ga, + i_gb, + i_b, + ): + ga_pos, ga_quat = geoms_state[i_ga, i_b].pos, geoms_state[i_ga, i_b].quat + gb_pos, gb_quat = geoms_state[i_gb, i_b].pos, geoms_state[i_gb, i_b].quat margin = gs.ti_float(0.0) is_return = False - tolerance = self._func_compute_tolerance(i_ga, i_gb, i_b) + tolerance = self_unused._func_compute_tolerance(geoms_info, geoms_init_AABB, collider_info, i_ga, i_gb, i_b) # pos = self._solver.geoms_state[i_ga, i_b].pos - self._solver.geoms_state[i_gb, i_b].pos # for i in range(3): # if self._solver.terrain_xyz_maxmin[i] < pos[i] - r2 - margin or \ @@ -574,13 +713,11 @@ def _func_contact_mpr_terrain(self, i_ga, i_gb, i_b): # is_return = True if not is_return: - self._solver.geoms_state[i_ga, i_b].pos, self._solver.geoms_state[i_ga, i_b].quat = ( - gu.ti_transform_pos_quat_by_trans_quat( - ga_pos - self._solver.geoms_state[i_gb, i_b].pos, - ga_quat, - ti.Vector.zero(gs.ti_float, 3), - gu.ti_inv_quat(self._solver.geoms_state[i_gb, i_b].quat), - ) + geoms_state[i_ga, i_b].pos, geoms_state[i_ga, i_b].quat = gu.ti_transform_pos_quat_by_trans_quat( + ga_pos - geoms_state[i_gb, i_b].pos, + ga_quat, + ti.Vector.zero(gs.ti_float, 3), + gu.ti_inv_quat(geoms_state[i_gb, i_b].quat), ) for i_axis, i_m in ti.ndrange(3, 2): @@ -589,60 +726,67 @@ def _func_contact_mpr_terrain(self, i_ga, i_gb, i_b): direction[i_axis] = 1.0 else: direction[i_axis] = -1.0 - v1 = self._mpr.support_driver(direction, i_ga, i_b) - self.xyz_max_min[3 * i_m + i_axis, i_b] = v1[i_axis] + v1 = mpr.support_driver(direction, i_ga, i_b) + collider_state.xyz_max_min[3 * i_m + i_axis, i_b] = v1[i_axis] for i in ti.static(range(3)): - self.prism[i, i_b][2] = self._solver.terrain_xyz_maxmin[5] + collider_state.prism[i, i_b][2] = collider_state.terrain_xyz_maxmin[5] if ( - self._solver.terrain_xyz_maxmin[i] < self.xyz_max_min[i + 3, i_b] - margin - or self._solver.terrain_xyz_maxmin[i + 3] > self.xyz_max_min[i, i_b] + margin + collider_state.terrain_xyz_maxmin[i] < collider_state.xyz_max_min[i + 3, i_b] - margin + or collider_state.terrain_xyz_maxmin[i + 3] > collider_state.xyz_max_min[i, i_b] + margin ): is_return = True if not is_return: - sh = self._solver.terrain_scale[0] - r_min = gs.ti_int(ti.floor((self.xyz_max_min[3, i_b] - self._solver.terrain_xyz_maxmin[3]) / sh)) - r_max = gs.ti_int(ti.ceil((self.xyz_max_min[0, i_b] - self._solver.terrain_xyz_maxmin[3]) / sh)) - c_min = gs.ti_int(ti.floor((self.xyz_max_min[4, i_b] - self._solver.terrain_xyz_maxmin[4]) / sh)) - c_max = gs.ti_int(ti.ceil((self.xyz_max_min[1, i_b] - self._solver.terrain_xyz_maxmin[4]) / sh)) + sh = collider_state.terrain_scale[0] + r_min = gs.ti_int( + ti.floor((collider_state.xyz_max_min[3, i_b] - collider_state.terrain_xyz_maxmin[3]) / sh) + ) + r_max = gs.ti_int( + ti.ceil((collider_state.xyz_max_min[0, i_b] - collider_state.terrain_xyz_maxmin[3]) / sh) + ) + c_min = gs.ti_int( + ti.floor((collider_state.xyz_max_min[4, i_b] - collider_state.terrain_xyz_maxmin[4]) / sh) + ) + c_max = gs.ti_int( + ti.ceil((collider_state.xyz_max_min[1, i_b] - collider_state.terrain_xyz_maxmin[4]) / sh) + ) r_min = ti.max(0, r_min) c_min = ti.max(0, c_min) - r_max = ti.min(self._solver.terrain_rc[0] - 1, r_max) - c_max = ti.min(self._solver.terrain_rc[1] - 1, c_max) + r_max = ti.min(collider_state.terrain_rc[0] - 1, r_max) + c_max = ti.min(collider_state.terrain_rc[1] - 1, c_max) n_con = 0 for r in range(r_min, r_max): nvert = 0 for c in range(c_min, c_max + 1): for i in range(2): - if n_con < ti.static(self._n_contacts_per_pair): + if n_con < ti.static(collider_info.n_contacts_per_pair): nvert = nvert + 1 - self.add_prism_vert( - sh * (r + i) + self._solver.terrain_xyz_maxmin[3], - sh * c + self._solver.terrain_xyz_maxmin[4], - self._solver.terrain_hf[r + i, c] + margin, + self_unused.add_prism_vert( + collider_state, + sh * (r + i) + collider_state.terrain_xyz_maxmin[3], + sh * c + collider_state.terrain_xyz_maxmin[4], + collider_state.terrain_hf[r + i, c] + margin, i_b, ) if nvert > 2 and ( - self.prism[3, i_b][2] >= self.xyz_max_min[5, i_b] - or self.prism[4, i_b][2] >= self.xyz_max_min[5, i_b] - or self.prism[5, i_b][2] >= self.xyz_max_min[5, i_b] + collider_state.prism[3, i_b][2] >= collider_state.xyz_max_min[5, i_b] + or collider_state.prism[4, i_b][2] >= collider_state.xyz_max_min[5, i_b] + or collider_state.prism[5, i_b][2] >= collider_state.xyz_max_min[5, i_b] ): - center_a = gu.ti_transform_by_trans_quat( - self._solver.geoms_info[i_ga].center, ga_pos, ga_quat - ) + center_a = gu.ti_transform_by_trans_quat(geoms_info[i_ga].center, ga_pos, ga_quat) center_b = ti.Vector.zero(gs.ti_float, 3) for i_p in ti.static(range(6)): - center_b = center_b + self.prism[i_p, i_b] + center_b = center_b + collider_state.prism[i_p, i_b] center_b = center_b / 6.0 - self._solver.geoms_state[i_gb, i_b].pos = ti.Vector.zero(gs.ti_float, 3) - self._solver.geoms_state[i_gb, i_b].quat = gu.ti_identity_quat() + geoms_state[i_gb, i_b].pos = ti.Vector.zero(gs.ti_float, 3) + geoms_state[i_gb, i_b].quat = gu.ti_identity_quat() - is_col, normal, penetration, contact_pos = self._mpr.func_mpr_contact_from_centers( + is_col, normal, penetration, contact_pos = mpr.func_mpr_contact_from_centers( i_ga, i_gb, i_b, center_a, center_b ) if is_col: @@ -651,240 +795,333 @@ def _func_contact_mpr_terrain(self, i_ga, i_gb, i_b): contact_pos = contact_pos + gb_pos valid = True - i_c = self.n_contacts[i_b] + i_c = collider_state.n_contacts[i_b] for j in range(n_con): if ( - contact_pos - self.contact_data[i_c - j - 1, i_b].pos + contact_pos - collider_state.contact_data[i_c - j - 1, i_b].pos ).norm() < tolerance: valid = False break if valid: - self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal, + contact_pos, + penetration, + i_b, + ) n_con = n_con + 1 - self._solver.geoms_state[i_ga, i_b].pos, self._solver.geoms_state[i_ga, i_b].quat = ga_pos, ga_quat - self._solver.geoms_state[i_gb, i_b].pos, self._solver.geoms_state[i_gb, i_b].quat = gb_pos, gb_quat + geoms_state[i_ga, i_b].pos, geoms_state[i_ga, i_b].quat = ga_pos, ga_quat + geoms_state[i_gb, i_b].pos, geoms_state[i_gb, i_b].quat = gb_pos, gb_quat @ti.func - def add_prism_vert(self, x, y, z, i_b): - self.prism[0, i_b] = self.prism[1, i_b] - self.prism[1, i_b] = self.prism[2, i_b] - self.prism[3, i_b] = self.prism[4, i_b] - self.prism[4, i_b] = self.prism[5, i_b] - - self.prism[2, i_b][0] = x - self.prism[5, i_b][0] = x - self.prism[2, i_b][1] = y - self.prism[5, i_b][1] = y - self.prism[5, i_b][2] = z + def add_prism_vert(self_unused, collider_state: ti.template(), x, y, z, i_b): + collider_state.prism[0, i_b] = collider_state.prism[1, i_b] + collider_state.prism[1, i_b] = collider_state.prism[2, i_b] + collider_state.prism[3, i_b] = collider_state.prism[4, i_b] + collider_state.prism[4, i_b] = collider_state.prism[5, i_b] + + collider_state.prism[2, i_b][0] = x + collider_state.prism[5, i_b][0] = x + collider_state.prism[2, i_b][1] = y + collider_state.prism[5, i_b][1] = y + collider_state.prism[5, i_b][2] = z @ti.kernel - def _func_update_aabbs(self): - self._solver._func_update_geom_aabbs() + def _func_update_aabbs(self_unused, solver: ti.template()): + # FIXME: Remove [solver] when Rigid solver migration is done. + solver._func_update_geom_aabbs() @ti.func - def _func_check_collision_valid(self, i_ga, i_gb, i_b): - is_valid = self.collision_pair_validity[i_ga, i_gb] + def _func_check_collision_valid( + self_unused, + links_state: array_class.LinksState, + links_info: array_class.LinksInfo, + geoms_info: array_class.GeomsInfo, + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), + i_ga, + i_gb, + i_b, + ): + is_valid = collider_state.collision_pair_validity[i_ga, i_gb] # hibernated <-> fixed links - if ti.static(self._solver._use_hibernation): - i_la = self._solver.geoms_info[i_ga].link_idx - i_lb = self._solver.geoms_info[i_gb].link_idx - I_la = [i_la, i_b] if ti.static(self._solver._options.batch_links_info) else i_la - I_lb = [i_lb, i_b] if ti.static(self._solver._options.batch_links_info) else i_lb - - if (self._solver.links_state[i_la, i_b].hibernated and self._solver.links_info[I_lb].is_fixed) or ( - self._solver.links_state[i_lb, i_b].hibernated and self._solver.links_info[I_la].is_fixed + if ti.static(static_rigid_sim_config.use_hibernation): + i_la = geoms_info[i_ga].link_idx + i_lb = geoms_info[i_gb].link_idx + I_la = [i_la, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_la + I_lb = [i_lb, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_lb + + if (links_state[i_la, i_b].hibernated and links_info[I_lb].is_fixed) or ( + links_state[i_lb, i_b].hibernated and links_info[I_la].is_fixed ): is_valid = False return is_valid @ti.kernel - def _func_broad_phase(self): + def _func_broad_phase( + self_unused, + links_state: array_class.LinksState, + links_info: array_class.LinksInfo, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + rigid_global_info: ti.template(), + static_rigid_sim_config: ti.template(), + # we will use ColliderBroadPhaseBuffer as typing after Hugh adds array_struct feature to taichi + collider_state: ti.template(), + ): """ Sweep and Prune (SAP) for broad-phase collision detection. - This function sorts the geometry axis-aligned bounding boxes (AABBs) along a specified axis and checks for potential collision pairs based on the AABB overlap. + This function sorts the geometry axis-aligned bounding boxes (AABBs) along a specified axis and checks for + potential collision pairs based on the AABB overlap. """ + _B = collider_state.active_buffer.shape[1] + n_geoms = collider_state.active_buffer.shape[0] - ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL) - for i_b in range(self._solver._B): + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) + for i_b in range(_B): axis = 0 # copy updated geom aabbs to buffer for sorting - if self.first_time[i_b]: - for i in range(self._solver.n_geoms): - self.sort_buffer[2 * i, i_b].value = self._solver.geoms_state[i, i_b].aabb_min[axis] - self.sort_buffer[2 * i, i_b].i_g = i - self.sort_buffer[2 * i, i_b].is_max = 0 + if collider_state.first_time[i_b]: + for i in range(n_geoms): + collider_state.sort_buffer[2 * i, i_b].value = geoms_state[i, i_b].aabb_min[axis] + collider_state.sort_buffer[2 * i, i_b].i_g = i + collider_state.sort_buffer[2 * i, i_b].is_max = 0 - self.sort_buffer[2 * i + 1, i_b].value = self._solver.geoms_state[i, i_b].aabb_max[axis] - self.sort_buffer[2 * i + 1, i_b].i_g = i - self.sort_buffer[2 * i + 1, i_b].is_max = 1 + collider_state.sort_buffer[2 * i + 1, i_b].value = geoms_state[i, i_b].aabb_max[axis] + collider_state.sort_buffer[2 * i + 1, i_b].i_g = i + collider_state.sort_buffer[2 * i + 1, i_b].is_max = 1 - self._solver.geoms_state[i, i_b].min_buffer_idx = 2 * i - self._solver.geoms_state[i, i_b].max_buffer_idx = 2 * i + 1 + geoms_state[i, i_b].min_buffer_idx = 2 * i + geoms_state[i, i_b].max_buffer_idx = 2 * i + 1 - self.first_time[i_b] = False + 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 self._solver._use_hibernation): - for i in range(self._solver.n_geoms * 2): - if self.sort_buffer[i, i_b].is_max: - self.sort_buffer[i, i_b].value = self._solver.geoms_state[ - self.sort_buffer[i, i_b].i_g, i_b + if ti.static(not static_rigid_sim_config.use_hibernation): + for i in range(n_geoms * 2): + if collider_state.sort_buffer[i, i_b].is_max: + collider_state.sort_buffer[i, i_b].value = geoms_state[ + collider_state.sort_buffer[i, i_b].i_g, i_b ].aabb_max[axis] else: - self.sort_buffer[i, i_b].value = self._solver.geoms_state[ - self.sort_buffer[i, i_b].i_g, i_b + collider_state.sort_buffer[i, i_b].value = geoms_state[ + collider_state.sort_buffer[i, i_b].i_g, i_b ].aabb_min[axis] # insertion sort, which has complexity near O(n) for nearly sorted array - for i in range(1, 2 * self._solver.n_geoms): - key = self.sort_buffer[i, i_b] + for i in range(1, 2 * n_geoms): + key = collider_state.sort_buffer[i, i_b] j = i - 1 - while j >= 0 and key.value < self.sort_buffer[j, i_b].value: - self.sort_buffer[j + 1, i_b] = self.sort_buffer[j, i_b] + while j >= 0 and key.value < collider_state.sort_buffer[j, i_b].value: + collider_state.sort_buffer[j + 1, i_b] = collider_state.sort_buffer[j, i_b] - if ti.static(self._solver._use_hibernation): - if self.sort_buffer[j, i_b].is_max: - self._solver.geoms_state[self.sort_buffer[j, i_b].i_g, i_b].max_buffer_idx = j + 1 + if ti.static(static_rigid_sim_config.use_hibernation): + if collider_state.sort_buffer[j, i_b].is_max: + geoms_state[collider_state.sort_buffer[j, i_b].i_g, i_b].max_buffer_idx = j + 1 else: - self._solver.geoms_state[self.sort_buffer[j, i_b].i_g, i_b].min_buffer_idx = j + 1 + geoms_state[collider_state.sort_buffer[j, i_b].i_g, i_b].min_buffer_idx = j + 1 j -= 1 - self.sort_buffer[j + 1, i_b] = key + collider_state.sort_buffer[j + 1, i_b] = key - if ti.static(self._solver._use_hibernation): + if ti.static(static_rigid_sim_config.use_hibernation): if key.is_max: - self._solver.geoms_state[key.i_g, i_b].max_buffer_idx = j + 1 + geoms_state[key.i_g, i_b].max_buffer_idx = j + 1 else: - self._solver.geoms_state[key.i_g, i_b].min_buffer_idx = j + 1 + geoms_state[key.i_g, i_b].min_buffer_idx = j + 1 # sweep over the sorted AABBs to find potential collision pairs - self.n_broad_pairs[i_b] = 0 - if ti.static(not self._solver._use_hibernation): + 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 * self._solver.n_geoms): - if not self.sort_buffer[i, i_b].is_max: + for i in range(2 * n_geoms): + if not collider_state.sort_buffer[i, i_b].is_max: for j in range(n_active): - i_ga = self.active_buffer[j, i_b] - i_gb = self.sort_buffer[i, i_b].i_g + i_ga = collider_state.active_buffer[j, i_b] + i_gb = collider_state.sort_buffer[i, i_b].i_g if i_ga > i_gb: i_ga, i_gb = i_gb, i_ga - if not self._func_check_collision_valid(i_ga, i_gb, i_b): + if not self_unused._func_check_collision_valid( + links_state, + links_info, + geoms_info, + static_rigid_sim_config, + collider_state, + i_ga, + i_gb, + i_b, + ): continue - if not self._func_is_geom_aabbs_overlap(i_ga, i_gb, i_b): + if not self_unused._func_is_geom_aabbs_overlap(geoms_state, geoms_info, i_ga, i_gb, i_b): # Clear collision normal cache if not in contact - if ti.static(not self._solver._enable_mujoco_compatibility): + if ti.static(not static_rigid_sim_config.enable_mujoco_compatibility): # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1 - self.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) + collider_state.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) continue - i_p = self.n_broad_pairs[i_b] - if i_p == self._max_collision_pairs: + i_p = collider_state.n_broad_pairs[i_b] + if i_p == collider_state._max_collision_pairs[None]: # print(self._warn_msg_max_collision_pairs) break - self.broad_collision_pairs[i_p, i_b][0] = i_ga - self.broad_collision_pairs[i_p, i_b][1] = i_gb - self.n_broad_pairs[i_b] += 1 + collider_state.broad_collision_pairs[i_p, i_b][0] = i_ga + collider_state.broad_collision_pairs[i_p, i_b][1] = i_gb + collider_state.n_broad_pairs[i_b] += 1 - self.active_buffer[n_active, i_b] = self.sort_buffer[i, i_b].i_g + collider_state.active_buffer[n_active, i_b] = collider_state.sort_buffer[i, i_b].i_g n_active = n_active + 1 else: - i_g_to_remove = self.sort_buffer[i, i_b].i_g + i_g_to_remove = collider_state.sort_buffer[i, i_b].i_g for j in range(n_active): - if self.active_buffer[j, i_b] == i_g_to_remove: + if collider_state.active_buffer[j, i_b] == i_g_to_remove: if j < n_active - 1: for k in range(j, n_active - 1): - self.active_buffer[k, i_b] = self.active_buffer[k + 1, i_b] + collider_state.active_buffer[k, i_b] = collider_state.active_buffer[k + 1, i_b] n_active = n_active - 1 break else: - if self._solver.n_awake_dofs[i_b] > 0: + if rigid_global_info.n_awake_dofs[i_b] > 0: n_active_awake = 0 n_active_hib = 0 - for i in range(2 * self._solver.n_geoms): - is_incoming_geom_hibernated = self._solver.geoms_state[ - self.sort_buffer[i, i_b].i_g, i_b + for i in range(2 * n_geoms): + is_incoming_geom_hibernated = geoms_state[ + collider_state.sort_buffer[i, i_b].i_g, i_b ].hibernated - if not self.sort_buffer[i, i_b].is_max: + if not collider_state.sort_buffer[i, i_b].is_max: # both awake and hibernated geom check with active awake geoms for j in range(n_active_awake): - i_ga = self.active_buffer_awake[j, i_b] - i_gb = self.sort_buffer[i, i_b].i_g + i_ga = collider_state.active_buffer_awake[j, i_b] + i_gb = collider_state.sort_buffer[i, i_b].i_g if i_ga > i_gb: i_ga, i_gb = i_gb, i_ga - if not self._func_check_collision_valid(i_ga, i_gb, i_b): + if not self_unused._func_check_collision_valid( + links_state, + links_info, + geoms_info, + static_rigid_sim_config, + collider_state, + i_ga, + i_gb, + i_b, + ): continue - if not self._func_is_geom_aabbs_overlap(i_ga, i_gb, i_b): + if not self_unused._func_is_geom_aabbs_overlap( + geoms_state, geoms_info, i_ga, i_gb, i_b + ): # Clear collision normal cache if not in contact - if ti.static(not self._solver._enable_mujoco_compatibility): + if ti.static(not static_rigid_sim_config._enable_mujoco_compatibility): # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1 - self.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) + collider_state.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) continue - self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][0] = i_ga - self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][1] = i_gb - self.n_broad_pairs[i_b] = self.n_broad_pairs[i_b] + 1 + collider_state.broad_collision_pairs[collider_state.n_broad_pairs[i_b], i_b][0] = i_ga + collider_state.broad_collision_pairs[collider_state.n_broad_pairs[i_b], i_b][1] = i_gb + collider_state.n_broad_pairs[i_b] = collider_state.n_broad_pairs[i_b] + 1 # if incoming geom is awake, also need to check with hibernated geoms if not is_incoming_geom_hibernated: for j in range(n_active_hib): - i_ga = self.active_buffer_hib[j, i_b] - i_gb = self.sort_buffer[i, i_b].i_g + i_ga = collider_state.active_buffer_hib[j, i_b] + i_gb = collider_state.sort_buffer[i, i_b].i_g if i_ga > i_gb: i_ga, i_gb = i_gb, i_ga - if not self._func_check_collision_valid(i_ga, i_gb, i_b): + if not self_unused._func_check_collision_valid( + links_state, + links_info, + geoms_info, + static_rigid_sim_config, + collider_state, + i_ga, + i_gb, + i_b, + ): continue - if not self._func_is_geom_aabbs_overlap(i_ga, i_gb, i_b): + if not self_unused._func_is_geom_aabbs_overlap( + geoms_state, geoms_info, i_ga, i_gb, i_b + ): # Clear collision normal cache if not in contact # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1 - self.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) + collider_state.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) continue - self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][0] = i_ga - self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][1] = i_gb - self.n_broad_pairs[i_b] = self.n_broad_pairs[i_b] + 1 + collider_state.broad_collision_pairs[collider_state.n_broad_pairs[i_b], i_b][ + 0 + ] = i_ga + collider_state.broad_collision_pairs[collider_state.n_broad_pairs[i_b], i_b][ + 1 + ] = i_gb + collider_state.n_broad_pairs[i_b] = collider_state.n_broad_pairs[i_b] + 1 if is_incoming_geom_hibernated: - self.active_buffer_hib[n_active_hib, i_b] = self.sort_buffer[i, i_b].i_g + collider_state.active_buffer_hib[n_active_hib, i_b] = collider_state.sort_buffer[ + i, i_b + ].i_g n_active_hib = n_active_hib + 1 else: - self.active_buffer_awake[n_active_awake, i_b] = self.sort_buffer[i, i_b].i_g + collider_state.active_buffer_awake[n_active_awake, i_b] = collider_state.sort_buffer[ + i, i_b + ].i_g n_active_awake = n_active_awake + 1 else: - i_g_to_remove = self.sort_buffer[i, i_b].i_g + i_g_to_remove = collider_state.sort_buffer[i, i_b].i_g if is_incoming_geom_hibernated: for j in range(n_active_hib): - if self.active_buffer_hib[j, i_b] == i_g_to_remove: + if collider_state.active_buffer_hib[j, i_b] == i_g_to_remove: if j < n_active_hib - 1: for k in range(j, n_active_hib - 1): - self.active_buffer_hib[k, i_b] = self.active_buffer_hib[k + 1, i_b] + collider_state.active_buffer_hib[k, i_b] = ( + collider_state.active_buffer_hib[k + 1, i_b] + ) n_active_hib = n_active_hib - 1 break else: for j in range(n_active_awake): - if self.active_buffer_awake[j, i_b] == i_g_to_remove: + if collider_state.active_buffer_awake[j, i_b] == i_g_to_remove: if j < n_active_awake - 1: for k in range(j, n_active_awake - 1): - self.active_buffer_awake[k, i_b] = self.active_buffer_awake[k + 1, i_b] + collider_state.active_buffer_awake[k, i_b] = ( + collider_state.active_buffer_awake[k + 1, i_b] + ) n_active_awake = n_active_awake - 1 break @ti.kernel - def _func_narrow_phase_convex_vs_convex(self): + def _func_narrow_phase_convex_vs_convex( + self_unused, + links_state: array_class.LinksState, + links_info: array_class.LinksInfo, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + verts_info: array_class.VertsInfo, + rigid_global_info: ti.template(), + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), + collider_info: ti.template(), + mpr: ti.template(), + gjk: ti.template(), + sdf: ti.template(), + ): """ NOTE: for a single non-batched scene with a lot of collisioin pairs, it will be faster if we also parallelize over `self.n_collision_pairs`. However, parallelize over both B and collision_pairs (instead of only over B) leads to significantly slow performance for batched scene. @@ -893,101 +1130,194 @@ def _func_narrow_phase_convex_vs_convex(self): Updated NOTE & TODO: For a HUGE scene with numerous bodies, it's also reasonable to run on GPU. Let's save this for later. Update2: Now we use n_broad_pairs instead of n_collision_pairs, so we probably need to think about how to handle non-batched large scene better. """ - ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL) - for i_b in range(self._solver._B): - for i_pair in range(self.n_broad_pairs[i_b]): - i_ga = self.broad_collision_pairs[i_pair, i_b][0] - i_gb = self.broad_collision_pairs[i_pair, i_b][1] + _B = collider_state.active_buffer.shape[1] + + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) + for i_b in range(_B): + for i_pair in range(collider_state.n_broad_pairs[i_b]): + i_ga = collider_state.broad_collision_pairs[i_pair, i_b][0] + i_gb = collider_state.broad_collision_pairs[i_pair, i_b][1] - if self._solver.geoms_info[i_ga].type > self._solver.geoms_info[i_gb].type: + if geoms_info[i_ga].type > geoms_info[i_gb].type: i_ga, i_gb = i_gb, i_ga if ( - self._solver.geoms_info[i_ga].is_convex - and self._solver.geoms_info[i_gb].is_convex - and not self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.TERRAIN + geoms_info[i_ga].is_convex + and geoms_info[i_gb].is_convex + and not geoms_info[i_gb].type == gs.GEOM_TYPE.TERRAIN and not ( - self._solver._box_box_detection - and self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.BOX - and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX + static_rigid_sim_config.box_box_detection + and geoms_info[i_ga].type == gs.GEOM_TYPE.BOX + and geoms_info[i_gb].type == gs.GEOM_TYPE.BOX ) ): if ti.static(sys.platform == "darwin"): - self._func_convex_convex_contact(i_ga, i_gb, i_b) + self_unused._func_convex_convex_contact( + links_state, + links_info, + geoms_state, + geoms_info, + geoms_init_AABB, + verts_info, + rigid_global_info, + static_rigid_sim_config, + collider_state, + collider_info, + mpr, + gjk, + sdf, + i_ga, + i_gb, + i_b, + ) else: if not ( - self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE - and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX + geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE and geoms_info[i_gb].type == gs.GEOM_TYPE.BOX ): - self._func_convex_convex_contact(i_ga, i_gb, i_b) + self_unused._func_convex_convex_contact( + links_state, + links_info, + geoms_state, + geoms_info, + geoms_init_AABB, + verts_info, + rigid_global_info, + static_rigid_sim_config, + collider_state, + collider_info, + mpr, + gjk, + sdf, + i_ga, + i_gb, + i_b, + ) @ti.kernel - def _func_narrow_phase_convex_specializations(self): - ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL) - for i_b in range(self._solver._B): - for i_pair in range(self.n_broad_pairs[i_b]): - i_ga = self.broad_collision_pairs[i_pair, i_b][0] - i_gb = self.broad_collision_pairs[i_pair, i_b][1] - - if self._solver.geoms_info[i_ga].type > self._solver.geoms_info[i_gb].type: + def _func_narrow_phase_convex_specializations( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + verts_info: array_class.VertsInfo, + rigid_global_info: ti.template(), + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), + collider_info: ti.template(), + mpr: ti.template(), + ): + _B = collider_state.active_buffer.shape[1] + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) + for i_b in range(_B): + for i_pair in range(collider_state.n_broad_pairs[i_b]): + i_ga = collider_state.broad_collision_pairs[i_pair, i_b][0] + i_gb = collider_state.broad_collision_pairs[i_pair, i_b][1] + + if geoms_info[i_ga].type > geoms_info[i_gb].type: i_ga, i_gb = i_gb, i_ga if ti.static(sys.platform != "darwin"): - if ( - self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE - and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX - ): - self._func_plane_box_contact(i_ga, i_gb, i_b) + if geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE and geoms_info[i_gb].type == gs.GEOM_TYPE.BOX: + self_unused._func_plane_box_contact( + geoms_state, + geoms_info, + geoms_init_AABB, + verts_info, + static_rigid_sim_config, + collider_state, + collider_info, + mpr, + i_ga, + i_gb, + i_b, + ) - if ti.static(self._solver._box_box_detection): - if ( - self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.BOX - and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX - ): - self._func_box_box_contact(i_ga, i_gb, i_b) + if ti.static(static_rigid_sim_config.box_box_detection): + if geoms_info[i_ga].type == gs.GEOM_TYPE.BOX and geoms_info[i_gb].type == gs.GEOM_TYPE.BOX: + self_unused._func_box_box_contact( + geoms_state, geoms_info, collider_state, collider_info, i_ga, i_gb, i_b + ) @ti.kernel - def _func_narrow_phase_any_vs_terrain(self): + def _func_narrow_phase_any_vs_terrain( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + rigid_global_info: ti.template(), + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), + collider_info: ti.template(), + mpr: ti.template(), + ): """ NOTE: for a single non-batched scene with a lot of collisioin pairs, it will be faster if we also parallelize over `self.n_collision_pairs`. However, parallelize over both B and collisioin_pairs (instead of only over B) leads to significantly slow performance for batched scene. We can treat B=0 and B>0 separately, but we will end up with messier code. Therefore, for a big non-batched scene, users are encouraged to simply use `gs.cpu` backend. Updated NOTE & TODO: For a HUGE scene with numerous bodies, it's also reasonable to run on GPU. Let's save this for later. Update2: Now we use n_broad_pairs instead of n_collision_pairs, so we probably need to think about how to handle non-batched large scene better. """ - ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL) - for i_b in range(self._solver._B): - for i_pair in range(self.n_broad_pairs[i_b]): - i_ga = self.broad_collision_pairs[i_pair, i_b][0] - i_gb = self.broad_collision_pairs[i_pair, i_b][1] - - if ti.static(self._has_terrain): - if self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.TERRAIN: + _B = collider_state.active_buffer.shape[1] + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) + for i_b in range(_B): + for i_pair in range(collider_state.n_broad_pairs[i_b]): + i_ga = collider_state.broad_collision_pairs[i_pair, i_b][0] + i_gb = collider_state.broad_collision_pairs[i_pair, i_b][1] + + if ti.static(collider_info.has_terrain): + if geoms_info[i_ga].type == gs.GEOM_TYPE.TERRAIN: i_ga, i_gb = i_gb, i_ga - if self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.TERRAIN: - self._func_contact_mpr_terrain(i_ga, i_gb, i_b) + if geoms_info[i_gb].type == gs.GEOM_TYPE.TERRAIN: + self_unused._func_contact_mpr_terrain( + geoms_state, + geoms_info, + geoms_init_AABB, + collider_state, + collider_info, + mpr, + i_ga, + i_gb, + i_b, + ) @ti.kernel - def _func_narrow_phase_nonconvex_vs_nonterrain(self): + def _func_narrow_phase_nonconvex_vs_nonterrain( + self_unused, + links_state: array_class.LinksState, + links_info: array_class.LinksInfo, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + verts_info: array_class.VertsInfo, + edges_info: array_class.EdgesInfo, + rigid_global_info: ti.template(), + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), + collider_info: ti.template(), + sdf: ti.template(), + ): """ NOTE: for a single non-batched scene with a lot of collisioin pairs, it will be faster if we also parallelize over `self.n_collision_pairs`. However, parallelize over both B and collisioin_pairs (instead of only over B) leads to significantly slow performance for batched scene. We can treat B=0 and B>0 separately, but we will end up with messier code. Therefore, for a big non-batched scene, users are encouraged to simply use `gs.cpu` backend. Updated NOTE & TODO: For a HUGE scene with numerous bodies, it's also reasonable to run on GPU. Let's save this for later. Update2: Now we use n_broad_pairs instead of n_collision_pairs, so we probably need to think about how to handle non-batched large scene better. """ - ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL) - for i_b in range(self._solver._B): - for i_pair in range(self.n_broad_pairs[i_b]): - i_ga = self.broad_collision_pairs[i_pair, i_b][0] - i_gb = self.broad_collision_pairs[i_pair, i_b][1] - - if ti.static(self._has_nonconvex_nonterrain): + _B = collider_state.active_buffer.shape[1] + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) + for i_b in range(_B): + for i_pair in range(collider_state.n_broad_pairs[i_b]): + i_ga = collider_state.broad_collision_pairs[i_pair, i_b][0] + i_gb = collider_state.broad_collision_pairs[i_pair, i_b][1] + + if ti.static(collider_info.has_nonconvex_nonterrain): if ( - not (self._solver.geoms_info[i_ga].is_convex and self._solver.geoms_info[i_gb].is_convex) - and self._solver.geoms_info[i_gb].type != gs.GEOM_TYPE.TERRAIN + not (geoms_info[i_ga].is_convex and geoms_info[i_gb].is_convex) + and geoms_info[i_gb].type != gs.GEOM_TYPE.TERRAIN ): is_col = False - tolerance = self._func_compute_tolerance(i_ga, i_gb, i_b) + tolerance = self_unused._func_compute_tolerance( + geoms_info, geoms_init_AABB, collider_info, i_ga, i_gb, i_b + ) for i in range(2): if i == 1: i_ga, i_gb = i_gb, i_ga @@ -997,36 +1327,62 @@ def _func_narrow_phase_nonconvex_vs_nonterrain(self): normal_i = ti.Vector.zero(gs.ti_float, 3) contact_pos_i = ti.Vector.zero(gs.ti_float, 3) if not is_col: - is_col_i, normal_i, penetration_i, contact_pos_i = self._func_contact_vertex_sdf( - i_ga, i_gb, i_b + is_col_i, normal_i, penetration_i, contact_pos_i = self_unused._func_contact_vertex_sdf( + geoms_state, geoms_info, verts_info, sdf, i_ga, i_gb, i_b ) if is_col_i: - self._func_add_contact(i_ga, i_gb, normal_i, contact_pos_i, penetration_i, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal_i, + contact_pos_i, + penetration_i, + i_b, + ) - if ti.static(self._solver._enable_multi_contact): + if ti.static(static_rigid_sim_config.enable_multi_contact): if not is_col and is_col_i: - ga_state = self._solver.geoms_state[i_ga, i_b] - gb_state = self._solver.geoms_state[i_gb, i_b] + ga_state = geoms_state[i_ga, i_b] + gb_state = geoms_state[i_gb, i_b] ga_pos, ga_quat = ga_state.pos, ga_state.quat gb_pos, gb_quat = gb_state.pos, gb_state.quat # Perturb geom_a around two orthogonal axes to find multiple contacts - axis_0, axis_1 = self._func_contact_orthogonals(i_ga, i_gb, normal_i, i_b) + axis_0, axis_1 = self_unused._func_contact_orthogonals( + links_state, + links_info, + geoms_state, + geoms_info, + geoms_init_AABB, + static_rigid_sim_config, + i_ga, + i_gb, + normal_i, + i_b, + ) n_con = 1 for i_rot in range(1, 5): axis = (2 * (i_rot % 2) - 1) * axis_0 + (1 - 2 * ((i_rot // 2) % 2)) * axis_1 - qrot = gu.ti_rotvec_to_quat(self._mc_perturbation * axis) - self._func_rotate_frame(i_ga, contact_pos_i, qrot, i_b) - self._func_rotate_frame(i_gb, contact_pos_i, gu.ti_inv_quat(qrot), i_b) + qrot = gu.ti_rotvec_to_quat(collider_info.mc_perturbation * axis) + self_unused._func_rotate_frame( + geoms_state, geoms_info, i_ga, contact_pos_i, qrot, i_b + ) + self_unused._func_rotate_frame( + geoms_state, geoms_info, i_gb, contact_pos_i, gu.ti_inv_quat(qrot), i_b + ) - is_col, normal, penetration, contact_pos = self._func_contact_vertex_sdf( - i_ga, i_gb, i_b + is_col, normal, penetration, contact_pos = self_unused._func_contact_vertex_sdf( + geoms_state, geoms_info, verts_info, sdf, i_ga, i_gb, i_b ) if is_col: - if ti.static(not self._solver._enable_mujoco_compatibility): + if ti.static(not static_rigid_sim_config.enable_mujoco_compatibility): # 1. Project the contact point on both geometries # 2. Revert the effect of small rotation # 3. Update contact point @@ -1049,8 +1405,8 @@ def _func_narrow_phase_nonconvex_vs_nonterrain(self): # First-order correction of the normal direction twist_rotvec = ti.math.clamp( normal.cross(normal_i), - -self._mc_perturbation, - self._mc_perturbation, + -collider_info.mc_perturbation, + collider_info.mc_perturbation, ) normal += twist_rotvec.cross(normal) @@ -1061,111 +1417,202 @@ def _func_narrow_phase_nonconvex_vs_nonterrain(self): repeated = False for i_c in range(n_con): if not repeated: - idx_prev = self.n_contacts[i_b] - 1 - i_c - prev_contact = self.contact_data[idx_prev, i_b].pos + idx_prev = collider_state.n_contacts[i_b] - 1 - i_c + prev_contact = collider_state.contact_data[idx_prev, i_b].pos if (contact_pos - prev_contact).norm() < tolerance: repeated = True if not repeated: if penetration > -tolerance: penetration = ti.max(penetration, 0.0) - self._func_add_contact( - i_ga, i_gb, normal, contact_pos, penetration, i_b + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal, + contact_pos, + penetration, + i_b, ) n_con += 1 - self._solver.geoms_state[i_ga, i_b].pos = ga_pos - self._solver.geoms_state[i_ga, i_b].quat = ga_quat - self._solver.geoms_state[i_gb, i_b].pos = gb_pos - self._solver.geoms_state[i_gb, i_b].quat = gb_quat + geoms_state[i_ga, i_b].pos = ga_pos + geoms_state[i_ga, i_b].quat = ga_quat + geoms_state[i_gb, i_b].pos = gb_pos + geoms_state[i_gb, i_b].quat = gb_quat if not is_col: # check edge-edge if vertex-face is not detected - is_col, normal, penetration, contact_pos = self._func_contact_edge_sdf(i_ga, i_gb, i_b) + is_col, normal, penetration, contact_pos = self_unused._func_contact_edge_sdf( + geoms_state, geoms_info, verts_info, edges_info, sdf, i_ga, i_gb, i_b + ) if is_col: - self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal, + contact_pos, + penetration, + i_b, + ) @ti.func - def _func_plane_box_contact(self, i_ga, i_gb, i_b): - ga_info = self._solver.geoms_info[i_ga] - gb_info = self._solver.geoms_info[i_gb] - ga_state = self._solver.geoms_state[i_ga, i_b] - gb_state = self._solver.geoms_state[i_gb, i_b] + def _func_plane_box_contact( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + verts_info: array_class.VertsInfo, + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), + collider_info: ti.template(), + mpr: ti.template(), + i_ga, + i_gb, + i_b, + ): + ga_info = geoms_info[i_ga] + gb_info = geoms_info[i_gb] + ga_state = geoms_state[i_ga, i_b] + gb_state = geoms_state[i_gb, i_b] plane_dir = ti.Vector([ga_info.data[0], ga_info.data[1], ga_info.data[2]], dt=gs.ti_float) plane_dir = gu.ti_transform_by_quat(plane_dir, ga_state.quat) normal = -plane_dir.normalized() - v1, _ = self._mpr.support_field._func_support_box(normal, i_gb, i_b) + v1, _ = mpr.support_field._func_support_box(normal, i_gb, i_b) penetration = normal.dot(v1 - ga_state.pos) if penetration > 0.0: contact_pos = v1 - 0.5 * penetration * normal - self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal, + contact_pos, + penetration, + i_b, + ) - if ti.static(self._solver._enable_multi_contact): + if ti.static(static_rigid_sim_config.enable_multi_contact): n_con = 1 contact_pos_0 = contact_pos - tolerance = self._func_compute_tolerance(i_ga, i_gb, i_b) + tolerance = self_unused._func_compute_tolerance( + geoms_info, geoms_init_AABB, collider_info, i_ga, i_gb, i_b + ) for i_v in range(gb_info.vert_start, gb_info.vert_end): - if n_con < ti.static(self._n_contacts_per_pair): + if n_con < ti.static(collider_info.n_contacts_per_pair): pos_corner = gu.ti_transform_by_trans_quat( - self._solver.verts_info[i_v].init_pos, gb_state.pos, gb_state.quat + verts_info[i_v].init_pos, gb_state.pos, gb_state.quat ) penetration = normal.dot(pos_corner - ga_state.pos) if penetration > 0.0: contact_pos = pos_corner - 0.5 * penetration * normal if (contact_pos - contact_pos_0).norm() > tolerance: - self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal, + contact_pos, + penetration, + i_b, + ) n_con = n_con + 1 @ti.func - def _func_add_contact(self, i_ga, i_gb, normal, contact_pos, penetration, i_b): - i_c = self.n_contacts[i_b] - if i_c == self._max_contact_pairs: + def _func_add_contact( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + collider_state: ti.template(), + collider_info: ti.template(), + i_ga, + i_gb, + normal: ti.types.vector(3), + contact_pos: ti.types.vector(3), + penetration, + i_b, + ): + i_c = collider_state.n_contacts[i_b] + if i_c == collider_state._max_contact_pairs[None]: # FIXME: 'ti.static_print' cannot be used as it will be printed systematically, completely ignoring guard # condition, while 'print' is slowing down the kernel even if every called in practice... # print(self._warn_msg_max_collision_pairs) pass else: - ga_info = self._solver.geoms_info[i_ga] - gb_info = self._solver.geoms_info[i_gb] + ga_info = geoms_info[i_ga] + gb_info = geoms_info[i_gb] - friction_a = ga_info.friction * self._solver.geoms_state[i_ga, i_b].friction_ratio - friction_b = gb_info.friction * self._solver.geoms_state[i_gb, i_b].friction_ratio + friction_a = ga_info.friction * geoms_state[i_ga, i_b].friction_ratio + friction_b = gb_info.friction * geoms_state[i_gb, i_b].friction_ratio # b to a - self.contact_data[i_c, i_b].geom_a = i_ga - self.contact_data[i_c, i_b].geom_b = i_gb - self.contact_data[i_c, i_b].normal = normal - self.contact_data[i_c, i_b].pos = contact_pos - self.contact_data[i_c, i_b].penetration = penetration - self.contact_data[i_c, i_b].friction = ti.max(ti.max(friction_a, friction_b), 1e-2) - self.contact_data[i_c, i_b].sol_params = 0.5 * (ga_info.sol_params + gb_info.sol_params) - self.contact_data[i_c, i_b].link_a = ga_info.link_idx - self.contact_data[i_c, i_b].link_b = gb_info.link_idx - - self.n_contacts[i_b] = i_c + 1 + collider_state.contact_data[i_c, i_b].geom_a = i_ga + collider_state.contact_data[i_c, i_b].geom_b = i_gb + collider_state.contact_data[i_c, i_b].normal = normal + collider_state.contact_data[i_c, i_b].pos = contact_pos + collider_state.contact_data[i_c, i_b].penetration = penetration + collider_state.contact_data[i_c, i_b].friction = ti.max(ti.max(friction_a, friction_b), 1e-2) + collider_state.contact_data[i_c, i_b].sol_params = 0.5 * (ga_info.sol_params + gb_info.sol_params) + collider_state.contact_data[i_c, i_b].link_a = ga_info.link_idx + collider_state.contact_data[i_c, i_b].link_b = gb_info.link_idx + + collider_state.n_contacts[i_b] = i_c + 1 @ti.func - def _func_compute_tolerance(self, i_ga, i_gb, i_b): + def _func_compute_tolerance( + self_unused, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + collider_info: ti.template(), + i_ga, + i_gb, + i_b, + ): # Note that the original world-aligned bounding box is used to computed the absolute tolerance from the # relative one. This way, it is a constant that does not depends on the orientation of the geometry, which # makes sense since the scale of the geometries is an intrinsic property and not something that is supposed # to change dynamically. - aabb_size_b = (self._solver.geoms_init_AABB[i_gb, 7] - self._solver.geoms_init_AABB[i_gb, 0]).norm() + aabb_size_b = (geoms_init_AABB[i_gb, 7] - geoms_init_AABB[i_gb, 0]).norm() aabb_size = aabb_size_b - if self._solver.geoms_info[i_ga].type != gs.GEOM_TYPE.PLANE: - aabb_size_a = (self._solver.geoms_init_AABB[i_ga, 7] - self._solver.geoms_init_AABB[i_ga, 0]).norm() + if geoms_info[i_ga].type != gs.GEOM_TYPE.PLANE: + aabb_size_a = (geoms_init_AABB[i_ga, 7] - geoms_init_AABB[i_ga, 0]).norm() aabb_size = ti.min(aabb_size_a, aabb_size_b) - return 0.5 * self._mc_tolerance * aabb_size + return 0.5 * collider_info.mc_tolerance * aabb_size @ti.func - def _func_contact_orthogonals(self, i_ga, i_gb, normal, i_b): + def _func_contact_orthogonals( + self_unused, + links_state: array_class.LinksState, + links_info: array_class.LinksInfo, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + static_rigid_sim_config: ti.template(), + i_ga, + i_gb, + normal: ti.types.vector(3), + i_b, + ): axis_0 = ti.Vector.zero(gs.ti_float, 3) axis_1 = ti.Vector.zero(gs.ti_float, 3) - if ti.static(self._solver._enable_mujoco_compatibility): + if ti.static(static_rigid_sim_config.enable_mujoco_compatibility): # Choose between world axes Y or Z to avoid colinearity issue if ti.abs(normal[1]) < 0.5: axis_0[1] = 1.0 @@ -1183,16 +1630,16 @@ def _func_contact_orthogonals(self, i_ga, i_gb, normal, i_b): # the contact point. Basically, the smallest one between the two, which can be approximated # by the volume of their respective bounding box. i_g = i_gb - if self._solver.geoms_info[i_ga].type != gs.GEOM_TYPE.PLANE: - size_ga = self._solver.geoms_init_AABB[i_ga, 7] + if geoms_info[i_ga].type != gs.GEOM_TYPE.PLANE: + size_ga = geoms_init_AABB[i_ga, 7] volume_ga = size_ga[0] * size_ga[1] * size_ga[2] - size_gb = self._solver.geoms_init_AABB[i_gb, 7] + size_gb = geoms_init_AABB[i_gb, 7] volume_gb = size_gb[0] * size_gb[1] * size_gb[2] i_g = i_ga if volume_ga < volume_gb else i_gb # Compute orthogonal basis mixing principal inertia axes of geometry with contact normal - i_l = self._solver.geoms_info[i_g].link_idx - rot = gu.ti_quat_to_R(self._solver.links_state[i_l, i_b].i_quat) + i_l = geoms_info[i_g].link_idx + rot = gu.ti_quat_to_R(links_state[i_l, i_b].i_quat) axis_idx = gs.ti_int(0) axis_angle_max = gs.ti_float(0.0) for i in ti.static(range(3)): @@ -1208,31 +1655,58 @@ def _func_contact_orthogonals(self, i_ga, i_gb, normal, i_b): return axis_0, axis_1 @ti.func - def _func_convex_convex_contact(self, i_ga, i_gb, i_b): - if ( - self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE - and self._solver.geoms_info[i_gb].type == gs.GEOM_TYPE.BOX - ): + def _func_convex_convex_contact( + self_unused, + links_state: array_class.LinksState, + links_info: array_class.LinksInfo, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + geoms_init_AABB: array_class.GeomsInitAABB, + verts_info: array_class.VertsInfo, + rigid_global_info: ti.template(), + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), + collider_info: ti.template(), + mpr: ti.template(), + gjk: ti.template(), + sdf: ti.template(), + i_ga, + i_gb, + i_b, + ): + if geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE and geoms_info[i_gb].type == gs.GEOM_TYPE.BOX: if ti.static(sys.platform == "darwin"): - self._func_plane_box_contact(i_ga, i_gb, i_b) + self_unused._func_plane_box_contact( + geoms_state, + geoms_info, + geoms_init_AABB, + verts_info, + static_rigid_sim_config, + collider_state, + collider_info, + mpr, + i_ga, + i_gb, + i_b, + ) else: # Disabling multi-contact for pairs of decomposed geoms would speed up simulation but may cause physical # instabilities in the few cases where multiple contact points are actually need. Increasing the tolerance # criteria to get rid of redundant contact points seems to be a better option. multi_contact = ( - self._solver._enable_multi_contact + static_rigid_sim_config.enable_multi_contact # and not (self._solver.geoms_info[i_ga].is_decomposed and self._solver.geoms_info[i_gb].is_decomposed) - and self._solver.geoms_info[i_ga].type != gs.GEOM_TYPE.SPHERE - and self._solver.geoms_info[i_ga].type != gs.GEOM_TYPE.ELLIPSOID - and self._solver.geoms_info[i_gb].type != gs.GEOM_TYPE.SPHERE - and self._solver.geoms_info[i_gb].type != gs.GEOM_TYPE.ELLIPSOID + and geoms_info[i_ga].type != gs.GEOM_TYPE.SPHERE + and geoms_info[i_ga].type != gs.GEOM_TYPE.ELLIPSOID + and geoms_info[i_gb].type != gs.GEOM_TYPE.SPHERE + and geoms_info[i_gb].type != gs.GEOM_TYPE.ELLIPSOID ) - tolerance = self._func_compute_tolerance(i_ga, i_gb, i_b) + tolerance = self_unused._func_compute_tolerance(geoms_info, geoms_init_AABB, collider_info, i_ga, i_gb, i_b) # Backup state before local perturbation - ga_state = self._solver.geoms_state[i_ga, i_b] - gb_state = self._solver.geoms_state[i_gb, i_b] + ga_state = geoms_state[i_ga, i_b] + gb_state = geoms_state[i_gb, i_b] ga_pos, ga_quat = ga_state.pos, ga_state.quat gb_pos, gb_quat = gb_state.pos, gb_state.quat @@ -1257,29 +1731,33 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): # Perturbation axis must not be aligned with the principal axes of inertia the geometry, # otherwise it would be more sensitive to ill-conditionning. axis = (2 * (i_detection % 2) - 1) * axis_0 + (1 - 2 * ((i_detection // 2) % 2)) * axis_1 - qrot = gu.ti_rotvec_to_quat(self._mc_perturbation * axis) - self._func_rotate_frame(i_ga, contact_pos_0, qrot, i_b) - self._func_rotate_frame(i_gb, contact_pos_0, gu.ti_inv_quat(qrot), i_b) + qrot = gu.ti_rotvec_to_quat(collider_info.mc_perturbation * axis) + self_unused._func_rotate_frame(geoms_state, geoms_info, i_ga, contact_pos_0, qrot, i_b) + self_unused._func_rotate_frame( + geoms_state, geoms_info, i_gb, contact_pos_0, gu.ti_inv_quat(qrot), i_b + ) if (multi_contact and is_col_0) or (i_detection == 0): try_sdf = False - if self._solver.geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE: - ga_info = self._solver.geoms_info[i_ga] + if geoms_info[i_ga].type == gs.GEOM_TYPE.PLANE: + ga_info = geoms_info[i_ga] plane_dir = ti.Vector([ga_info.data[0], ga_info.data[1], ga_info.data[2]], dt=gs.ti_float) - plane_dir = gu.ti_transform_by_quat(plane_dir, self._solver.geoms_state[i_ga, i_b].quat) + plane_dir = gu.ti_transform_by_quat(plane_dir, geoms_state[i_ga, i_b].quat) normal = -plane_dir.normalized() - v1 = self._mpr.support_driver(normal, i_gb, i_b) - penetration = normal.dot(v1 - self._solver.geoms_state[i_ga, i_b].pos) + v1 = mpr.support_driver(normal, i_gb, i_b) + penetration = normal.dot(v1 - geoms_state[i_ga, i_b].pos) contact_pos = v1 - 0.5 * penetration * normal is_col = penetration > 0 else: ### MPR, MJ_MPR - if ti.static(self.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.MJ_MPR)): + if ti.static( + collider_info.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.MJ_MPR) + ): # Try using MPR before anything else is_mpr_updated = False is_mpr_guess_direction_available = True - normal_ws = self.contact_cache[i_ga, i_gb, i_b].normal + normal_ws = collider_state.contact_cache[i_ga, i_gb, i_b].normal for i_mpr in range(2): if i_mpr == 1: # Try without warm-start if no contact was detected using it. @@ -1292,7 +1770,7 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): is_mpr_updated = False if not is_mpr_updated: - is_col, normal, penetration, contact_pos = self._mpr.func_mpr_contact( + is_col, normal, penetration, contact_pos = mpr.func_mpr_contact( i_ga, i_gb, i_b, normal_ws ) is_mpr_updated = True @@ -1311,30 +1789,43 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): try_sdf = True ### GJK, MJ_GJK - elif ti.static(self.ccd_algorithm in (CCD_ALGORITHM_CODE.GJK, CCD_ALGORITHM_CODE.MJ_GJK)): - self._gjk.func_gjk_contact(i_ga, i_gb, i_b) + elif ti.static( + collider_info.ccd_algorithm in (CCD_ALGORITHM_CODE.GJK, CCD_ALGORITHM_CODE.MJ_GJK) + ): + gjk.func_gjk_contact(i_ga, i_gb, i_b) - is_col = self._gjk.is_col[i_b] == 1 - penetration = self._gjk.penetration[i_b] - n_contacts = self._gjk.n_contacts[i_b] + is_col = gjk.is_col[i_b] == 1 + penetration = gjk.penetration[i_b] + n_contacts = gjk.n_contacts[i_b] if is_col: - if self._gjk.multi_contact_flag[i_b]: + if gjk.multi_contact_flag[i_b]: # Used MuJoCo's multi-contact algorithm to find multiple contact points. Therefore, # add the discovered contact points and stop multi-contact search. for i_c in range(n_contacts): # Ignore contact points if the number of contacts exceeds the limit. - if i_c < ti.static(self._n_contacts_per_pair): - contact_pos = self._gjk.contact_pos[i_b, i_c] - normal = self._gjk.normal[i_b, i_c] - self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b) + if i_c < ti.static(collider_info.n_contacts_per_pair): + contact_pos = gjk.contact_pos[i_b, i_c] + normal = gjk.normal[i_b, i_c] + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal, + contact_pos, + penetration, + i_b, + ) break else: - contact_pos = self._gjk.contact_pos[i_b, 0] - normal = self._gjk.normal[i_b, 0] + contact_pos = gjk.contact_pos[i_b, 0] + normal = gjk.normal[i_b, 0] - if ti.static(self.ccd_algorithm == CCD_ALGORITHM_CODE.MPR): + if ti.static(collider_info.ccd_algorithm == CCD_ALGORITHM_CODE.MPR): if try_sdf: is_col_a = False is_col_b = False @@ -1357,8 +1848,14 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): # ) # ) # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = i_va - is_col_i, normal_i, penetration_i, contact_pos_i = self._func_contact_vertex_sdf( - i_ga if i_sdf == 0 else i_gb, i_gb if i_sdf == 0 else i_ga, i_b + is_col_i, normal_i, penetration_i, contact_pos_i = self_unused._func_contact_vertex_sdf( + geoms_state, + geoms_info, + verts_info, + sdf, + i_ga if i_sdf == 0 else i_gb, + i_gb if i_sdf == 0 else i_ga, + i_b, ) if i_sdf == 0: is_col_a = is_col_i @@ -1374,7 +1871,10 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): # MPR cannot handle collision detection for fully enclosed geometries. Falling back to SDF. # Note that SDF does not take into account to direction of interest. As such, it cannot be # used reliably for anything else than the point of deepest penetration. - prefer_sdf = self._mc_tolerance * penetration >= self._mpr_to_sdf_overlap_ratio * tolerance + prefer_sdf = ( + collider_info.mc_tolerance * penetration + >= collider_info.mpr_to_sdf_overlap_ratio * tolerance + ) if is_col_a and ( not is_col_b or penetration_a >= max(penetration_b, (not prefer_sdf) * penetration) @@ -1394,21 +1894,43 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): if i_detection == 0: is_col_0, normal_0, penetration_0, contact_pos_0 = is_col, normal, penetration, contact_pos if is_col_0: - self._func_add_contact(i_ga, i_gb, normal_0, contact_pos_0, penetration_0, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal_0, + contact_pos_0, + penetration_0, + i_b, + ) if multi_contact: # perturb geom_a around two orthogonal axes to find multiple contacts - axis_0, axis_1 = self._func_contact_orthogonals(i_ga, i_gb, normal, i_b) + axis_0, axis_1 = self_unused._func_contact_orthogonals( + links_state, + links_info, + geoms_state, + geoms_info, + geoms_init_AABB, + static_rigid_sim_config, + i_ga, + i_gb, + normal, + i_b, + ) n_con = 1 - if ti.static(self.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.GJK)): - self.contact_cache[i_ga, i_gb, i_b].normal = normal + if ti.static(collider_info.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.GJK)): + collider_state.contact_cache[i_ga, i_gb, i_b].normal = normal else: # Clear collision normal cache if not in contact # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1 - self.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) + collider_state.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0) elif multi_contact and is_col_0 > 0 and is_col > 0: - if ti.static(self.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.GJK)): + if ti.static(collider_info.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.GJK)): # 1. Project the contact point on both geometries # 2. Revert the effect of small rotation # 3. Update contact point @@ -1435,7 +1957,7 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): # original one, up to the scale of the perturbation, then apply first-order Taylor expension of # Rodrigues' rotation formula. twist_rotvec = ti.math.clamp( - normal.cross(normal_0), -self._mc_perturbation, self._mc_perturbation + normal.cross(normal_0), -collider_info.mc_perturbation, collider_info.mc_perturbation ) normal += twist_rotvec.cross(normal) @@ -1445,7 +1967,7 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): # dynamics since zero-penetration contact points should not induce any force. penetration = normal.dot(contact_point_b - contact_point_a) - elif ti.static(self.ccd_algorithm == CCD_ALGORITHM_CODE.MJ_GJK): + elif ti.static(collider_info.ccd_algorithm == CCD_ALGORITHM_CODE.MJ_GJK): # Only change penetration to the initial one, because the normal vector could change abruptly # under MuJoCo's GJK-EPA. penetration = penetration_0 @@ -1454,35 +1976,61 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b): repeated = False for i_c in range(n_con): if not repeated: - idx_prev = self.n_contacts[i_b] - 1 - i_c - prev_contact = self.contact_data[idx_prev, i_b].pos + idx_prev = collider_state.n_contacts[i_b] - 1 - i_c + prev_contact = collider_state.contact_data[idx_prev, i_b].pos if (contact_pos - prev_contact).norm() < tolerance: repeated = True if not repeated: if penetration > -tolerance: penetration = ti.max(penetration, 0.0) - self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + normal, + contact_pos, + penetration, + i_b, + ) n_con = n_con + 1 - self._solver.geoms_state[i_ga, i_b].pos = ga_pos - self._solver.geoms_state[i_ga, i_b].quat = ga_quat - self._solver.geoms_state[i_gb, i_b].pos = gb_pos - self._solver.geoms_state[i_gb, i_b].quat = gb_quat + geoms_state[i_ga, i_b].pos = ga_pos + geoms_state[i_ga, i_b].quat = ga_quat + geoms_state[i_gb, i_b].pos = gb_pos + geoms_state[i_gb, i_b].quat = gb_quat @ti.func - def _func_rotate_frame(self, i_g, contact_pos, qrot, i_b): - self._solver.geoms_state[i_g, i_b].quat = gu.ti_transform_quat_by_quat( - self._solver.geoms_state[i_g, i_b].quat, qrot - ) + def _func_rotate_frame( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + i_g, + contact_pos: ti.types.vector(3), + qrot: ti.types.vector(4), + i_b, + ): + geoms_state[i_g, i_b].quat = gu.ti_transform_quat_by_quat(geoms_state[i_g, i_b].quat, qrot) - rel = contact_pos - self._solver.geoms_state[i_g, i_b].pos + rel = contact_pos - geoms_state[i_g, i_b].pos vec = gu.ti_transform_by_quat(rel, qrot) vec = vec - rel - self._solver.geoms_state[i_g, i_b].pos = self._solver.geoms_state[i_g, i_b].pos - vec + geoms_state[i_g, i_b].pos = geoms_state[i_g, i_b].pos - vec @ti.func - def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): + def _func_box_box_contact( + self_unused, + geoms_state: array_class.GeomsState, + geoms_info: array_class.GeomsInfo, + collider_state: ti.template(), + collider_info: ti.template(), + i_ga, + i_gb, + i_b, + ): """ Use Mujoco's box-box contact detection algorithm for more stable collision detection. @@ -1502,10 +2050,10 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): margin2 = margin * margin rotmore = ti.Matrix.zero(gs.ti_float, 3, 3) - ga_info = self._solver.geoms_info[i_ga] - gb_info = self._solver.geoms_info[i_gb] - ga_state = self._solver.geoms_state[i_ga, i_b] - gb_state = self._solver.geoms_state[i_gb, i_b] + ga_info = geoms_info[i_ga] + gb_info = geoms_info[i_gb] + ga_state = geoms_state[i_ga, i_b] + gb_state = geoms_state[i_gb, i_b] size1 = ti.Vector([ga_info.data[0], ga_info.data[1], ga_info.data[2]], dt=gs.ti_float) / 2 size2 = ti.Vector([gb_info.data[0], gb_info.data[1], gb_info.data[2]], dt=gs.ti_float) / 2 @@ -1693,42 +2241,42 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): lp = lp + rt[i, :] * s[i] * (1 if (clcorner & (1 << i)) else -1) m, k = 0, 0 - self.box_pts[m, i_b] = lp + collider_state.box_pts[m, i_b] = lp m = m + 1 for i in ti.static(range(3)): if ti.abs(r[2, i]) < 0.5: - self.box_pts[m, i_b] = rt[i, :] * s[i] * (-2 if (clcorner & (1 << i)) else 2) + collider_state.box_pts[m, i_b] = rt[i, :] * s[i] * (-2 if (clcorner & (1 << i)) else 2) m = m + 1 - self.box_pts[3, i_b] = self.box_pts[0, i_b] + self.box_pts[1, i_b] - self.box_pts[4, i_b] = self.box_pts[0, i_b] + self.box_pts[2, i_b] - self.box_pts[5, i_b] = self.box_pts[3, i_b] + self.box_pts[2, i_b] + collider_state.box_pts[3, i_b] = collider_state.box_pts[0, i_b] + collider_state.box_pts[1, i_b] + collider_state.box_pts[4, i_b] = collider_state.box_pts[0, i_b] + collider_state.box_pts[2, i_b] + collider_state.box_pts[5, i_b] = collider_state.box_pts[3, i_b] + collider_state.box_pts[2, i_b] if m > 1: - self.box_lines[k, i_b][0:3] = self.box_pts[0, i_b] - self.box_lines[k, i_b][3:6] = self.box_pts[1, i_b] + collider_state.box_lines[k, i_b][0:3] = collider_state.box_pts[0, i_b] + collider_state.box_lines[k, i_b][3:6] = collider_state.box_pts[1, i_b] k = k + 1 if m > 2: - self.box_lines[k, i_b][0:3] = self.box_pts[0, i_b] - self.box_lines[k, i_b][3:6] = self.box_pts[2, i_b] + collider_state.box_lines[k, i_b][0:3] = collider_state.box_pts[0, i_b] + collider_state.box_lines[k, i_b][3:6] = collider_state.box_pts[2, i_b] k = k + 1 - self.box_lines[k, i_b][0:3] = self.box_pts[3, i_b] - self.box_lines[k, i_b][3:6] = self.box_pts[2, i_b] + collider_state.box_lines[k, i_b][0:3] = collider_state.box_pts[3, i_b] + collider_state.box_lines[k, i_b][3:6] = collider_state.box_pts[2, i_b] k = k + 1 - self.box_lines[k, i_b][0:3] = self.box_pts[4, i_b] - self.box_lines[k, i_b][3:6] = self.box_pts[1, i_b] + collider_state.box_lines[k, i_b][0:3] = collider_state.box_pts[4, i_b] + collider_state.box_lines[k, i_b][3:6] = collider_state.box_pts[1, i_b] k = k + 1 for i in range(k): for q in ti.static(range(2)): - a = self.box_lines[i, i_b][0 + q] - b = self.box_lines[i, i_b][3 + q] - c = self.box_lines[i, i_b][1 - q] - d = self.box_lines[i, i_b][4 - q] + a = collider_state.box_lines[i, i_b][0 + q] + b = collider_state.box_lines[i, i_b][3 + q] + c = collider_state.box_lines[i, i_b][1 - q] + d = collider_state.box_lines[i, i_b][4 - q] if ti.abs(b) > gs.EPS: for _j in ti.static(range(2)): j = 2 * _j - 1 @@ -1737,14 +2285,15 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): if 0 <= c1 and c1 <= 1: c2 = c + d * c1 if ti.abs(c2) <= ss[1 - q]: - self.box_points[n, i_b] = ( - self.box_lines[i, i_b][0:3] + self.box_lines[i, i_b][3:6] * c1 + collider_state.box_points[n, i_b] = ( + collider_state.box_lines[i, i_b][0:3] + + collider_state.box_lines[i, i_b][3:6] * c1 ) n = n + 1 - a = self.box_pts[1, i_b][0] - b = self.box_pts[2, i_b][0] - c = self.box_pts[1, i_b][1] - d = self.box_pts[2, i_b][1] + a = collider_state.box_pts[1, i_b][0] + b = collider_state.box_pts[2, i_b][0] + c = collider_state.box_pts[1, i_b][1] + d = collider_state.box_pts[2, i_b][1] c1 = a * d - b * c if m > 2: @@ -1752,35 +2301,37 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): llx = lx if (i // 2) else -lx lly = ly if (i % 2) else -ly - x = llx - self.box_pts[0, i_b][0] - y = lly - self.box_pts[0, i_b][1] + x = llx - collider_state.box_pts[0, i_b][0] + y = lly - collider_state.box_pts[0, i_b][1] u = (x * d - y * b) / c1 v = (y * a - x * c) / c1 if 0 < u and u < 1 and 0 < v and v < 1: - self.box_points[n, i_b] = ti.Vector( + collider_state.box_points[n, i_b] = ti.Vector( [ llx, lly, - self.box_pts[0, i_b][2] + u * self.box_pts[1, i_b][2] + v * self.box_pts[2, i_b][2], + collider_state.box_pts[0, i_b][2] + + u * collider_state.box_pts[1, i_b][2] + + v * collider_state.box_pts[2, i_b][2], ] ) n = n + 1 for i in range(1 << (m - 1)): - tmp1 = self.box_pts[0 if i == 0 else i + 2, i_b] + tmp1 = collider_state.box_pts[0 if i == 0 else i + 2, i_b] if not (i and (tmp1[0] <= -lx or tmp1[0] >= lx or tmp1[1] <= -ly or tmp1[1] >= ly)): - self.box_points[n, i_b] = tmp1 + collider_state.box_points[n, i_b] = tmp1 n = n + 1 m = n n = 0 for i in range(m): - if self.box_points[i, i_b][2] <= margin: - self.box_points[n, i_b] = self.box_points[i, i_b] - self.box_depth[n, i_b] = self.box_points[n, i_b][2] - self.box_points[n, i_b][2] = self.box_points[n, i_b][2] * 0.5 + if collider_state.box_points[i, i_b][2] <= margin: + collider_state.box_points[n, i_b] = collider_state.box_points[i, i_b] + collider_state.box_depth[n, i_b] = collider_state.box_points[n, i_b][2] + collider_state.box_points[n, i_b][2] = collider_state.box_points[n, i_b][2] * 0.5 n = n + 1 r = (mat2 if q2 else mat1) @ rotmore.transpose() p = pos2 if q2 else pos1 @@ -1790,11 +2341,22 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): ) normal_0 = tmp2 for i in range(n): - dist = self.box_points[i, i_b][2] - self.box_points[i, i_b][2] = self.box_points[i, i_b][2] + hz - tmp2 = r @ self.box_points[i, i_b] + dist = collider_state.box_points[i, i_b][2] + collider_state.box_points[i, i_b][2] = collider_state.box_points[i, i_b][2] + hz + tmp2 = r @ collider_state.box_points[i, i_b] contact_pos = tmp2 + p - self._func_add_contact(i_ga, i_gb, -normal_0, contact_pos, -dist, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + -normal_0, + contact_pos, + -dist, + i_b, + ) else: code = code - 12 @@ -1874,40 +2436,40 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): p[2] = p[2] - hz n = 0 - self.box_points[n, i_b] = p + collider_state.box_points[n, i_b] = p - self.box_points[n, i_b] = self.box_points[n, i_b] + rt[ax1, :] * size2[ax1] * ( + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] + rt[ax1, :] * size2[ax1] * ( 1 if (cle2 & (1 << ax1)) else -1 ) - self.box_points[n, i_b] = self.box_points[n, i_b] + rt[ax2, :] * size2[ax2] * ( + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] + rt[ax2, :] * size2[ax2] * ( 1 if (cle2 & (1 << ax2)) else -1 ) - self.box_points[n + 1, i_b] = self.box_points[n, i_b] - self.box_points[n, i_b] = self.box_points[n, i_b] + rt[q2, :] * size2[q2] + collider_state.box_points[n + 1, i_b] = collider_state.box_points[n, i_b] + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] + rt[q2, :] * size2[q2] n = 1 - self.box_points[n, i_b] = self.box_points[n, i_b] - rt[q2, :] * size2[q2] + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] - rt[q2, :] * size2[q2] n = 2 - self.box_points[n, i_b] = p - self.box_points[n, i_b] = self.box_points[n, i_b] + rt[ax1, :] * size2[ax1] * ( + collider_state.box_points[n, i_b] = p + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] + rt[ax1, :] * size2[ax1] * ( -1 if (cle2 & (1 << ax1)) else 1 ) - self.box_points[n, i_b] = self.box_points[n, i_b] + rt[ax2, :] * size2[ax2] * ( + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] + rt[ax2, :] * size2[ax2] * ( 1 if (cle2 & (1 << ax2)) else -1 ) - self.box_points[n + 1, i_b] = self.box_points[n, i_b] - self.box_points[n, i_b] = self.box_points[n, i_b] + rt[q2, :] * size2[q2] + collider_state.box_points[n + 1, i_b] = collider_state.box_points[n, i_b] + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] + rt[q2, :] * size2[q2] n = 3 - self.box_points[n, i_b] = self.box_points[n, i_b] - rt[q2, :] * size2[q2] + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] - rt[q2, :] * size2[q2] n = 4 - self.box_axi[0, i_b] = self.box_points[0, i_b] - self.box_axi[1, i_b] = self.box_points[1, i_b] - self.box_points[0, i_b] - self.box_axi[2, i_b] = self.box_points[2, i_b] - self.box_points[0, i_b] + collider_state.box_axi[0, i_b] = collider_state.box_points[0, i_b] + collider_state.box_axi[1, i_b] = collider_state.box_points[1, i_b] - collider_state.box_points[0, i_b] + collider_state.box_axi[2, i_b] = collider_state.box_points[2, i_b] - collider_state.box_points[0, i_b] if ti.abs(rnorm[2]) < gs.EPS: is_return = True @@ -1915,90 +2477,110 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): innorm = (1 / rnorm[2]) * (-1 if in_ else 1) for i in ti.static(range(4)): - c1 = -self.box_points[i, i_b][2] / rnorm[2] - self.box_pu[i, i_b] = self.box_points[i, i_b] - self.box_points[i, i_b] = self.box_points[i, i_b] + c1 * rnorm - - self.box_ppts2[i, 0, i_b] = self.box_points[i, i_b][0] - self.box_ppts2[i, 1, i_b] = self.box_points[i, i_b][1] - self.box_pts[0, i_b] = self.box_points[0, i_b] - self.box_pts[1, i_b] = self.box_points[1, i_b] - self.box_points[0, i_b] - self.box_pts[2, i_b] = self.box_points[2, i_b] - self.box_points[0, i_b] + c1 = -collider_state.box_points[i, i_b][2] / rnorm[2] + collider_state.box_pu[i, i_b] = collider_state.box_points[i, i_b] + collider_state.box_points[i, i_b] = collider_state.box_points[i, i_b] + c1 * rnorm + + collider_state.box_ppts2[i, 0, i_b] = collider_state.box_points[i, i_b][0] + collider_state.box_ppts2[i, 1, i_b] = collider_state.box_points[i, i_b][1] + collider_state.box_pts[0, i_b] = collider_state.box_points[0, i_b] + collider_state.box_pts[1, i_b] = ( + collider_state.box_points[1, i_b] - collider_state.box_points[0, i_b] + ) + collider_state.box_pts[2, i_b] = ( + collider_state.box_points[2, i_b] - collider_state.box_points[0, i_b] + ) m = 3 k = 0 n = 0 if m > 1: - self.box_lines[k, i_b][0:3] = self.box_pts[0, i_b] - self.box_lines[k, i_b][3:6] = self.box_pts[1, i_b] - self.box_linesu[k, i_b][0:3] = self.box_axi[0, i_b] - self.box_linesu[k, i_b][3:6] = self.box_axi[1, i_b] + collider_state.box_lines[k, i_b][0:3] = collider_state.box_pts[0, i_b] + collider_state.box_lines[k, i_b][3:6] = collider_state.box_pts[1, i_b] + collider_state.box_linesu[k, i_b][0:3] = collider_state.box_axi[0, i_b] + collider_state.box_linesu[k, i_b][3:6] = collider_state.box_axi[1, i_b] k = k + 1 if m > 2: - self.box_lines[k, i_b][0:3] = self.box_pts[0, i_b] - self.box_lines[k, i_b][3:6] = self.box_pts[2, i_b] - self.box_linesu[k, i_b][0:3] = self.box_axi[0, i_b] - self.box_linesu[k, i_b][3:6] = self.box_axi[2, i_b] + collider_state.box_lines[k, i_b][0:3] = collider_state.box_pts[0, i_b] + collider_state.box_lines[k, i_b][3:6] = collider_state.box_pts[2, i_b] + collider_state.box_linesu[k, i_b][0:3] = collider_state.box_axi[0, i_b] + collider_state.box_linesu[k, i_b][3:6] = collider_state.box_axi[2, i_b] k = k + 1 - self.box_lines[k, i_b][0:3] = self.box_pts[0, i_b] + self.box_pts[1, i_b] - self.box_lines[k, i_b][3:6] = self.box_pts[2, i_b] - self.box_linesu[k, i_b][0:3] = self.box_axi[0, i_b] + self.box_axi[1, i_b] - self.box_linesu[k, i_b][3:6] = self.box_axi[2, i_b] + collider_state.box_lines[k, i_b][0:3] = ( + collider_state.box_pts[0, i_b] + collider_state.box_pts[1, i_b] + ) + collider_state.box_lines[k, i_b][3:6] = collider_state.box_pts[2, i_b] + collider_state.box_linesu[k, i_b][0:3] = ( + collider_state.box_axi[0, i_b] + collider_state.box_axi[1, i_b] + ) + collider_state.box_linesu[k, i_b][3:6] = collider_state.box_axi[2, i_b] k = k + 1 - self.box_lines[k, i_b][0:3] = self.box_pts[0, i_b] + self.box_pts[2, i_b] - self.box_lines[k, i_b][3:6] = self.box_pts[1, i_b] - self.box_linesu[k, i_b][0:3] = self.box_axi[0, i_b] + self.box_axi[2, i_b] - self.box_linesu[k, i_b][3:6] = self.box_axi[1, i_b] + collider_state.box_lines[k, i_b][0:3] = ( + collider_state.box_pts[0, i_b] + collider_state.box_pts[2, i_b] + ) + collider_state.box_lines[k, i_b][3:6] = collider_state.box_pts[1, i_b] + collider_state.box_linesu[k, i_b][0:3] = ( + collider_state.box_axi[0, i_b] + collider_state.box_axi[2, i_b] + ) + collider_state.box_linesu[k, i_b][3:6] = collider_state.box_axi[1, i_b] k = k + 1 for i in range(k): for q in ti.static(range(2)): - a = self.box_lines[i, i_b][q] - b = self.box_lines[i, i_b][q + 3] - c = self.box_lines[i, i_b][1 - q] - d = self.box_lines[i, i_b][4 - q] + a = collider_state.box_lines[i, i_b][q] + b = collider_state.box_lines[i, i_b][q + 3] + c = collider_state.box_lines[i, i_b][1 - q] + d = collider_state.box_lines[i, i_b][4 - q] if ti.abs(b) > gs.EPS: for _j in ti.static(range(2)): j = 2 * _j - 1 - if n < self.box_MAXCONPAIR: + if n < collider_info.box_MAXCONPAIR: l = s[q] * j c1 = (l - a) / b if 0 <= c1 and c1 <= 1: c2 = c + d * c1 if (ti.abs(c2) <= s[1 - q]) and ( - (self.box_linesu[i, i_b][2] + self.box_linesu[i, i_b][5] * c1) * innorm + ( + collider_state.box_linesu[i, i_b][2] + + collider_state.box_linesu[i, i_b][5] * c1 + ) + * innorm <= margin ): - self.box_points[n, i_b] = ( - self.box_linesu[i, i_b][0:3] * 0.5 - + c1 * 0.5 * self.box_linesu[i, i_b][3:6] + collider_state.box_points[n, i_b] = ( + collider_state.box_linesu[i, i_b][0:3] * 0.5 + + c1 * 0.5 * collider_state.box_linesu[i, i_b][3:6] + ) + collider_state.box_points[n, i_b][q] = ( + collider_state.box_points[n, i_b][q] + 0.5 * l ) - self.box_points[n, i_b][q] = self.box_points[n, i_b][q] + 0.5 * l - self.box_points[n, i_b][1 - q] = ( - self.box_points[n, i_b][1 - q] + 0.5 * c2 + collider_state.box_points[n, i_b][1 - q] = ( + collider_state.box_points[n, i_b][1 - q] + 0.5 * c2 + ) + collider_state.box_depth[n, i_b] = ( + collider_state.box_points[n, i_b][2] * innorm * 2 ) - self.box_depth[n, i_b] = self.box_points[n, i_b][2] * innorm * 2 n = n + 1 nl = n - a = self.box_pts[1, i_b][0] - b = self.box_pts[2, i_b][0] - c = self.box_pts[1, i_b][1] - d = self.box_pts[2, i_b][1] + a = collider_state.box_pts[1, i_b][0] + b = collider_state.box_pts[2, i_b][0] + c = collider_state.box_pts[1, i_b][1] + d = collider_state.box_pts[2, i_b][1] c1 = a * d - b * c for i in range(4): - if n < self.box_MAXCONPAIR: + if n < collider_info.box_MAXCONPAIR: llx = lx if (i // 2) else -lx lly = ly if (i % 2) else -ly - x = llx - self.box_pts[0, i_b][0] - y = lly - self.box_pts[0, i_b][1] + x = llx - collider_state.box_pts[0, i_b][0] + y = lly - collider_state.box_pts[0, i_b][1] u = (x * d - y * b) / c1 v = (y * a - x * c) / c1 @@ -2011,27 +2593,31 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): u = ti.math.clamp(u, 0, 1) v = ti.math.clamp(v, 0, 1) - tmp1 = self.box_pu[0, i_b] * (1 - u - v) + self.box_pu[1, i_b] * u + self.box_pu[2, i_b] * v - self.box_points[n, i_b][0] = llx - self.box_points[n, i_b][1] = lly - self.box_points[n, i_b][2] = 0 + tmp1 = ( + collider_state.box_pu[0, i_b] * (1 - u - v) + + collider_state.box_pu[1, i_b] * u + + collider_state.box_pu[2, i_b] * v + ) + collider_state.box_points[n, i_b][0] = llx + collider_state.box_points[n, i_b][1] = lly + collider_state.box_points[n, i_b][2] = 0 - tmp2 = self.box_points[n, i_b] - tmp1 + tmp2 = collider_state.box_points[n, i_b] - tmp1 c2 = tmp2.dot(tmp2) if not (tmp1[2] > 0 and c2 > margin2): - self.box_points[n, i_b] = self.box_points[n, i_b] + tmp1 - self.box_points[n, i_b] = self.box_points[n, i_b] * 0.5 + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] + tmp1 + collider_state.box_points[n, i_b] = collider_state.box_points[n, i_b] * 0.5 - self.box_depth[n, i_b] = ti.sqrt(c2) * (-1 if tmp1[2] < 0 else 1) + collider_state.box_depth[n, i_b] = ti.sqrt(c2) * (-1 if tmp1[2] < 0 else 1) n = n + 1 nf = n for i in range(4): - if n < self.box_MAXCONPAIR: - x, y = self.box_ppts2[i, 0, i_b], self.box_ppts2[i, 1, i_b] + if n < collider_info.box_MAXCONPAIR: + x, y = collider_state.box_ppts2[i, 0, i_b], collider_state.box_ppts2[i, 1, i_b] if nl == 0: if (nf != 0) and (x < -lx or x > lx) and (y < -ly or y > ly): @@ -2041,30 +2627,37 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): c1 = 0 for j in ti.static(range(2)): - if self.box_ppts2[i, j, i_b] < -s[j]: - c1 = c1 + (self.box_ppts2[i, j, i_b] + s[j]) ** 2 - elif self.box_ppts2[i, j, i_b] > s[j]: - c1 = c1 + (self.box_ppts2[i, j, i_b] - s[j]) ** 2 + if collider_state.box_ppts2[i, j, i_b] < -s[j]: + c1 = c1 + (collider_state.box_ppts2[i, j, i_b] + s[j]) ** 2 + elif collider_state.box_ppts2[i, j, i_b] > s[j]: + c1 = c1 + (collider_state.box_ppts2[i, j, i_b] - s[j]) ** 2 - c1 = c1 + (self.box_pu[i, i_b][2] * innorm) ** 2 + c1 = c1 + (collider_state.box_pu[i, i_b][2] * innorm) ** 2 - if self.box_pu[i, i_b][2] > 0 and c1 > margin2: + if collider_state.box_pu[i, i_b][2] > 0 and c1 > margin2: continue tmp1 = ti.Vector( - [self.box_ppts2[i, 0, i_b] * 0.5, self.box_ppts2[i, 1, i_b] * 0.5, 0], dt=gs.ti_float + [ + collider_state.box_ppts2[i, 0, i_b] * 0.5, + collider_state.box_ppts2[i, 1, i_b] * 0.5, + 0, + ], + dt=gs.ti_float, ) for j in ti.static(range(2)): - if self.box_ppts2[i, j, i_b] < -s[j]: + if collider_state.box_ppts2[i, j, i_b] < -s[j]: tmp1[j] = -s[j] * 0.5 - elif self.box_ppts2[i, j, i_b] > s[j]: + elif collider_state.box_ppts2[i, j, i_b] > s[j]: tmp1[j] = s[j] * 0.5 - tmp1 = tmp1 + self.box_pu[i, i_b] * 0.5 - self.box_points[n, i_b] = tmp1 + tmp1 = tmp1 + collider_state.box_pu[i, i_b] * 0.5 + collider_state.box_points[n, i_b] = tmp1 - self.box_depth[n, i_b] = ti.sqrt(c1) * (-1 if self.box_pu[i, i_b][2] < 0 else 1) + collider_state.box_depth[n, i_b] = ti.sqrt(c1) * ( + -1 if collider_state.box_pu[i, i_b][2] < 0 else 1 + ) n = n + 1 r = mat1 @ rotmore.transpose() @@ -2073,11 +2666,22 @@ def _func_box_box_contact(self, i_ga: ti.i32, i_gb: ti.i32, i_b: ti.i32): normal_0 = tmp1 * (-1 if in_ else 1) for i in range(n): - dist = self.box_depth[i, i_b] - self.box_points[i, i_b][2] = self.box_points[i, i_b][2] + hz - tmp2 = r @ self.box_points[i, i_b] + dist = collider_state.box_depth[i, i_b] + collider_state.box_points[i, i_b][2] = collider_state.box_points[i, i_b][2] + hz + tmp2 = r @ collider_state.box_points[i, i_b] contact_pos = tmp2 + pos1 - self._func_add_contact(i_ga, i_gb, -normal_0, contact_pos, -dist, i_b) + self_unused._func_add_contact( + geoms_state, + geoms_info, + collider_state, + collider_info, + i_ga, + i_gb, + -normal_0, + contact_pos, + -dist, + i_b, + ) def get_contacts(self, as_tensor: bool = True, to_torch: bool = True): # Early return if already pre-computed @@ -2086,7 +2690,7 @@ def get_contacts(self, as_tensor: bool = True, to_torch: bool = True): return contacts_info.copy() # Find out how much dynamic memory must be allocated - n_contacts = tuple(self.n_contacts.to_numpy()) + n_contacts = tuple(self._collider_state.n_contacts.to_numpy()) n_envs = len(n_contacts) n_contacts_max = max(n_contacts) if as_tensor: @@ -2104,7 +2708,15 @@ def get_contacts(self, as_tensor: bool = True, to_torch: bool = True): # Copy contact data if n_contacts_max > 0: - self._kernel_get_contacts(as_tensor, iout, fout) + self._kernel_get_contacts( + self._solver._rigid_global_info, + self._solver._static_rigid_sim_config, + self._collider_state, + self._collider_info, + as_tensor, + iout, + fout, + ) # Build structured view (no copy) if as_tensor: @@ -2154,31 +2766,41 @@ def get_contacts(self, as_tensor: bool = True, to_torch: bool = True): return contacts_info.copy() @ti.kernel - def _kernel_get_contacts(self, is_padded: ti.template(), iout: ti.types.ndarray(), fout: ti.types.ndarray()): + def _kernel_get_contacts( + self_unused, + rigid_global_info: ti.template(), + static_rigid_sim_config: ti.template(), + collider_state: ti.template(), + collider_info: ti.template(), + is_padded: ti.template(), + iout: ti.types.ndarray(), + fout: ti.types.ndarray(), + ): + _B = collider_state.active_buffer.shape[1] n_contacts_max = gs.ti_int(0) - for i_b in range(self._solver._B): - n_contacts = self.n_contacts[i_b] + for i_b in range(_B): + n_contacts = collider_state.n_contacts[i_b] if n_contacts > n_contacts_max: n_contacts_max = n_contacts - ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL) - for i_b in range(self._solver._B): + ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL) + for i_b in range(_B): i_c_start = gs.ti_int(0) if ti.static(is_padded): i_c_start = i_b * n_contacts_max else: for j_b in range(i_b): - i_c_start = i_c_start + self.n_contacts[j_b] + i_c_start = i_c_start + collider_state.n_contacts[j_b] - for i_c_ in range(self.n_contacts[i_b]): + for i_c_ in range(collider_state.n_contacts[i_b]): i_c = i_c_start + i_c_ - iout[i_c, 0] = self.contact_data[i_c_, i_b].link_a - iout[i_c, 1] = self.contact_data[i_c_, i_b].link_b - iout[i_c, 2] = self.contact_data[i_c_, i_b].geom_a - iout[i_c, 3] = self.contact_data[i_c_, i_b].geom_b - fout[i_c, 0] = self.contact_data[i_c_, i_b].penetration + iout[i_c, 0] = collider_state.contact_data[i_c_, i_b].link_a + iout[i_c, 1] = collider_state.contact_data[i_c_, i_b].link_b + iout[i_c, 2] = collider_state.contact_data[i_c_, i_b].geom_a + iout[i_c, 3] = collider_state.contact_data[i_c_, i_b].geom_b + fout[i_c, 0] = collider_state.contact_data[i_c_, i_b].penetration for j in ti.static(range(3)): - fout[i_c, 1 + j] = self.contact_data[i_c_, i_b].pos[j] - fout[i_c, 4 + j] = self.contact_data[i_c_, i_b].normal[j] - fout[i_c, 7 + j] = self.contact_data[i_c_, i_b].force[j] + fout[i_c, 1 + j] = collider_state.contact_data[i_c_, i_b].pos[j] + fout[i_c, 4 + j] = collider_state.contact_data[i_c_, i_b].normal[j] + fout[i_c, 7 + j] = collider_state.contact_data[i_c_, i_b].force[j] diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp.py b/genesis/engine/solvers/rigid/constraint_solver_decomp.py index 1c273a771f..7c2d7b2772 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp.py @@ -27,7 +27,7 @@ def __init__(self, rigid_solver: "RigidSolver"): # 4 constraints per contact, 1 constraints per joint limit (upper and lower, if not inf), and 3 constraints per equality self.len_constraints = ( - 5 * self._collider._max_contact_pairs + 5 * rigid_solver.collider._collider_state._max_contact_pairs[None] + np.logical_not(np.isinf(self._solver.dofs_info.limit.to_numpy()[:, 0])).sum() + self._solver.n_equalities_candidate * 6 ) @@ -122,8 +122,8 @@ def _kernel_clear(self, envs_idx: ti.types.ndarray()): def add_collision_constraints(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): - for i_col in range(self._collider.n_contacts[i_b]): - contact_data = self._collider.contact_data[i_col, i_b] + for i_col in range(self._collider._collider_state.n_contacts[i_b]): + contact_data = self._collider._collider_state.contact_data[i_col, i_b] link_a = contact_data.link_a link_b = contact_data.link_b link_a_maybe_batch = [link_a, i_b] if ti.static(self._solver._options.batch_links_info) else link_a @@ -778,8 +778,8 @@ def _func_update_contact_force(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): const_start = self.n_constraints_equality[i_b] - for i_c in range(self._collider.n_contacts[i_b]): - contact_data = self._collider.contact_data[i_c, i_b] + for i_c in range(self._collider._collider_state.n_contacts[i_b]): + contact_data = self._collider._collider_state.contact_data[i_c, i_b] force = ti.Vector.zero(gs.ti_float, 3) d1, d2 = gu.ti_orthogonals(contact_data.normal) @@ -788,7 +788,7 @@ def _func_update_contact_force(self): n = d * contact_data.friction - contact_data.normal force += n * self.efc_force[i_c * 4 + i_dir + const_start, i_b] - self._collider.contact_data[i_c, i_b].force = force + self._collider._collider_state.contact_data[i_c, i_b].force = force self._solver.links_state[contact_data.link_a, i_b].contact_force = ( self._solver.links_state[contact_data.link_a, i_b].contact_force - force diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py index 45e61e4ec1..88503c36b9 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py @@ -24,7 +24,7 @@ def __init__(self, rigid_solver): # 4 constraints per contact and 1 constraints per joint limit (upper and lower, if not inf) self.len_constraints = ( - 5 * self._collider._max_contact_pairs + 5 * self._collider._collider_state._max_contact_pairs[None] + np.logical_not(np.isinf(self._solver.dofs_info.limit.to_numpy()[:, 0])).sum() ) self.len_constraints_ = max(1, self.len_constraints) @@ -133,7 +133,7 @@ def add_collision_constraints(self, island, i_b): i_col_ = self.contact_island.island_col[island, i_b].start + i_island_col i_col = self.contact_island.constraint_id[i_col_, i_b] - contact_data = self._collider.contact_data[i_col, i_b] + contact_data = self._collider._collider_state.contact_data[i_col, i_b] link_a = contact_data.link_a link_b = contact_data.link_b link_a_maybe_batch = [link_a, i_b] if ti.static(self._solver._options.batch_links_info) else link_a @@ -523,7 +523,7 @@ def _func_update_contact_force(self, island, i_b): i_col_ = self.contact_island.island_col[island, i_b].start + i_island_col i_col = self.contact_island.constraint_id[i_col_, i_b] - contact_data = self._collider.contact_data[i_col, i_b] + contact_data = self._collider._collider_state.contact_data[i_col, i_b] force = ti.Vector.zero(gs.ti_float, 3) d1, d2 = gu.ti_orthogonals(contact_data.normal) @@ -531,7 +531,7 @@ def _func_update_contact_force(self, island, i_b): d = (2 * (i % 2) - 1) * (d1 if i < 2 else d2) n = d * contact_data.friction - contact_data.normal force += n * self.efc_force[i_island_col * 4 + i, i_b] - self._collider.contact_data[i_col, i_b].force = force + self._collider._collider_state.contact_data[i_col, i_b].force = force self._solver.links_state[contact_data.link_a, i_b].contact_force = ( self._solver.links_state[contact_data.link_a, i_b].contact_force - force diff --git a/genesis/engine/solvers/rigid/contact_island.py b/genesis/engine/solvers/rigid/contact_island.py index bda0692078..39c5773a45 100644 --- a/genesis/engine/solvers/rigid/contact_island.py +++ b/genesis/engine/solvers/rigid/contact_island.py @@ -17,16 +17,22 @@ def __init__(self, collider): start=gs.ti_int, ) - self.ci_edges = ti.field(dtype=gs.ti_int, shape=self.solver._batch_shape((self.collider._max_contact_pairs, 2))) + self.ci_edges = ti.field( + dtype=gs.ti_int, shape=self.solver._batch_shape((self.collider._collider_state._max_contact_pairs[None], 2)) + ) - self.edge_id = ti.field(dtype=gs.ti_int, shape=self.solver._batch_shape((self.collider._max_contact_pairs * 2))) + self.edge_id = ti.field( + dtype=gs.ti_int, + shape=self.solver._batch_shape((self.collider._collider_state._max_contact_pairs[None] * 2)), + ) self.constraint_list = ti.field( - dtype=gs.ti_int, shape=self.solver._batch_shape((self.collider._max_contact_pairs)) + dtype=gs.ti_int, shape=self.solver._batch_shape((self.collider._collider_state._max_contact_pairs[None])) ) self.constraint_id = ti.field( - dtype=gs.ti_int, shape=self.solver._batch_shape((self.collider._max_contact_pairs * 2)) + dtype=gs.ti_int, + shape=self.solver._batch_shape((self.collider._collider_state._max_contact_pairs[None] * 2)), ) self.entity_edge = struct_agg_list.field( @@ -86,8 +92,8 @@ def add_edge(self, link_a, link_b, i_b): def add_island(self): ti.loop_config(serialize=self.solver._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self.solver._B): - for i_col in range(self.collider.n_contacts[i_b]): - impact = self.collider.contact_data[i_col, i_b] + for i_col in range(self.collider._collider_state.n_contacts[i_b]): + impact = self.collider._collider_state.contact_data[i_col, i_b] self.add_edge(impact.link_a, impact.link_b, i_b) def construct(self): @@ -101,8 +107,8 @@ def construct(self): def postprocess_island(self): ti.loop_config(serialize=self.solver._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self.solver._B): - for i_col in range(self.collider.n_contacts[i_b]): - impact = self.collider.contact_data[i_col, i_b] + for i_col in range(self.collider._collider_state.n_contacts[i_b]): + impact = self.collider._collider_state.contact_data[i_col, i_b] link_a = impact.link_a link_b = impact.link_b link_a_maybe_batch = [link_a, i_b] if ti.static(self.solver._options.batch_links_info) else link_a @@ -129,7 +135,7 @@ def postprocess_island(self): self.island_hibernated[i, i_b] = 1 - for i_col in range(self.collider.n_contacts[i_b]): + for i_col in range(self.collider._collider_state.n_contacts[i_b]): island = self.constraint_list[i_col, i_b] self.constraint_id[self.island_col[island, i_b].curr, i_b] = i_col self.island_col[island, i_b].curr = self.island_col[island, i_b].curr + 1 diff --git a/genesis/engine/solvers/rigid/gjk_decomp.py b/genesis/engine/solvers/rigid/gjk_decomp.py index c6215c2114..be8fb64a3d 100644 --- a/genesis/engine/solvers/rigid/gjk_decomp.py +++ b/genesis/engine/solvers/rigid/gjk_decomp.py @@ -2979,7 +2979,7 @@ def support_driver(self, direction, i_g, i_b, i_o, shrink_sphere): elif geom_type == gs.GEOM_TYPE.BOX: v, vid = self.support_field._func_support_box(direction, i_g, i_b) elif geom_type == gs.GEOM_TYPE.TERRAIN: - if ti.static(self._solver.collider._has_terrain): + if ti.static(self._solver.collider._collider_info.has_terrain): v, vid = self.support_field._func_support_prism(direction, i_g, i_b) elif geom_type == gs.GEOM_TYPE.MESH and self._enable_mujoco_compatibility: # If mujoco-compatible, do exhaustive search for the vertex diff --git a/genesis/engine/solvers/rigid/mpr_decomp.py b/genesis/engine/solvers/rigid/mpr_decomp.py index c83d20bf64..f7575efba5 100644 --- a/genesis/engine/solvers/rigid/mpr_decomp.py +++ b/genesis/engine/solvers/rigid/mpr_decomp.py @@ -179,7 +179,7 @@ def support_driver(self, direction, i_g, i_b): elif geom_type == gs.GEOM_TYPE.BOX: v, _ = self.support_field._func_support_box(direction, i_g, i_b) elif geom_type == gs.GEOM_TYPE.TERRAIN: - if ti.static(self._solver.collider._has_terrain): + if ti.static(self._solver.collider._collider_info.has_terrain): v, _ = self.support_field._func_support_prism(direction, i_g, i_b) else: v, _ = self.support_field._func_support_world(direction, i_g, i_b) diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index 00db730bb6..b8fa9f9117 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -72,6 +72,12 @@ class StaticRigidSimConfig: # store static arguments here para_level: int = 0 use_hibernation: bool = False + batch_links_info: bool = False + enable_mujoco_compatibility: bool = False + enable_multi_contact: bool = True + enable_self_collision: bool = True + enable_adjacent_collision: bool = False + box_box_detection: bool = False def __init__(self, scene: "Scene", sim: "Simulator", options: RigidOptions) -> None: super().__init__(scene, sim, options) @@ -223,6 +229,12 @@ def build(self): self._static_rigid_sim_config = self.StaticRigidSimConfig( para_level=self.sim._para_level, use_hibernation=getattr(self, "_use_hibernation", False), + batch_links_info=getattr(self._options, "batch_links_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), + box_box_detection=getattr(self, "_box_box_detection", False), ) # when the migration is finished, we will remove the about two lines # and initizlize the awake_dofs and n_awake_dofs in _rigid_global_info directly @@ -230,6 +242,7 @@ def build(self): n_dofs=self.n_dofs_, n_entities=self.n_entities_, n_geoms=self.n_geoms_, + _B=self._B, f_batch=self._batch_shape, ) @@ -1378,7 +1391,7 @@ def _init_sdf(self): def _init_collider(self): self.collider = Collider(self) - if self.collider._has_terrain: + if self.collider._collider_info.has_terrain: links_idx = self.geoms_info.link_idx.to_numpy()[self.geoms_info.type.to_numpy() == gs.GEOM_TYPE.TERRAIN] entity = self._entities[self.links_info.entity_idx.to_numpy()[links_idx[0]]] @@ -1819,10 +1832,10 @@ def _kernel_detect_collision(self): def detect_collision(self, env_idx=0): # TODO: support batching self._kernel_detect_collision() - n_collision = self.collider.n_contacts.to_numpy()[env_idx] + n_collision = self.collider._collider_state.n_contacts.to_numpy()[env_idx] collision_pairs = np.empty((n_collision, 2), dtype=np.int32) - collision_pairs[:, 0] = self.collider.contact_data.geom_a.to_numpy()[:n_collision, env_idx] - collision_pairs[:, 1] = self.collider.contact_data.geom_b.to_numpy()[:n_collision, env_idx] + collision_pairs[:, 0] = self.collider._collider_state.contact_data.geom_a.to_numpy()[:n_collision, env_idx] + collision_pairs[:, 1] = self.collider._collider_state.contact_data.geom_b.to_numpy()[:n_collision, env_idx] return collision_pairs @ti.kernel @@ -1883,7 +1896,7 @@ def _func_constraint_clear(self): self.n_contacts[i_b] = self.n_contacts_hibernated[i_b] else: - self.collider.n_contacts.fill(0) + self.collider._collider_state.n_contacts.fill(0) def _batch_array(self, arr, first_dim=False): if first_dim: @@ -2565,12 +2578,12 @@ def _func_hibernate(self): else: # update collider sort_buffer for i_g in range(self.entities_info[i_e].geom_start, self.entities_info[i_e].geom_end): - self.collider.sort_buffer[self.geoms_state[i_g, i_b].min_buffer_idx, i_b].value = ( - self.geoms_state[i_g, i_b].aabb_min[0] - ) - self.collider.sort_buffer[self.geoms_state[i_g, i_b].max_buffer_idx, i_b].value = ( - self.geoms_state[i_g, i_b].aabb_max[0] - ) + self.collider._collider_state.sort_buffer[ + self.geoms_state[i_g, i_b].min_buffer_idx, i_b + ].value = self.geoms_state[i_g, i_b].aabb_min[0] + self.collider._collider_state.sort_buffer[ + self.geoms_state[i_g, i_b].max_buffer_idx, i_b + ].value = self.geoms_state[i_g, i_b].aabb_max[0] @ti.func def _func_aggregate_awake_entities(self): diff --git a/genesis/engine/solvers/rigid/sdf.py b/genesis/engine/solvers/rigid/sdf.py index 4dbf3eafaf..0de6e81a55 100644 --- a/genesis/engine/solvers/rigid/sdf.py +++ b/genesis/engine/solvers/rigid/sdf.py @@ -104,7 +104,7 @@ def true_sdf_grad_sdf(self, pos_sdf, geom_idx): """ sdf_grad_sdf = ti.Vector.zero(gs.ti_float, 3) if self.solver.geoms_info[geom_idx].type == gs.GEOM_TYPE.TERRAIN: # Terrain uses finite difference - if ti.static(self.solver.collider._has_terrain): # for speed up compilation + if ti.static(self.solver.collider._collider_info._has_terrain): # for speed up compilation # since we are in sdf frame, delta can be a relatively big value delta = gs.ti_float(1e-2) diff --git a/genesis/engine/solvers/rigid/sdf_decomp.py b/genesis/engine/solvers/rigid/sdf_decomp.py index c55e78e75d..a80ca6d26e 100644 --- a/genesis/engine/solvers/rigid/sdf_decomp.py +++ b/genesis/engine/solvers/rigid/sdf_decomp.py @@ -104,7 +104,7 @@ def true_sdf_grad_sdf(self, pos_sdf, geom_idx): """ sdf_grad_sdf = ti.Vector.zero(gs.ti_float, 3) if self.solver.geoms_info[geom_idx].type == gs.GEOM_TYPE.TERRAIN: # Terrain uses finite difference - if ti.static(self.solver.collider._has_terrain): # for speed up compilation + if ti.static(self.solver.collider._collider_info.has_terrain): # for speed up compilation # since we are in sdf frame, delta can be a relatively big value delta = gs.ti_float(1e-2) diff --git a/genesis/engine/solvers/rigid/support_field_decomp.py b/genesis/engine/solvers/rigid/support_field_decomp.py index 0bb8c1ffeb..af20a1cfa5 100644 --- a/genesis/engine/solvers/rigid/support_field_decomp.py +++ b/genesis/engine/solvers/rigid/support_field_decomp.py @@ -202,14 +202,14 @@ def _func_support_prism(self, d, i_g, i_b): istart = 0 ibest = istart - best = self.solver.collider.prism[istart, i_b].dot(d) + best = self.solver.collider._collider_state.prism[istart, i_b].dot(d) for i in range(istart + 1, istart + 3): - dot = self.solver.collider.prism[i, i_b].dot(d) + dot = self.solver.collider._collider_state.prism[i, i_b].dot(d) if dot > best: ibest = i best = dot - return self.solver.collider.prism[ibest, i_b], ibest + return self.solver.collider._collider_state.prism[ibest, i_b], ibest @ti.func def _func_support_box(self, d, i_g, i_b): diff --git a/tests/test_rigid_physics.py b/tests/test_rigid_physics.py index 9945fd2440..a5128c771a 100644 --- a/tests/test_rigid_physics.py +++ b/tests/test_rigid_physics.py @@ -1668,8 +1668,8 @@ def test_path_planning_avoidance(show_viewer): # Check if the cube is colliding with the robot scene.rigid_solver._kernel_forward_dynamics() scene.rigid_solver._func_constraint_force() - for i in range(scene.rigid_solver.collider.n_contacts.to_numpy()[0]): - contact_data = scene.rigid_solver.collider.contact_data[i, 0] + for i in range(scene.rigid_solver.collider._collider_state.n_contacts.to_numpy()[0]): + contact_data = scene.rigid_solver.collider._collider_state.contact_data[i, 0] if any(i_g in tuple(range(len(cubes))) for i_g in (contact_data.link_a, contact_data.link_b)): max_penetration = max(max_penetration, contact_data.penetration) @@ -2028,7 +2028,7 @@ def test_convexify(euler, backend, show_viewer, gjk_collision): # Make sure that all the geometries in the scene are convex assert gs_sim.rigid_solver.geoms_info.is_convex.to_numpy().all() - assert not gs_sim.rigid_solver.collider._has_nonconvex_nonterrain + assert not gs_sim.rigid_solver.collider._collider_info.has_nonconvex_nonterrain # There should be only one geometry for the apple as it can be convexify without decomposition, # but for the others it is hard to tell... Let's use some reasonable guess. @@ -2563,7 +2563,7 @@ def test_drone_advanced(show_viewer): drone.set_propellels_rpm(torch.full((4,), 50000.0)) scene.step() if i > 350: - assert scene.rigid_solver.collider.n_contacts.to_numpy()[0] == 2 + assert scene.rigid_solver.collider._collider_state.n_contacts.to_numpy()[0] == 2 assert_allclose(scene.rigid_solver.get_dofs_velocity(), 0, tol=2e-3) # Push the drones symmetrically and wait for them to collide @@ -2573,7 +2573,7 @@ def test_drone_advanced(show_viewer): for drone in drones: drone.set_propellels_rpm(torch.full((4,), 50000.0)) scene.step() - if scene.rigid_solver.collider.n_contacts.to_numpy()[0] > 2: + if scene.rigid_solver.collider._collider_state.n_contacts.to_numpy()[0] > 2: break else: raise AssertionError @@ -2634,7 +2634,7 @@ def test_data_accessor(n_envs, batched, tol): for _ in range(400): gs_sim.step() - gs_n_contacts = gs_sim.rigid_solver.collider.n_contacts.to_numpy() + gs_n_contacts = gs_sim.rigid_solver.collider._collider_state.n_contacts.to_numpy() assert len(gs_n_contacts) == max(n_envs, 1) for as_tensor in (False, True): for to_torch in (False, True): diff --git a/tests/utils.py b/tests/utils.py index 24b951b5d5..5917727336 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -762,7 +762,7 @@ def check_mujoco_data_consistency( mj_qfrc_actuator = mj_sim.data.qfrc_actuator assert_allclose(gs_qfrc_actuator, mj_qfrc_actuator[mj_dofs_idx], tol=tol) - gs_n_contacts = gs_sim.rigid_solver.collider.n_contacts.to_numpy()[0] + gs_n_contacts = gs_sim.rigid_solver.collider._collider_state.n_contacts.to_numpy()[0] mj_n_contacts = mj_sim.data.ncon assert gs_n_contacts == mj_n_contacts gs_n_constraints = gs_sim.rigid_solver.constraint_solver.n_constraints.to_numpy()[0] @@ -770,7 +770,7 @@ def check_mujoco_data_consistency( assert gs_n_constraints == mj_n_constraints if gs_n_constraints: - gs_contact_pos = gs_sim.rigid_solver.collider.contact_data.pos.to_numpy()[:gs_n_contacts, 0] + gs_contact_pos = gs_sim.rigid_solver.collider._collider_state.contact_data.pos.to_numpy()[:gs_n_contacts, 0] mj_contact_pos = mj_sim.data.contact.pos # Sort based on the axis with the largest variation max_var_axis = 0 @@ -785,10 +785,14 @@ def check_mujoco_data_consistency( gs_sidx = np.argsort(gs_contact_pos[:, max_var_axis]) mj_sidx = np.argsort(mj_contact_pos[:, max_var_axis]) assert_allclose(gs_contact_pos[gs_sidx], mj_contact_pos[mj_sidx], tol=tol) - gs_contact_normal = gs_sim.rigid_solver.collider.contact_data.normal.to_numpy()[:gs_n_contacts, 0] + gs_contact_normal = gs_sim.rigid_solver.collider._collider_state.contact_data.normal.to_numpy()[ + :gs_n_contacts, 0 + ] mj_contact_normal = -mj_sim.data.contact.frame[:, :3] assert_allclose(gs_contact_normal[gs_sidx], mj_contact_normal[mj_sidx], tol=tol) - gs_penetration = gs_sim.rigid_solver.collider.contact_data.penetration.to_numpy()[:gs_n_contacts, 0] + gs_penetration = gs_sim.rigid_solver.collider._collider_state.contact_data.penetration.to_numpy()[ + :gs_n_contacts, 0 + ] mj_penetration = -mj_sim.data.contact.dist assert_allclose(gs_penetration[gs_sidx], mj_penetration[mj_sidx], tol=tol)