robolib.physics.collision

  1"""collision.py - Collision submodule between basic shapes."""
  2import numpy as np
  3from microecs import Entity, QueryResult
  4from robolib.utils import Point3D, logger, fmt
  5from robolib.components import ColliderKinds
  6
  7Key3D = tuple[int, int, int]
  8
  9def make_grid_3d(entities: list[Entity], cell_size: Point3D) -> dict[Key3D, list[Entity]]:
 10    """
 11    Creates a 3d grid given all the objects based on their bounding boxes. Cell size is provided from outside.
 12    Output dict is provided (not returned), so we can create it outside (e.g. prefill fixed ones)!
 13    """
 14    res: dict[Key3D, list[Entity]] = {}
 15    assert (cell_size > 0).all(), cell_size
 16    for entity in entities:
 17        top_left, bottom_right = entity.collider_bbox[0:2]
 18        first_cell = (top_left / cell_size).astype(int) # Note: possible bug with negative positions??
 19        last_cell = (bottom_right / cell_size).astype(int)
 20        nx, ny, nz = last_cell - first_cell + 1
 21
 22        for i in range(nx):
 23            for j in range(ny):
 24                for k in range(nz):
 25                    key = (first_cell[0].item() + i, first_cell[1].item() + j, first_cell[2].item() + k)
 26                    res.setdefault(key, []).append(entity)
 27    return res
 28
 29def make_collision_cell_size(entities: QueryResult) -> Point3D:
 30    """the collision cell size is twice them edian of all bboxes of all collidables with a model"""
 31    bboxes = entities.model_bbox * entities.scale[..., None]
 32    diffs = (bboxes[:, 1] - bboxes[:, 0]).numpy()
 33    median = np.median(diffs, axis=0) # median across all 3 dimensions
 34    res = median * 2 # heuristic: twice the median
 35    logger.info(f"Median bbox size: {median}. Cell size: {res}")
 36    assert (res > 0).all(), bboxes.numpy()
 37    return res
 38
 39def _get_closest_point_on_box(top_left: Point3D, bottom_right: Point3D, other_point: Point3D) -> Point3D:
 40    _f = lambda l, p, r: l if p < l else min(p, r) # returns one of 3 cases: l, p or r on 1D line: [pl, [l, p, r], pr]
 41    px = _f(top_left[0], other_point[0], bottom_right[0])
 42    py = _f(top_left[1], other_point[1], bottom_right[1])
 43    pz = _f(top_left[2], other_point[2], bottom_right[2])
 44    return np.float32([px, py, pz])
 45
 46def sphere_sphere_collision(pos1: Point3D, radius1: float, pos2: Point3D, radius2: float, eps: float=1e-5) -> bool:
 47    """checks if two spheres are colliding"""
 48    max_dist = radius1 + radius2
 49    diff_vec = pos1 - pos2
 50    dist: float = ((diff_vec**2).sum() - max_dist**2).item() # instead of np.linalg.norm(a.pos-b.pos) < max_dist
 51    if check := dist < -eps:
 52        logger.log_every_s(f"Collision. pos1={fmt(pos1)} radius1={radius1:.3f} "
 53                           f"pos2={fmt(pos2)} radius2={radius2:.3f}", "DEBUG", True)
 54    return check
 55
 56def sphere_aabb_collision(center: Point3D, radius: float, top_left: Point3D,
 57                          bottom_right: Point3D, eps: float=1e-5) -> bool:
 58    """checks if a sphere is colliding with an aabb"""
 59    xr, yr, zr = _get_closest_point_on_box(top_left, bottom_right, center)
 60    dist_sq = (center[0] - xr)**2 + (center[1] - yr)**2 + (center[2] - zr)**2
 61    dist: float = (dist_sq - radius ** 2).item()
 62    if check := dist < -eps:
 63        logger.log_every_s(f"Collision. center={fmt(center)} radius={radius:.3f} top_left={fmt(top_left)} "
 64                           f"bottom_right={fmt(bottom_right)}", "DEBUG", True)
 65    return check
 66
 67def aabb_aabb_collision(top_left1: Point3D, bottom_right1: Point3D,
 68                        top_left2: Point3D, bottom_right2: Point3D, eps: float=1e-5) -> bool:
 69    """
 70    Checks if two aabbs are colliding. We do this by doing one check on each ax.
 71    We are thus reducing the complexity from 3D to three independent 1D checks.
 72
 73    For 1D, assume the two lines given by top_left and bottom_right: [tl1, br1], [tl2, br2]. There are 5 cases:
 74    - tl1, br1, tl2, br2 -> not overlapping
 75    - tl1, tl2, br1, br2 -> overlapping
 76    - tl1, tl2, br2, br1 -> overlapping
 77    - tl2, tl1, br1, br2 -> overlapping
 78    - tl2, tl1, br2, br1 -> overlapping
 79    - tl2, br2, tl1, br1 -> not overlapping
 80
 81    In all the cases, we can look at the midpoints c1, c2 and the actual distances between the two corners:
 82    h1 = (br1 - tl1) / 2, h2 = (br2 - tl2) / 2. The lines are overlapping only if dist(c1, c2) < h1 + h2 across all dims
 83    """
 84
 85    h1 = (bottom_right1 - top_left1) / 2
 86    h2 = (bottom_right2 - top_left2) / 2
 87
 88    c1 = top_left1 + h1
 89    c2 = top_left2 + h2
 90
 91    # signed separation on each axis, like the sphere checks above: < 0 means overlapping on that axis
 92    gap = abs(c1 - c2) - (h1 + h2)
 93    if check := (gap[0] < -eps) and (gap[1] < -eps) and (gap[2] < -eps):
 94        logger.log_every_s(f"Collision. {top_left1=} {bottom_right1=} {top_left2=} {bottom_right2=}", "DEBUG", True)
 95    return check
 96
 97def check_collision(e1: Entity, e2: Entity) -> bool:
 98    """checks the collision between two entities with the HasCollision component"""
 99    if e1.collider_kind == ColliderKinds.SPHERE and e2.collider_kind == ColliderKinds.SPHERE:
100        return sphere_sphere_collision(e1.candidate_pose[0:3, 3], e1.collider_radii.item(),
101                                       e2.candidate_pose[0:3, 3], e2.collider_radii.item())
102    elif e1.collider_kind == ColliderKinds.SPHERE and e2.collider_kind == ColliderKinds.AABB:
103        return sphere_aabb_collision(e1.candidate_pose[0:3, 3], e1.collider_radii.item(),
104                                     e2.collider_bbox[0], e2.collider_bbox[1])
105    elif e1.collider_kind == ColliderKinds.AABB and e2.collider_kind == ColliderKinds.SPHERE:
106        return sphere_aabb_collision(e2.candidate_pose[0:3, 3], e2.collider_radii.item(),
107                                     e1.collider_bbox[0], e1.collider_bbox[1])
108    elif e1.collider_kind == ColliderKinds.AABB and e2.collider_kind == ColliderKinds.AABB:
109        return aabb_aabb_collision(e1.collider_bbox[0], e1.collider_bbox[1],
110                                   e2.collider_bbox[0], e2.collider_bbox[1])
111    else:
112        raise NotImplementedError((e1.collider_kind, e2.collider_kind))