Simplify render example. Remove inactive dependencies.

Author: amacati

crazyflow/sim/visualize.py

@@ -71,6 +71,8 @@ def draw_points(sim: Sim, points: NDArray, rgba: NDArray | None = None, size: fl
 
 def _rotation_matrix_from_points(p1: NDArray, p2: NDArray) -> R:
     """Generate rotation matrices that align their z-axis to p2-p1."""
+    p1, p2 = p1.copy(), p2.copy()  # Make sure we don't modify the original arrays
+    p2[np.linalg.norm(p2 - p1, axis=-1) < 1e-6] += 1e-6
     z_axis = (v := p2 - p1) / np.linalg.norm(v, axis=-1, keepdims=True)
     random_vector = np.random.rand(*z_axis.shape)
     x_axis = (v := np.cross(random_vector, z_axis)) / np.linalg.norm(v, axis=-1, keepdims=True)

examples/render.py

@@ -1,14 +1,10 @@
 from collections import deque
 
-import einops
-import mujoco
 import numpy as np
-from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
-from numpy.typing import NDArray
-from scipy.spatial.transform import Rotation as R
 
 from crazyflow.constants import GRAVITY, MASS
 from crazyflow.sim import Physics, Sim
+from crazyflow.sim.visualize import draw_line
 
 
 def main():
@@ -18,60 +14,23 @@ def main():
     fps = 60
     cmd = np.zeros((sim.n_worlds, sim.n_drones, 4))
     cmd[..., 0] = MASS * GRAVITY * 1.2
+    rgbas = np.random.default_rng(0).uniform(0, 1, (n_drones, 4))
+    rgbas[..., 3] = 1.0
 
-    pos = deque(maxlen=16)
-    rot = deque(maxlen=15)
+    pos = deque(maxlen=15)
 
     for i in range(int(5 * sim.control_freq)):
         sim.attitude_control(cmd)
         sim.step(sim.freq // sim.control_freq)
         if i % 20 == 0:
             pos.append(sim.data.states.pos[0, :])
-            if len(pos) > 1:
-                rot.append(rotation_matrix_from_points(pos[-2], pos[-1]))
         if ((i * fps) % sim.control_freq) < fps:
-            render_traces(sim.viewer, pos, rot)
+            lines = np.array(pos)
+            for i in range(n_drones):
+                draw_line(sim, lines[:, i, :], rgbas[i, :], start_size=0.3, end_size=3.0)
             sim.render()
     sim.close()
 
 
-def render_traces(viewer: MujocoRenderer, pos: deque[NDArray], rot: deque[R]):
-    """Render traces of the drone trajectories."""
-    if len(pos) < 2 or viewer is None:
-        return
-
-    n_trace, n_drones = len(rot), len(pos[0])
-    pos = np.array(pos)
-    sizes = np.zeros((n_trace, n_drones, 3))
-    rgbas = np.zeros((n_trace, n_drones, 4))
-    sizes[..., 2] = np.linalg.norm(pos[1:] - pos[:-1], axis=-1)
-    sizes[..., :2] = np.linspace(1.0, 5.0, n_trace)[:, None, None]
-    mats = np.zeros((n_trace, n_drones, 9))
-    for i in range(n_trace):
-        mats[i, :] = einops.rearrange(rot[i].as_matrix(), "d n m -> d (n m)")
-    rgbas = np.zeros((n_trace, n_drones, 4))
-    np.random.seed(0)  # Ensure consistent colors
-    rgbas[..., :3] = np.random.uniform(0, 1, (1, n_drones, 3))
-    rgbas[..., 3] = np.linspace(0, 1, n_trace)[:, None]
-
-    for i in range(n_trace):
-        for j in range(n_drones):
-            viewer.viewer.add_marker(
-                type=mujoco.mjtGeom.mjGEOM_LINE,
-                size=sizes[i, j],
-                pos=pos[i][j],
-                mat=mats[i, j],
-                rgba=rgbas[i, j],
-            )
-
-
-def rotation_matrix_from_points(p1: NDArray, p2: NDArray) -> R:
-    z_axis = (v := p2 - p1) / np.linalg.norm(v, axis=-1, keepdims=True)
-    random_vector = np.random.rand(*z_axis.shape)
-    x_axis = (v := np.cross(random_vector, z_axis)) / np.linalg.norm(v, axis=-1, keepdims=True)
-    y_axis = np.cross(z_axis, x_axis)
-    return R.from_matrix(np.stack((x_axis, y_axis, z_axis), axis=-1))
-
-
 if __name__ == "__main__":
     main()

pixi.lock

@@ -28,7 +28,6 @@ dependencies = [
     "ml_collections",
     "casadi",
     "numpy",
-    # "drone-models @ git+https://github.com/utiasDSL/drone-models.git@08ffb18",
 ]
 requires-python = ">=3.10"