Skip to content

Utilities

Rotation Conversions

rot_utils

Rotation utility functions for conversions between representations.

quaternion_to_matrix(quat: np.ndarray) -> np.ndarray

Convert quaternion to 3x3 rotation matrix.

Input

quat: Quaternion in (w, x, y, z) format, shape (4,)

Output: rotation matrix, shape (3, 3)

Source code in drone_planning/utils/rot_utils.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def quaternion_to_matrix(quat: np.ndarray) -> np.ndarray:
    """
    Convert quaternion to 3x3 rotation matrix.

    Input:
        quat: Quaternion in (w, x, y, z) format, shape (4,)
    Output:
        rotation matrix, shape (3, 3)
    """
    quat = np.asarray(quat, dtype=np.float64)
    w, x, y, z = quat
    # Normalize
    norm = np.sqrt(w * w + x * x + y * y + z * z)
    w, x, y, z = w / norm, x / norm, y / norm, z / norm

    return np.array(
        [
            [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
            [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
            [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
        ]
    )

matrix_to_quaternion(rot: np.ndarray) -> np.ndarray

Convert 3x3 rotation matrix to quaternion.

Input

rot: Rotation matrix, shape (3, 3)

Output: quaternion in (w, x, y, z) format, shape (4,)

Source code in drone_planning/utils/rot_utils.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def matrix_to_quaternion(rot: np.ndarray) -> np.ndarray:
    """
    Convert 3x3 rotation matrix to quaternion.

    Input:
        rot: Rotation matrix, shape (3, 3)
    Output:
        quaternion in (w, x, y, z) format, shape (4,)
    """
    rot = np.asarray(rot, dtype=np.float64)
    trace = np.trace(rot)

    if trace > 0:
        s = 0.5 / np.sqrt(trace + 1.0)
        w = 0.25 / s
        x = (rot[2, 1] - rot[1, 2]) * s
        y = (rot[0, 2] - rot[2, 0]) * s
        z = (rot[1, 0] - rot[0, 1]) * s
    elif rot[0, 0] > rot[1, 1] and rot[0, 0] > rot[2, 2]:
        s = 2.0 * np.sqrt(1.0 + rot[0, 0] - rot[1, 1] - rot[2, 2])
        w = (rot[2, 1] - rot[1, 2]) / s
        x = 0.25 * s
        y = (rot[0, 1] + rot[1, 0]) / s
        z = (rot[0, 2] + rot[2, 0]) / s
    elif rot[1, 1] > rot[2, 2]:
        s = 2.0 * np.sqrt(1.0 + rot[1, 1] - rot[0, 0] - rot[2, 2])
        w = (rot[0, 2] - rot[2, 0]) / s
        x = (rot[0, 1] + rot[1, 0]) / s
        y = 0.25 * s
        z = (rot[1, 2] + rot[2, 1]) / s
    else:
        s = 2.0 * np.sqrt(1.0 + rot[2, 2] - rot[0, 0] - rot[1, 1])
        w = (rot[1, 0] - rot[0, 1]) / s
        x = (rot[0, 2] + rot[2, 0]) / s
        y = (rot[1, 2] + rot[2, 1]) / s
        z = 0.25 * s

    return np.array([w, x, y, z])

rpy_to_matrix(roll: float, pitch: float, yaw: float) -> np.ndarray

Convert roll-pitch-yaw (XYZ Euler angles) to rotation matrix.

Input

roll: Rotation around X axis (radians) pitch: Rotation around Y axis (radians) yaw: Rotation around Z axis (radians)

Output: rotation matrix, shape (3, 3)

Source code in drone_planning/utils/rot_utils.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def rpy_to_matrix(roll: float, pitch: float, yaw: float) -> np.ndarray:
    """
    Convert roll-pitch-yaw (XYZ Euler angles) to rotation matrix.

    Input:
        roll: Rotation around X axis (radians)
        pitch: Rotation around Y axis (radians)
        yaw: Rotation around Z axis (radians)
    Output:
        rotation matrix, shape (3, 3)
    """
    cr, sr = np.cos(roll), np.sin(roll)
    cp, sp = np.cos(pitch), np.sin(pitch)
    cy, sy = np.cos(yaw), np.sin(yaw)

    return np.array(
        [
            [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
            [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
            [-sp, cp * sr, cp * cr],
        ]
    )

matrix_to_rpy(rot: np.ndarray) -> tuple[float, float, float]

Convert rotation matrix to roll-pitch-yaw (XYZ Euler angles).

Input

rot: Rotation matrix, shape (3, 3)

Output: (roll, pitch, yaw) in radians

Source code in drone_planning/utils/rot_utils.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def matrix_to_rpy(rot: np.ndarray) -> tuple[float, float, float]:
    """
    Convert rotation matrix to roll-pitch-yaw (XYZ Euler angles).

    Input:
        rot: Rotation matrix, shape (3, 3)
    Output:
        (roll, pitch, yaw) in radians
    """
    rot = np.asarray(rot, dtype=np.float64)

    if abs(rot[2, 0]) < 1.0 - 1e-10:
        pitch = -np.arcsin(rot[2, 0])
        roll = np.arctan2(rot[2, 1] / np.cos(pitch), rot[2, 2] / np.cos(pitch))
        yaw = np.arctan2(rot[1, 0] / np.cos(pitch), rot[0, 0] / np.cos(pitch))
    else:
        # Gimbal lock
        yaw = 0.0
        if rot[2, 0] < 0:
            pitch = np.pi / 2
            roll = np.arctan2(rot[0, 1], rot[0, 2])
        else:
            pitch = -np.pi / 2
            roll = np.arctan2(-rot[0, 1], -rot[0, 2])

    return roll, pitch, yaw

axis_angle_to_matrix(axis: np.ndarray, angle: float) -> np.ndarray

Convert axis-angle representation to rotation matrix (Rodrigues formula).

Input

axis: Unit vector representing rotation axis, shape (3,) angle: Rotation angle in radians

Output: rotation matrix, shape (3, 3)

Source code in drone_planning/utils/rot_utils.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def axis_angle_to_matrix(axis: np.ndarray, angle: float) -> np.ndarray:
    """
    Convert axis-angle representation to rotation matrix (Rodrigues formula).

    Input:
        axis: Unit vector representing rotation axis, shape (3,)
        angle: Rotation angle in radians
    Output:
        rotation matrix, shape (3, 3)
    """
    axis = np.asarray(axis, dtype=np.float64)
    axis = axis / np.linalg.norm(axis)

    K = np.array(
        [[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]]
    )

    return np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)

matrix_to_axis_angle(rot: np.ndarray) -> tuple[np.ndarray, float]

Convert rotation matrix to axis-angle representation.

Input

rot: Rotation matrix, shape (3, 3)

Output: (axis, angle) where axis is unit vector shape (3,), angle in radians

Source code in drone_planning/utils/rot_utils.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def matrix_to_axis_angle(rot: np.ndarray) -> tuple[np.ndarray, float]:
    """
    Convert rotation matrix to axis-angle representation.

    Input:
        rot: Rotation matrix, shape (3, 3)
    Output:
        (axis, angle) where axis is unit vector shape (3,), angle in radians
    """
    rot = np.asarray(rot, dtype=np.float64)

    angle = np.arccos(np.clip((np.trace(rot) - 1) / 2, -1, 1))

    if angle < 1e-10:
        return np.array([1.0, 0.0, 0.0]), 0.0

    if abs(angle - np.pi) < 1e-10:
        # Find column of (R + I) with largest norm
        B = rot + np.eye(3)
        col_norms = np.linalg.norm(B, axis=0)
        idx = np.argmax(col_norms)
        axis = B[:, idx] / col_norms[idx]
        return axis, angle

    axis = np.array(
        [rot[2, 1] - rot[1, 2], rot[0, 2] - rot[2, 0], rot[1, 0] - rot[0, 1]]
    )
    axis = axis / (2 * np.sin(angle))

    return axis, angle

Video Recorder

video_recorder

Offscreen MP4 recorder for :class:PyBulletEnv.

Records MP4 clips from a headless PyBullet session by rendering frames via getCameraImage and piping them into ffmpeg. Works with a DIRECT connection (no GUI window required), which makes it safe to use inside build pipelines, CI, or a doc-rendering driver.

Typical usage::

from drone_planning.utils.video_recorder import CameraView, VideoRecorder

with VideoRecorder(env, "output.mp4", camera=CameraView()) as rec:
    rec.hold(seconds=0.5)            # pause on the start pose
    rec.play(configurations)         # one frame per configuration
    rec.hold(seconds=1.0)            # pause on the final pose

The recorder assumes ffmpeg is on PATH. Frames are rendered with the CPU TinyRenderer.

CameraView(target: tuple[float, float, float] = (2.6, 0.0, 1.0), distance: float = 6.0, yaw: float = -35.0, pitch: float = -25.0, roll: float = 0.0, fov: float = 45.0, near: float = 0.05, far: float = 30.0) dataclass

Static orbit-camera parameters passed to computeViewMatrixFromYawPitchRoll.

target points at a world-space location, distance is the orbit radius, and yaw/pitch orbit around target in degrees. fov is the vertical field of view of the projection. The camera does not move during a clip.

VideoRecorder(env: 'PyBulletEnv', output_path: Path | str, fps: int = 30, resolution: tuple[int, int] = (960, 540), camera: CameraView = CameraView(), light_direction: tuple[float, float, float] = (0.55, -0.4, 1.3), crf: int = 28) dataclass

Render frames from a :class:PyBulletEnv into an H.264 MP4.

The recorder owns a running ffmpeg subprocess while the context is open; frames are written directly to its stdin in rgba format.

capture() -> None

Render the current env state as a single frame into the clip.

Source code in drone_planning/utils/video_recorder.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def capture(self) -> None:
    """Render the current env state as a single frame into the clip."""
    w, h = self.resolution
    img = self.env.sim.client.getCameraImage(
        width=w,
        height=h,
        viewMatrix=self._view,
        projectionMatrix=self._proj,
        shadow=1,
        lightDirection=list(self.light_direction),
        renderer=pb.ER_TINY_RENDERER,
    )
    rgba = np.asarray(img[2], dtype=np.uint8).reshape(h, w, 4)
    self._proc.stdin.write(rgba.tobytes())

hold(seconds: float = 0.5) -> None

Freeze on the current configuration for seconds.

Source code in drone_planning/utils/video_recorder.py 
125
126
127
128
def hold(self, seconds: float = 0.5) -> None:
    """Freeze on the current configuration for *seconds*."""
    for _ in range(max(1, int(round(seconds * self.fps)))):
        self.capture()

play(configurations: np.ndarray) -> None

Capture one frame per full-DOF configuration.

Source code in drone_planning/utils/video_recorder.py 
130
131
132
133
134
def play(self, configurations: np.ndarray) -> None:
    """Capture one frame per full-DOF configuration."""
    for q in configurations:
        self.env.set_configuration(q)
        self.capture()