robolib.utils

55def get_project_root() -> Path:
56    """The project root (this file lives at src/robolib/utils/utils.py, so root is 3 levels up)"""
57    return Path(__file__).parents[3]

51def fmt(arr: np.ndarray, precision: int=3) -> str:
52    """formats a numpy array for logging as tuple"""
53    return "(" + ", ".join(f"{float(x):.{precision}g}" for x in arr) + ")"

23@dataclass
24class FPVData:
25    """The first perso view camera data. Latest frame (as bytes), a lock (for concurrency) and a way to compress it"""
26    frame: bytes
27    frame_shape: tuple[int, int, int]
28    lock: threading.Lock = field(default_factory=threading.Lock)
29    frame_id: int = 0 # used for de-dup on client side
30    _frame_compressed: bytes | None = None
31
32    @property
33    def frame_compressed(self) -> bytes:
34        """compress the frame on demand when the client requests a get_state_with_frame. NOTE: use lock on call."""
35        if self._frame_compressed is None:
36            self._frame_compressed = zlib.compress(self.frame, level=1)
37        return self._frame_compressed

59def make_arr(*data):
60    """makes a numpy array from variable data, e.g. make_arr(1,2,3) <=> np.array([1,2,3], "float32")"""
61    return np.float32(data)

39class FixedSizeDict(OrderedDict):
40    """An dict with a fixed size. Useful for caching purposes."""
41    def __init__(self, *args, maxlen: int = 0, **kwargs):
42        self._maxlen = maxlen
43        super().__init__(*args, **kwargs)
44
45    def __setitem__(self, key, value):
46        existed_before_add = key in self
47        super().__setitem__(key, value)
48        if not existed_before_add and len(self) > self._maxlen:
49            self.popitem(False)

19@dataclass
20class RlModelWithTexture:
21    """basic raylib model with an optional texture applied to it"""
22    model: RlModel
23    texture: RlTexture | None
24    _bbox: tuple[np.ndarray, np.ndarray] | None = None
25
26    def __post_init__(self):
27        if self.texture is not None:
28            self.model.materials.maps[rl.MATERIAL_MAP_ALBEDO].texture = self.texture
29
30    @staticmethod
31    def from_path(model_path: Path | str, texture_path: Path | str | None = None) -> RlModelWithTexture:
32        """builds a RLModelWithTexture from a path"""
33        model_path = Path(model_path)
34        texture_path = texture_path and Path(texture_path)
35        assert not model_path.is_absolute(), model_path
36        assert texture_path is None or not texture_path.is_absolute(), texture_path
37        assert model_path.exists(), f"Model path: '{model_path}' doesn't exist."
38        logger.info(f"Loading model from '{model_path}'")
39        model = rl.LoadModel(str(model_path).encode())
40        texture = rl.LoadTexture(str(texture_path).encode()) if texture_path is not None else None
41        return RlModelWithTexture(model=model, texture=texture)
42
43    @property
44    def bbox(self) -> tuple[np.ndarray, np.ndarray]:
45        """gets the top left and bottom right 3D points of this mesh. Used for collisions. Cahed once after 1st call"""
46        if self._bbox is None:
47            mins, maxs = np.float32([1<<31, 1<<31, 1<<31]), np.float32([-1<<31, -1<<31, -1<<31])
48            for i in range(self.model.meshCount):
49                mesh: RlMesh = self.model.meshes[i] # of type struct Mesh &
50                buf = rl.ffi.buffer(mesh.vertices, mesh.vertexCount * 3 * 4)
51                vertices = np.frombuffer(buf, "float32").reshape(-1, 3)
52                mins = np.minimum(mins, np.min(vertices, axis=0))
53                maxs = np.maximum(maxs, np.max(vertices, axis=0))
54            self._bbox = mins, maxs
55        return self._bbox
56
57    def __del__(self):
58        if not rl.IsWindowReady(): # in case rl.CloseWindow() was called
59            return
60        if self.model is not None:
61            rl.UnloadModel(self.model)
62        if self.texture is not None:
63            rl.UnloadTexture(self.texture)

