robolib.physics

14def integrate_velocity_into_pose(pose: Arr, velocity: Arr, dt: float) -> np.ndarray:
15    """integrates pose[t+1] = pose[t] + v[t] * dt with v=(v, w) twist. Uses rodrigues formula (batched)."""
16    pose = pose.numpy() if isinstance(pose, QRField) else pose                 # (N, 4, 4)
17    velocity = velocity.numpy() if isinstance(velocity, QRField) else velocity # (N, 6) [m/s]
18    res = np.float32(pose) @ btrexp(bskewa(velocity * dt)).astype("float32") # (N, 4, 4) [m]
19    if random.random() >= 0.99: # renormalize every now and then to ensure rotation matrix is still reliable
20        for i in range(len(pose)):
21            res[i, 0:3, 0:3] = renormalize_rotation_matrix(res[i, 0:3, 0:3])
22    return res

24def motion_level1(motion_input: Arr, max_velocities: Arr) -> Arr:
25    """level 1 physics: velocity is the motion input. Returns it after clamping"""
26    velocity = np.clip(motion_input, -1, 1) * max_velocities              # [m/s]
27    return velocity

29def motion_level2(motion_input: Arr, dt: float, current_velocity: Arr,
30                  max_accelerations: Arr, drag_coefficient: Arr) -> tuple[Arr, Arr]:
31    """
32    level 2 physics: integrates acceleration (motion input) + drag (ct) into the current velocity (A = (a[t] - k * v[t])
33    - v[t+1] = v[t] + dt * (a[t] - k * v[t])
34    """
35    clamped_acceleration = np.clip(motion_input, -1, 1) * max_accelerations
36    drag = current_velocity * drag_coefficient                            # [m/s**2]; (N, 6) * (N, )
37    acceleration = clamped_acceleration - drag                            # [m/s**2]; (N, 6) - (N, 6)
38    velocity = current_velocity + dt * acceleration                       # [m/s]
39    # clamp off super small velocities so they don't go "forever" in the tests
40    velocity = np.where(np.abs(velocity) < EPS, 0, velocity)              # [m/s]
41    return velocity, acceleration

10def make_grid_3d(entities: list[Entity], cell_size: Point3D) -> dict[Key3D, list[Entity]]:
11    """
12    Creates a 3d grid given all the objects based on their bounding boxes. Cell size is provided from outside.
13    Output dict is provided (not returned), so we can create it outside (e.g. prefill fixed ones)!
14    """
15    res: dict[Key3D, list[Entity]] = {}
16    assert (cell_size > 0).all(), cell_size
17    for entity in entities:
18        top_left, bottom_right = entity.collider_bbox[0:2]
19        first_cell = (top_left / cell_size).astype(int) # Note: possible bug with negative positions??
20        last_cell = (bottom_right / cell_size).astype(int)
21        nx, ny, nz = last_cell - first_cell + 1
22
23        for i in range(nx):
24            for j in range(ny):
25                for k in range(nz):
26                    key = (first_cell[0].item() + i, first_cell[1].item() + j, first_cell[2].item() + k)
27                    res.setdefault(key, []).append(entity)
28    return res

30def make_collision_cell_size(entities: QueryResult) -> Point3D:
31    """the collision cell size is twice them edian of all bboxes of all collidables with a model"""
32    bboxes = entities.model_bbox * entities.scale[..., None]
33    diffs = (bboxes[:, 1] - bboxes[:, 0]).numpy()
34    median = np.median(diffs, axis=0) # median across all 3 dimensions
35    res = median * 2 # heuristic: twice the median
36    logger.info(f"Median bbox size: {median}. Cell size: {res}")
37    assert (res > 0).all(), bboxes.numpy()
38    return res

 98def check_collision(e1: Entity, e2: Entity) -> bool:
 99    """checks the collision between two entities with the HasCollision component"""