73def rl_get_device_and_renderer() -> tuple[str, str]:
74    """returns the device and renderer used by raylib for logging"""
75    gl_renderer = 0x1F01
76    gl_get_string = rl.ffi.cast("unsigned char *(*)(unsigned int)", rl.rlGetProcAddress(b"glGetString"))
77    gstr = lambda e: rl.ffi.string(rl.ffi.cast("char*", gl_get_string(e))).decode()
78    renderer = gstr(gl_renderer)
79    software = ("llvmpipe", "softpipe", "swrast", "lavapipe")
80    device = "CPU (software rasterizer)" if any(s in renderer.lower() for s in software) else "GPU"
81    return device, renderer

65def vec3_from(arr: np.ndarray) -> RlVector3:
66    """convers the numpy array to a raylib vector"""
67    return rl.ffi.from_buffer("Vector3 *", np.ascontiguousarray(arr))[0]

69def rl_vec3_to_np(vec: RlVector3) -> np.ndarray:
70    """returns a np array from a raylib vec3"""
71    return np.array([vec.x, vec.y, vec.z], "float32")

15class Clock:
16    """clock used for physics with fixed DT in main loops"""
17
18    def __init__(self, dt: float, max_ticks: int):
19        self.dt = dt
20        self.max_ticks = max_ticks
21        self.prev_time = time.perf_counter()
22        self.accumulator = 0
23        self.tick_count = -1
24
25    def tick(self):
26        """tick once by adding the delta between prev frame and now"""
27        now = time.perf_counter()
28        frame_time = now - self.prev_time
29        self.prev_time = now
30        self.accumulator += frame_time
31        self.tick_count += 1
32
33    def drain(self):
34        """drain the accumulator. in main loop: for _ in clock.drain(): ..."""
35        n_ticks = 0
36        while self.accumulator >= self.dt and n_ticks < self.max_ticks:
37            yield
38            self.accumulator -= self.dt
39            n_ticks += 1
40        self.accumulator = min(self.accumulator, self.dt) # Drop residual debt instead of it piling up across frames
41
42    def wait(self):
43        """waits the leftover time in case the previous tick ran too fast to maintain consistent FPS"""
44        time.sleep(max(self.dt - (time.perf_counter() - self.prev_time), 0))
45
46    def wait_and_tick(self):
47        """calls wait() then tick(). Put this at the beginning of the main loop :)"""
48        self.wait()
49        self.tick()

11def norm(x: np.ndarray) -> np.ndarray:
12    """L2 normalize a vector"""
13    return (x / np.linalg.norm(x, ord=2)).astype("float32")

15def rot_from_forward_up(forward: Point3D, up: Point3D) -> Rot3x3:
16    """right hand rule: x x y = z, y x z = x, z x x = y; z = forward. (3, ) + (3, ) -> (3, 3)"""
17    z = norm(forward) # first, normalize forward, as we care about it's orientation only. z = ||z||
18    y_init = np.float32(up) # up = Y
19    x = norm(np.cross(y_init, z)) # second build x=||y_init x z|| as the cross of existing forward and up
20    y = np.cross(z, x) # third, get the new y as the cross of the already built x and z.
21    R = np.column_stack([x, y, z]) # NOTE: somehow it's left-handed and we make it right-handed here
22    assert is_rot(R), f"\n-{R=}\n-{R@R.T=}"
23    return R

46def pose_from_position_target_up(position: Point3D, target: Point3D, up: Point3D) -> Pose4x4:
47    """returns a 4x4 pose matrix from initial position+target+up (3, ) + (3, ) + (3, ) -> (4, 4)"""
48    pose = np.eye(4, dtype="float32")
49    pose[0:3, 0:3] = rot_from_forward_up(forward=np.float32(target) - position, up=up)
50    pose[0:3, 3] = position
51    return pose

25def pose_to_trans_euler(pose: Pose4x4) -> Point6D:
26    """converts a pose to a 6DoF translation + euler vector, mostly for printing reasons (4, 4) -> (6, )"""
27    return np.float32([*pose[0:3, 3], *tr2rpy(pose)])

29def pose_from_trans_euler(position_rpy: Point6D) -> Pose4x4:
30    """converts a position + a euler vector to a Pose. (6, ) -> (4, 4)"""
31    res = rpy2tr(position_rpy[-3:]).astype("float32")
32    res[0:3, 3] = position_rpy[0:3]
33    return res

35def renormalize_rotation_matrix(rot: Rot3x3, tol: float=1e10) -> Rot3x3:
36    """uses SVD to renormalize a rotation matrix if it drifted too much (i.e. during trajectory updates)"""
37    if is_rot(rot, tol=tol):
38        return rot
39    logger.debug(f"Renormalizing R as err is {np.linalg.norm((rot @ rot.T) - np.eye(3)):.7f}")
40    U, _, Vt = np.linalg.svd(rot)
41    D = np.diag([1, 1, np.linalg.det(U @ Vt)]) # force det=1 (not -1)
42    new_rot = U @ D @ Vt
43    assert is_rot(new_rot, tol=1e10), new_rot
44    return new_rot

53def make_radius(size: Point3D) -> float:
54    """creates a radius from a size and a scale"""
55    return math.sqrt(size[0]**2 + size[1]**2 + size[2]**2) / 2

178def btrexp(S: Mat4x4) -> Pose4x4:
179    """
180    Batched matrix exponential of an se(3) element -> SE(3) (a screw motion).
181    ``S`` is the batched (Nx)4x4 augmented skew-symmetric matrix produced by ``bskewa``.
182        >>> btrexp(bskewa(np.array([[1, 0, 0, 0, 0, 0]])))
183    """
184    n = len(S)
185    tw = bvexa(S) # (N, 6)
186    v = tw[:, 0:3] # (3, )
187    w = tw[:, 3:6] # (3, )
188
189    w_norm = np.linalg.norm(w, axis=1) # (N, )
190    v_norm = np.linalg.norm(v, axis=1) # (N, )
191    # theta is the screw magnitude: rotation dominates, else it's a pure translation
192    theta = np.where(w_norm >= 20 * _eps, w_norm, v_norm)[:, None] # (N, 1)
193
194    with np.errstate(divide="ignore", invalid="ignore"):   # 0/0 on zero-twist is expected; nan_to_num handles it
195        v = np.nan_to_num(v / theta, nan=0) # (N, 3)
196        w = np.nan_to_num(w / theta, nan=0) # (N, 3)
197
198    skw = bskew(w)                                                      # (N, 3, 3)
199    theta_n = theta[..., None]                                          # (N, 1, 1)
200    st, ct = np.sin(theta_n), np.cos(theta_n)                           # (N, 1, 1), (N, 1, 1)
201    eye = np.eye(3)[None].repeat(n, axis=0)                             # (N, 3, 3)
202    R = eye + st * skw + (1.0 - ct) * skw @ skw                         # (N, 3, 3) Rodrigues
203    V = eye * theta_n + (1.0 - ct) * skw + (theta_n - st) * skw @ skw   # (N, 3, 3) translation map
204
205    T = np.eye(4)[None].repeat(n, axis=0)                               # (N, 4, 4)
206    T[:, 0:3, 0:3] = R
207    T[:, 0:3, 3] = np.einsum("bij,bj->bi", V, v)
208    return T

167def bskewa(v: Point6D) -> Mat4x4:
168    """batched skewa"""
169    assert len(v.shape) == 2 and v.shape[1] == 6, v.shape
170    zero = np.zeros((len(v), ), "float32")
171    res = np.array([
172        [ zero,     -v[:, 5],    v[:, 4],  v[:, 0]],
173        [ v[:, 5],   zero,      -v[:, 3],  v[:, 1]],
174        [-v[:, 4],   v[:, 3],    zero,     v[:, 2]],
175        [ zero,      zero,       zero,     zero]], dtype=np.float64) # (4, 4, N)
176    return np.permute_dims(res, (2, 0, 1)) # (N, 4, 4)