100    if e1.collider_kind == ColliderKinds.SPHERE and e2.collider_kind == ColliderKinds.SPHERE:
101        return sphere_sphere_collision(e1.candidate_pose[0:3, 3], e1.collider_radii.item(),
102                                       e2.candidate_pose[0:3, 3], e2.collider_radii.item())
103    elif e1.collider_kind == ColliderKinds.SPHERE and e2.collider_kind == ColliderKinds.AABB:
104        return sphere_aabb_collision(e1.candidate_pose[0:3, 3], e1.collider_radii.item(),
105                                     e2.collider_bbox[0], e2.collider_bbox[1])
106    elif e1.collider_kind == ColliderKinds.AABB and e2.collider_kind == ColliderKinds.SPHERE:
107        return sphere_aabb_collision(e2.candidate_pose[0:3, 3], e2.collider_radii.item(),
108                                     e1.collider_bbox[0], e1.collider_bbox[1])
109    elif e1.collider_kind == ColliderKinds.AABB and e2.collider_kind == ColliderKinds.AABB:
110        return aabb_aabb_collision(e1.collider_bbox[0], e1.collider_bbox[1],
111                                   e2.collider_bbox[0], e2.collider_bbox[1])
112    else:
113        raise NotImplementedError((e1.collider_kind, e2.collider_kind))

68def aabb_aabb_collision(top_left1: Point3D, bottom_right1: Point3D,
69                        top_left2: Point3D, bottom_right2: Point3D, eps: float=1e-5) -> bool:
70    """
71    Checks if two aabbs are colliding. We do this by doing one check on each ax.
72    We are thus reducing the complexity from 3D to three independent 1D checks.
73
74    For 1D, assume the two lines given by top_left and bottom_right: [tl1, br1], [tl2, br2]. There are 5 cases:
75    - tl1, br1, tl2, br2 -> not overlapping
76    - tl1, tl2, br1, br2 -> overlapping
77    - tl1, tl2, br2, br1 -> overlapping
78    - tl2, tl1, br1, br2 -> overlapping
79    - tl2, tl1, br2, br1 -> overlapping
80    - tl2, br2, tl1, br1 -> not overlapping
81
82    In all the cases, we can look at the midpoints c1, c2 and the actual distances between the two corners:
83    h1 = (br1 - tl1) / 2, h2 = (br2 - tl2) / 2. The lines are overlapping only if dist(c1, c2) < h1 + h2 across all dims
84    """
85
86    h1 = (bottom_right1 - top_left1) / 2
87    h2 = (bottom_right2 - top_left2) / 2
88
89    c1 = top_left1 + h1
90    c2 = top_left2 + h2
91
92    # signed separation on each axis, like the sphere checks above: < 0 means overlapping on that axis
93    gap = abs(c1 - c2) - (h1 + h2)
94    if check := (gap[0] < -eps) and (gap[1] < -eps) and (gap[2] < -eps):
95        logger.log_every_s(f"Collision. {top_left1=} {bottom_right1=} {top_left2=} {bottom_right2=}", "DEBUG", True)
96    return check

57def sphere_aabb_collision(center: Point3D, radius: float, top_left: Point3D,
58                          bottom_right: Point3D, eps: float=1e-5) -> bool:
59    """checks if a sphere is colliding with an aabb"""
60    xr, yr, zr = _get_closest_point_on_box(top_left, bottom_right, center)
61    dist_sq = (center[0] - xr)**2 + (center[1] - yr)**2 + (center[2] - zr)**2
62    dist: float = (dist_sq - radius ** 2).item()
63    if check := dist < -eps:
64        logger.log_every_s(f"Collision. center={fmt(center)} radius={radius:.3f} top_left={fmt(top_left)} "
65                           f"bottom_right={fmt(bottom_right)}", "DEBUG", True)
66    return check

47def sphere_sphere_collision(pos1: Point3D, radius1: float, pos2: Point3D, radius2: float, eps: float=1e-5) -> bool:
48    """checks if two spheres are colliding"""
49    max_dist = radius1 + radius2
50    diff_vec = pos1 - pos2
51    dist: float = ((diff_vec**2).sum() - max_dist**2).item() # instead of np.linalg.norm(a.pos-b.pos) < max_dist
52    if check := dist < -eps:
53        logger.log_every_s(f"Collision. pos1={fmt(pos1)} radius1={radius1:.3f} "
54                           f"pos2={fmt(pos2)} radius2={radius2:.3f}", "DEBUG", True)
55    return check