diff --git a/media/showcase/multi_robot_3d.gif b/media/showcase/multi_robot_3d.gif index ef78274d..06734e6c 100644 Binary files a/media/showcase/multi_robot_3d.gif and b/media/showcase/multi_robot_3d.gif differ diff --git a/src/cbfkit/utils/visualization.py b/src/cbfkit/utils/visualization.py index 0eb1ba01..d087a5eb 100644 --- a/src/cbfkit/utils/visualization.py +++ b/src/cbfkit/utils/visualization.py @@ -7,7 +7,6 @@ matplotlib (MP4/GIF), and Manim (high-quality MP4) backends. """ -import warnings from typing import Any, List, Optional import numpy as np @@ -173,6 +172,7 @@ def visualize_3d_multi_robot( include_min_distance_to_obstacles_plot: bool = False, threshold: Optional[float] = None, backend: str = "plotly", + safety_radius: float = 0.25, ): """Animate a 3D multi-robot system with optional distance subplots. @@ -192,6 +192,14 @@ def visualize_3d_multi_robot( ``"plotly"`` (default), ``"matplotlib"``, or ``"manim"``. Manim accepts a quality suffix: ``"manim-low"`` (default), ``"manim-medium"``, ``"manim-high"``, ``"manim-production"``. + threshold : float, optional + Minimum-separation threshold drawn as a dashed reference line on the + inter-robot distance panel. Pass the same value the controller + enforces so the plot can be read as satisfied / violated. + safety_radius : float + Radius of the per-robot safety bubble drawn by the manim backend. Two + bubbles touch at ``2 * safety_radius``; pass ``threshold / 2`` to keep + the drawing consistent with the enforced constraint. """ from cbfkit.utils.visualizations.helpers_3d import _compute_distance_metrics @@ -276,20 +284,6 @@ def visualize_3d_multi_robot( from cbfkit.utils.visualizations.manim_3d_multi_robot import render_multi_robot_3d _require_manim() - if include_min_distance_plot: - warnings.warn( - "Manim backend does not support inline subplot panels. " - "include_min_distance_plot will be ignored.", - UserWarning, - stacklevel=2, - ) - if include_min_distance_to_obstacles_plot: - warnings.warn( - "Manim backend does not support inline subplot panels. " - "include_min_distance_to_obstacles_plot will be ignored.", - UserWarning, - stacklevel=2, - ) save_path = animation_filename if save_animation else None return render_multi_robot_3d( states=states, @@ -307,9 +301,11 @@ def visualize_3d_multi_robot( ellipse_rotations=ellipse_rotations, save_path=save_path, quality=quality, + safety_radius=safety_radius, goal_dists=goal_dists, min_dists=min_dists if include_min_distance_plot else None, obs_dists=obs_dists if include_min_distance_to_obstacles_plot else None, + threshold=threshold, ) else: from cbfkit.utils.visualizations.matplotlib_3d_multi_robot import _visualize_3d_matplotlib diff --git a/src/cbfkit/utils/visualizations/helpers_3d.py b/src/cbfkit/utils/visualizations/helpers_3d.py index e5c7fee4..724d02ad 100644 --- a/src/cbfkit/utils/visualizations/helpers_3d.py +++ b/src/cbfkit/utils/visualizations/helpers_3d.py @@ -29,7 +29,7 @@ def _point_to_ellipsoid_distance(p, c, r, R): return -np.min(r) u_normalized = u / norm_u - inv_r_squared = np.diag(1.0 / (r ** 2)) + inv_r_squared = np.diag(1.0 / (r**2)) K = u_normalized.T @ R @ inv_r_squared @ R.T @ u_normalized s = 1.0 / np.sqrt(K) x = c + s * u_normalized @@ -60,9 +60,17 @@ def _ellipsoid_mesh(center, radii, rotation, n=12): return pts[:, 0], pts[:, 1], pts[:, 2], ii, jj, kk -def _compute_distance_metrics(states, desired_states, num_robots, sdim, - ellipse_centers, ellipse_radii, ellipse_rotations, - include_min_dist, include_obs_dist): +def _compute_distance_metrics( + states, + desired_states, + num_robots, + sdim, + ellipse_centers, + ellipse_radii, + ellipse_rotations, + include_min_dist, + include_obs_dist, +): """Pre-compute distance arrays used by both backends.""" N = len(states) @@ -71,15 +79,19 @@ def _compute_distance_metrics(states, desired_states, num_robots, sdim, for i in range(num_robots): idx = sdim * i goal_dists[:, i] = np.linalg.norm( - states[:, idx:idx + 3] - desired_states[idx:idx + 3], axis=1, + states[:, idx : idx + 3] - desired_states[idx : idx + 3], + axis=1, ) - # Min inter-robot distances (vectorized over robots) + # Min inter-robot distances (vectorized over robots). + # Only the first 3 columns of each robot's block are positions; including + # the remaining state (e.g. velocity) would report sqrt(|dp|^2 + |dv|^2), + # which overstates separation exactly when robots converge at speed. min_dists = None if include_min_dist: min_dists = np.zeros((N, num_robots)) for t in range(N): - positions = states[t, :].reshape(num_robots, sdim) + positions = states[t, : num_robots * sdim].reshape(num_robots, sdim)[:, :3] diffs = positions[:, np.newaxis, :] - positions[np.newaxis, :, :] dists = np.linalg.norm(diffs, axis=2) np.fill_diagonal(dists, np.inf) @@ -93,11 +105,14 @@ def _compute_distance_metrics(states, desired_states, num_robots, sdim, for t in range(N): for i in range(num_robots): idx = sdim * i - p = states[t, idx:idx + 3] + p = states[t, idx : idx + 3] dmin = np.inf for j in range(num_obs): d = _point_to_ellipsoid_distance( - p, ellipse_centers[j], ellipse_radii[j], ellipse_rotations[j], + p, + ellipse_centers[j], + ellipse_radii[j], + ellipse_rotations[j], ) dmin = min(dmin, d) obs_dists[t, i] = dmin diff --git a/src/cbfkit/utils/visualizations/manim_3d_multi_robot.py b/src/cbfkit/utils/visualizations/manim_3d_multi_robot.py index 1c6a7700..520ba020 100644 --- a/src/cbfkit/utils/visualizations/manim_3d_multi_robot.py +++ b/src/cbfkit/utils/visualizations/manim_3d_multi_robot.py @@ -27,6 +27,7 @@ RIGHT, UP, Axes, + DashedLine, Dot3D, Line, Sphere, @@ -86,21 +87,34 @@ def _build_chart_panel( num_robots: int, panel_width: float = 3.8, panel_height: float = 2.0, + threshold: float | None = None, ): """Build a 2D Axes with pre-drawn static line segments for each robot. - Returns ``(axes, title_mob, robot_lines)`` where ``robot_lines[i]`` is a - :class:`VGroup` of ``Line`` segments for robot *i* that will be revealed - progressively by an updater. + Returns ``(axes, title_mob, x_label, robot_lines, panel)`` where + ``robot_lines[i]`` is a :class:`VGroup` of ``Line`` segments for robot *i* + that will be revealed progressively by an updater. Parameters ---------- data : np.ndarray Shape ``(N, num_robots)`` — one column per robot. + threshold : float, optional + Safety threshold; drawn as a dashed red horizontal line so a viewer can + tell satisfaction from violation. Matches the plotly / matplotlib + backends, which draw the same reference line. """ N = len(data) t_max = (N - 1) * dt y_max = float(np.nanmax(data)) * 1.15 + if not np.isfinite(y_max): + # All-NaN data, or the all-inf min_dists a single robot produces. + # max() below would propagate NaN from its left argument. + y_max = 0.0 + if threshold is not None and threshold >= 0: + # Keep the line on-screen even when every sample sits below it + # (i.e. when the constraint is violated throughout). + y_max = max(y_max, float(threshold) * 1.15) if y_max < 1e-6: y_max = 1.0 @@ -133,6 +147,15 @@ def _build_chart_panel( robot_lines.append(segs) panel = VGroup(axes, title_mob, x_label, *robot_lines) + if threshold is not None and threshold >= 0: + panel.add( + DashedLine( + axes.c2p(0, float(threshold)), + axes.c2p(t_max, float(threshold)), + stroke_width=2, + color="#ff0000", + ) + ) return axes, title_mob, x_label, robot_lines, panel @@ -166,6 +189,8 @@ class MultiRobot3DScene(ThreeDScene): goal_dists: np.ndarray | None = None min_dists: np.ndarray | None = None obs_dists: np.ndarray | None = None + # Safety threshold drawn on the inter-robot panel (None = no reference line) + threshold: float | None = None # Scale factor to fit data into Manim's coordinate system (default ~7 units) _scale: float = 1.0 @@ -302,16 +327,16 @@ def updater(mob): panel_configs = [] if has_goal: - panel_configs.append(("Dist to Goal", self.goal_dists)) + panel_configs.append(("Dist to Goal", self.goal_dists, None)) if has_min: - panel_configs.append(("Min Inter-Robot Dist", self.min_dists)) + panel_configs.append(("Min Inter-Robot Dist", self.min_dists, self.threshold)) if has_obs: - panel_configs.append(("Min Obstacle Dist", self.obs_dists)) + panel_configs.append(("Min Obstacle Dist", self.obs_dists, None)) panel_height = min(2.0, 5.5 / max(n_panels, 1)) panel_gap = 0.4 - for p_idx, (p_title, p_data) in enumerate(panel_configs): + for p_idx, (p_title, p_data, p_threshold) in enumerate(panel_configs): _, _, _, robot_lines, panel = _build_chart_panel( title_text=p_title, data=p_data, @@ -319,6 +344,7 @@ def updater(mob): num_robots=n_robots, panel_width=3.8, panel_height=panel_height, + threshold=p_threshold, ) all_robot_lines.append(robot_lines) panels_group.add(panel) @@ -430,6 +456,7 @@ def render_multi_robot_3d( goal_dists: np.ndarray | None = None, min_dists: np.ndarray | None = None, obs_dists: np.ndarray | None = None, + threshold: float | None = None, ) -> str: """Render a multi-robot 3D animation using Manim. @@ -440,7 +467,11 @@ def render_multi_robot_3d( desired_states : np.ndarray ``(state_dim,)`` goal vector. safety_radius : float - Per-robot collision avoidance bubble radius. + Per-robot collision avoidance bubble radius. Two bubbles touch when + the robots are ``2 * safety_radius`` apart, so this should be half the + minimum separation the controller actually enforces. + threshold : float, optional + Minimum-separation threshold marked on the inter-robot distance panel. ellipse_centers : list, optional List of obstacle center positions ``[x, y, z]``. ellipse_radii : list, optional @@ -466,13 +497,21 @@ def render_multi_robot_3d( """ _require_manim() + import glob import os + import shutil # Configure Manim output config.quality = quality if save_path: - config.output_file = os.path.basename(save_path) + # Manim appends its own extension, so hand it the stem: passing + # "anim.gif" produced "anim.gif.mp4". The format must be set + # explicitly or a .gif request silently renders MP4. + stem, ext = os.path.splitext(os.path.basename(save_path)) + config.output_file = stem config.media_dir = os.path.dirname(save_path) or "./media" + if ext.lower() == ".gif": + config.format = "gif" # Inject data into the scene class MultiRobot3DScene.states = np.asarray(states) @@ -492,9 +531,31 @@ def render_multi_robot_3d( MultiRobot3DScene.goal_dists = goal_dists MultiRobot3DScene.min_dists = min_dists MultiRobot3DScene.obs_dists = obs_dists + MultiRobot3DScene.threshold = threshold scene = MultiRobot3DScene() scene.render() - # Return path to rendered file - return str(scene.renderer.file_writer.movie_file_path) + rendered = str(scene.renderer.file_writer.movie_file_path) + if not save_path: + return rendered + + # For GIF output the file writer still reports the .mp4 name, so fall back + # to the newest matching file under the media dir. + if not os.path.exists(rendered): + pattern = os.path.join(config.media_dir, "videos", "**", f"*{ext or '.mp4'}") + candidates = glob.glob(pattern, recursive=True) + if not candidates: + raise FileNotFoundError( + f"Manim did not produce a {ext or '.mp4'} file under {config.media_dir!r}." + ) + rendered = max(candidates, key=os.path.getmtime) + + # Manim writes under /videos//; copy to the path the + # caller asked for so the returned path is the file that actually exists. + # Without this the caller's path silently keeps whatever was there before. + if os.path.abspath(rendered) != os.path.abspath(save_path): + out_dir = os.path.dirname(os.path.abspath(save_path)) + os.makedirs(out_dir, exist_ok=True) + shutil.copy2(rendered, save_path) + return os.path.abspath(save_path) diff --git a/src/cbfkit/utils/visualizations/matplotlib_3d_multi_robot.py b/src/cbfkit/utils/visualizations/matplotlib_3d_multi_robot.py index 559faa78..025fbe9a 100644 --- a/src/cbfkit/utils/visualizations/matplotlib_3d_multi_robot.py +++ b/src/cbfkit/utils/visualizations/matplotlib_3d_multi_robot.py @@ -8,12 +8,27 @@ def _visualize_3d_matplotlib( - states, desired_states, desired_state_radius, num_robots, - ellipse_centers, ellipse_radii, ellipse_rotations, - x_lim, y_lim, z_lim, dt, sdim, title, save_animation, - animation_filename, include_min_distance_plot, - include_min_distance_to_obstacles_plot, threshold, - goal_dists, min_dists, obs_dists, + states, + desired_states, + desired_state_radius, + num_robots, + ellipse_centers, + ellipse_radii, + ellipse_rotations, + x_lim, + y_lim, + z_lim, + dt, + sdim, + title, + save_animation, + animation_filename, + include_min_distance_plot, + include_min_distance_to_obstacles_plot, + threshold, + goal_dists, + min_dists, + obs_dists, ): import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation @@ -61,8 +76,12 @@ def _visualize_3d_matplotlib( color = colors[i % num_robots] ax_traj.scatter( - desired_states[idx], desired_states[idx + 1], desired_states[idx + 2], - color=color, s=50, label=f"Desired State {i + 1}", + desired_states[idx], + desired_states[idx + 1], + desired_states[idx + 2], + color=color, + s=50, + label=f"Desired State {i + 1}", ) u = np.linspace(0, 2 * np.pi, 50) v = np.linspace(0, np.pi, 50) @@ -89,7 +108,15 @@ def _visualize_3d_matplotlib( ax_min_dist.set_ylabel("Minimum Distance [m]") ax_min_dist.set_title("Min Distance Between Robots") ax_min_dist.grid(True) - ax_min_dist.set_ylim(0, float(np.max(min_dists)) * 1.1) + # Include the threshold in the range, else a run that violates the + # constraint throughout renders as a clean plot with the reference + # line clipped off-screen. + md_top = float(np.max(min_dists)) * 1.1 + if not np.isfinite(md_top): + md_top = 0.0 + if threshold is not None and threshold >= 0: + md_top = max(md_top, float(threshold) * 1.15) + ax_min_dist.set_ylim(0, md_top if md_top > 1e-6 else 1.0) if threshold is not None: ax_min_dist.axhline(y=threshold, color="red", linestyle="--", label="Threshold") for i in range(num_robots): @@ -145,8 +172,12 @@ def update(num): return artists ani = FuncAnimation( - fig, update, frames=N, - init_func=init, blit=True, interval=dt * 1000, + fig, + update, + frames=N, + init_func=init, + blit=True, + interval=dt * 1000, ) plt.tight_layout() diff --git a/src/cbfkit/utils/visualizations/plotly_3d_multi_robot.py b/src/cbfkit/utils/visualizations/plotly_3d_multi_robot.py index 8378a3a8..e5433f07 100644 --- a/src/cbfkit/utils/visualizations/plotly_3d_multi_robot.py +++ b/src/cbfkit/utils/visualizations/plotly_3d_multi_robot.py @@ -8,24 +8,48 @@ # Plotly tab10-equivalent colours _PLOTLY_TAB10 = [ - "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", - "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf", + "#1f77b4", + "#ff7f0e", + "#2ca02c", + "#d62728", + "#9467bd", + "#8c564b", + "#e377c2", + "#7f7f7f", + "#bcbd22", + "#17becf", ] def _visualize_3d_plotly( - states, desired_states, desired_state_radius, num_robots, - ellipse_centers, ellipse_radii, ellipse_rotations, - x_lim, y_lim, z_lim, dt, sdim, title, save_animation, - animation_filename, include_min_distance_plot, - include_min_distance_to_obstacles_plot, threshold, - goal_dists, min_dists, obs_dists, + states, + desired_states, + desired_state_radius, + num_robots, + ellipse_centers, + ellipse_radii, + ellipse_rotations, + x_lim, + y_lim, + z_lim, + dt, + sdim, + title, + save_animation, + animation_filename, + include_min_distance_plot, + include_min_distance_to_obstacles_plot, + threshold, + goal_dists, + min_dists, + obs_dists, ): from cbfkit.utils.animators.helpers import ( _compute_plotly_frame_step, _plotly_animation_controls, ) from cbfkit.utils.animators.deps import _require_plotly + _require_plotly() import plotly.graph_objects as go @@ -37,7 +61,10 @@ def _visualize_3d_plotly( # 3D scenes are heavier to render; use a higher floor than 2D frame_indices, frame_duration_ms = _compute_plotly_frame_step( - dt, N, max_frames=200, min_frame_ms=80, + dt, + N, + max_frames=200, + min_frame_ms=80, ) # --- subplot layout --- @@ -55,7 +82,8 @@ def _visualize_3d_plotly( widths = [2] + [1] * (num_cols - 1) fig = make_subplots( - rows=1, cols=num_cols, + rows=1, + cols=num_cols, specs=[specs_row], column_widths=widths, subplot_titles=titles, @@ -64,46 +92,68 @@ def _visualize_3d_plotly( # --- static 3D elements --- for i in range(num_robots): idx = sdim * i - gx, gy, gz = float(desired_states[idx]), float(desired_states[idx + 1]), float(desired_states[idx + 2]) + gx, gy, gz = ( + float(desired_states[idx]), + float(desired_states[idx + 1]), + float(desired_states[idx + 2]), + ) fig.add_trace( go.Scatter3d( - x=[gx], y=[gy], z=[gz], + x=[gx], + y=[gy], + z=[gz], mode="markers", marker=dict(size=8, color=colors[i], symbol="diamond"), name=f"Goal {i + 1}", showlegend=True, ), - row=1, col=1, + row=1, + col=1, ) sx, sy, sz, si, sj, sk = _ellipsoid_mesh( np.array([gx, gy, gz]), np.array([desired_state_radius] * 3), - np.eye(3), n=15, + np.eye(3), + n=15, ) fig.add_trace( go.Mesh3d( - x=sx.tolist(), y=sy.tolist(), z=sz.tolist(), - i=si, j=sj, k=sk, - color=colors[i], opacity=0.15, + x=sx.tolist(), + y=sy.tolist(), + z=sz.tolist(), + i=si, + j=sj, + k=sk, + color=colors[i], + opacity=0.15, showlegend=False, ), - row=1, col=1, + row=1, + col=1, ) # Ellipsoid obstacles if ellipse_centers is not None and ellipse_radii is not None and ellipse_rotations is not None: for ec, er, erot in zip(ellipse_centers, ellipse_radii, ellipse_rotations): mx, my, mz, ii, jj, kk = _ellipsoid_mesh( - np.asarray(ec), np.asarray(er), np.asarray(erot), + np.asarray(ec), + np.asarray(er), + np.asarray(erot), ) fig.add_trace( go.Mesh3d( - x=mx.tolist(), y=my.tolist(), z=mz.tolist(), - i=ii, j=jj, k=kk, - color="black", opacity=0.2, + x=mx.tolist(), + y=my.tolist(), + z=mz.tolist(), + i=ii, + j=jj, + k=kk, + color="black", + opacity=0.2, showlegend=False, ), - row=1, col=1, + row=1, + col=1, ) n_static = len(fig.data) @@ -112,23 +162,28 @@ def _visualize_3d_plotly( for i in range(num_robots): fig.add_trace( go.Scatter3d( - x=[], y=[], z=[], + x=[], + y=[], + z=[], mode="lines", line=dict(color=colors[i], width=3), name=f"Robot {i + 1}", showlegend=True, ), - row=1, col=1, + row=1, + col=1, ) fig.add_trace( go.Scatter( - x=[], y=[], + x=[], + y=[], mode="lines", line=dict(color=colors[i], width=2), name=f"Robot {i + 1}", showlegend=False, ), - row=1, col=2, + row=1, + col=2, ) if include_min_distance_plot: @@ -136,18 +191,23 @@ def _visualize_3d_plotly( for i in range(num_robots): fig.add_trace( go.Scatter( - x=[], y=[], + x=[], + y=[], mode="lines", line=dict(color=colors[i], width=2), name=f"Robot {i + 1}", showlegend=False, ), - row=1, col=col_md, + row=1, + col=col_md, ) if threshold is not None: fig.add_hline( - y=threshold, line_dash="dash", line_color="red", - row=1, col=col_md, + y=threshold, + line_dash="dash", + line_color="red", + row=1, + col=col_md, ) if include_min_distance_to_obstacles_plot: @@ -155,13 +215,15 @@ def _visualize_3d_plotly( for i in range(num_robots): fig.add_trace( go.Scatter( - x=[], y=[], + x=[], + y=[], mode="lines", line=dict(color=colors[i], width=2), name=f"Robot {i + 1}", showlegend=False, ), - row=1, col=col_od, + row=1, + col=col_od, ) n_animated = len(fig.data) - n_static @@ -219,7 +281,9 @@ def _visualize_3d_plotly( # --- scene + axis layout --- menus, sliders = _plotly_animation_controls( - frames, frame_duration_ms, button_y=-0.25, + frames, + frame_duration_ms, + button_y=-0.25, ) fig.update_layout( scene=dict( @@ -243,7 +307,15 @@ def _visualize_3d_plotly( if include_min_distance_plot: fig.update_xaxes(title_text="Time [s]", row=1, col=3) fig.update_yaxes(title_text="Distance [m]", row=1, col=3) - fig.update_yaxes(range=[0, float(np.max(min_dists)) * 1.1], row=1, col=3) + # Include the threshold in the range, else a run that violates the + # constraint throughout renders as a clean plot with the reference + # line clipped off-screen. + md_top = float(np.max(min_dists)) * 1.1 + if not np.isfinite(md_top): + md_top = 0.0 + if threshold is not None and threshold >= 0: + md_top = max(md_top, float(threshold) * 1.15) + fig.update_yaxes(range=[0, md_top if md_top > 1e-6 else 1.0], row=1, col=3) if include_min_distance_to_obstacles_plot: col_od = 3 if not include_min_distance_plot else 4 diff --git a/tests/test_utils/test_visualization_helpers_3d.py b/tests/test_utils/test_visualization_helpers_3d.py new file mode 100644 index 00000000..4e28fbbb --- /dev/null +++ b/tests/test_utils/test_visualization_helpers_3d.py @@ -0,0 +1,102 @@ +"""Tests for the backend-shared 3D distance metrics. + +``_compute_distance_metrics`` feeds the safety panels of every 3D backend +(plotly, matplotlib, manim), so it runs in CI without any optional extra. +""" + +import numpy as np + +from cbfkit.utils.visualizations.helpers_3d import _compute_distance_metrics + + +def _min_dists(states, num_robots, sdim): + """Run the helper and return the per-step minimum over all robots.""" + goals = np.zeros(sdim * num_robots) + _, min_dists, _ = _compute_distance_metrics( + np.asarray(states, dtype=float), + goals, + num_robots, + sdim, + None, + None, + None, + True, + False, + ) + return min_dists.min(axis=1) + + +class TestMinInterRobotDistance: + """Separation must be measured over positions only. + + Regression: the helper used to reshape to the full per-robot state width and + take the norm across every column, reporting sqrt(|dp|^2 + |dv|^2). That + overstates separation precisely when robots converge at speed -- the only + regime the safety panel exists to monitor. + """ + + def test_coincident_robots_report_zero_despite_velocity(self): + # Both robots at the origin -- a collision -- closing at 3 m/s each. + # The velocity columns must not rescue the reported distance. + states = [[0, 0, 0, 3, 0, 0, 0, 0, 0, -3, 0, 0]] + assert _min_dists(states, num_robots=2, sdim=6)[0] == 0.0 + + def test_velocity_columns_do_not_inflate_separation(self): + # Robots 1.0 apart in x, with deliberately large opposing velocities. + states = [[0, 0, 0, 5, -4, 2, 1, 0, 0, -5, 4, -2]] + assert _min_dists(states, num_robots=2, sdim=6)[0] == 1.0 + + def test_matches_position_only_norm_over_a_trajectory(self): + rng = np.random.default_rng(0) + n_steps, num_robots, sdim = 12, 3, 6 + states = rng.normal(0, 5, size=(n_steps, num_robots * sdim)) + + got = _min_dists(states, num_robots, sdim) + + pos = states.reshape(n_steps, num_robots, sdim)[:, :, :3] + expected = np.array( + [ + min( + np.linalg.norm(pos[t, i] - pos[t, j]) + for i in range(num_robots) + for j in range(num_robots) + if i != j + ) + for t in range(n_steps) + ] + ) + assert np.allclose(got, expected) + + def test_position_only_state_layout_still_works(self): + # sdim == 3 has no velocity columns; behaviour must be unchanged. + states = [[0, 0, 0, 2, 0, 0]] + assert _min_dists(states, num_robots=2, sdim=3)[0] == 2.0 + + +class TestGoalAndObstacleDistances: + """These already sliced positions correctly -- pin that they still do.""" + + def test_goal_distance_ignores_velocity_columns(self): + states = np.array([[0.0, 0.0, 0.0, 9.0, 9.0, 9.0]]) + goals = np.array([3.0, 4.0, 0.0, 0.0, 0.0, 0.0]) + goal_dists, _, _ = _compute_distance_metrics( + states, goals, 1, 6, None, None, None, False, False + ) + assert goal_dists[0, 0] == 5.0 + + def test_obstacle_distance_ignores_velocity_columns(self): + states = np.array([[10.0, 0.0, 0.0, 7.0, 7.0, 7.0]]) + goals = np.zeros(6) + _, _, obs_dists = _compute_distance_metrics( + states, + goals, + 1, + 6, + [np.array([0.0, 0.0, 0.0])], + [np.array([2.0, 2.0, 2.0])], + [np.eye(3)], + False, + True, + ) + # Point 10 from centre, sphere radius 2 -> 8 to the surface. + assert obs_dists[0, 0] == 8.0 diff --git a/tests/test_utils/test_visualization_manim.py b/tests/test_utils/test_visualization_manim.py index 6bfb2e2e..4b921239 100644 --- a/tests/test_utils/test_visualization_manim.py +++ b/tests/test_utils/test_visualization_manim.py @@ -5,7 +5,7 @@ import numpy as np import pytest -import cbfkit.utils.animator as animator_module +from cbfkit.utils.animators import deps from cbfkit.utils.visualization import _parse_manim_backend, visualize_3d_multi_robot @@ -29,12 +29,15 @@ class TestManimQualityParsing: def test_bare_manim_defaults_to_low(self): assert _parse_manim_backend("manim") == "low_quality" - @pytest.mark.parametrize("suffix,expected", [ - ("low", "low_quality"), - ("medium", "medium_quality"), - ("high", "high_quality"), - ("production", "production_quality"), - ]) + @pytest.mark.parametrize( + "suffix,expected", + [ + ("low", "low_quality"), + ("medium", "medium_quality"), + ("high", "high_quality"), + ("production", "production_quality"), + ], + ) def test_quality_suffixes(self, suffix, expected): assert _parse_manim_backend(f"manim-{suffix}") == expected @@ -44,8 +47,10 @@ def test_invalid_suffix_raises(self): def test_quality_passed_to_render(self, monkeypatch): """Ensure the quality kwarg reaches render_multi_robot_3d.""" - if not animator_module._HAS_MANIM: - pytest.skip("manim not installed") + # Nothing here needs manim itself: the renderer is mocked and only the + # dispatch path is under test. Satisfy the gate rather than skipping -- + # manim is excluded from the [dev] extra, so a skip means "never in CI". + monkeypatch.setattr(deps, "_HAS_MANIM", True) states, goals = _make_synthetic_data() captured = {} @@ -71,7 +76,9 @@ def mock_render(**kwargs): class TestManimBackendDispatch: def test_manim_backend_raises_import_error_when_missing(self, monkeypatch): """Without manim installed, backend='manim' should raise ImportError.""" - monkeypatch.setattr(animator_module, "_HAS_MANIM", False) + # deps._HAS_MANIM is what _require_manim() consults; patching the + # re-export on cbfkit.utils.animator does not affect the gate. + monkeypatch.setattr(deps, "_HAS_MANIM", False) states, goals = _make_synthetic_data() with pytest.raises(ImportError, match=r"cbfkit\[manim\]"): visualize_3d_multi_robot( @@ -82,18 +89,25 @@ def test_manim_backend_raises_import_error_when_missing(self, monkeypatch): backend="manim", ) - def test_manim_backend_warns_on_subplot_features(self, monkeypatch): - """Subplot features should emit warnings when using manim backend.""" - # Skip if manim is not installed - if not animator_module._HAS_MANIM: - pytest.skip("manim not installed") + def test_manim_backend_forwards_subplot_data(self, monkeypatch): + """The manim backend renders the distance panels rather than dropping them. + + It previously warned that ``include_min_distance_plot`` "will be + ignored" while forwarding the data and drawing the panel anyway, so the + warning told users the opposite of what happened. + """ + # Only the dispatch path is under test and the renderer is mocked, so + # satisfy the gate rather than skipping (manim is not in the [dev] + # extra, so skipping here would mean this never runs in CI). + monkeypatch.setattr(deps, "_HAS_MANIM", True) states, goals = _make_synthetic_data() # We can't actually render without a display, so mock render_multi_robot_3d - import cbfkit.utils.visualization as vis_module + recorded = {} def mock_render(**kwargs): + recorded.update(kwargs) return "/tmp/mock_output.mp4" monkeypatch.setattr( @@ -111,9 +125,18 @@ def mock_render(**kwargs): backend="manim", include_min_distance_plot=True, include_min_distance_to_obstacles_plot=True, + threshold=0.5, + safety_radius=0.25, ellipse_centers=[np.array([0.0, 0.0, 0.0])], ellipse_radii=[np.array([1.0, 1.0, 1.0])], ellipse_rotations=[np.eye(3)], ) - manim_warnings = [x for x in w if "Manim backend" in str(x.message)] - assert len(manim_warnings) == 2 + # Stronger than matching the old text: the path should be silent. + assert [str(x.message) for x in w] == [] + + # The panels are actually drawn, so the data must reach the renderer. + assert recorded["min_dists"] is not None + assert recorded["obs_dists"] is not None + # ...along with the reference line and bubble size that make them readable. + assert recorded["threshold"] == 0.5 + assert recorded["safety_radius"] == 0.25 diff --git a/tutorials/multi_robot_3d_reachavoid.py b/tutorials/multi_robot_3d_reachavoid.py index 3c4b6816..bafed47f 100644 --- a/tutorials/multi_robot_3d_reachavoid.py +++ b/tutorials/multi_robot_3d_reachavoid.py @@ -42,7 +42,8 @@ ACTUATION_LIMITS = jnp.full((CONTROL_DIM,), 10) # Barrier Parameters -D_MIN_SQUARED = 0.1 # Minimum distance squared between robots +D_MIN_SQUARED = 0.25 # Minimum separation squared between robots +D_MIN = float(np.sqrt(D_MIN_SQUARED)) # => robots must stay 0.5 apart OBS_RADIUS = 8.0 # Obstacle sphere radius OBS_RADIUS_SQUARED = OBS_RADIUS**2 # Obstacle radius squared for barrier @@ -56,18 +57,22 @@ INITIAL_STATE = np.zeros(STATE_DIM) goals = np.zeros(STATE_DIM) +# Seeded so the scenario -- and therefore the safety margins it demonstrates -- +# is the same on every run. +rng = np.random.default_rng(0) + radius = 15 for i in range(NUM_ROBOTS): angle = 2 * np.pi * i / NUM_ROBOTS idx = STATE_DIM_PER_ROBOT * i # Initial positions (distributed in a circle with some noise) - INITIAL_STATE[idx] = radius * np.cos(angle) + np.random.normal(0.5, 1) # x - INITIAL_STATE[idx + 1] = radius * np.sin(angle) + np.random.normal(0.5, 1) # y - INITIAL_STATE[idx + 2] = np.random.normal(-3, 3) # z + INITIAL_STATE[idx] = radius * np.cos(angle) + rng.normal(0.5, 1) # x + INITIAL_STATE[idx + 1] = radius * np.sin(angle) + rng.normal(0.5, 1) # y + INITIAL_STATE[idx + 2] = rng.normal(-3, 3) # z # Initial velocities - INITIAL_STATE[idx + 3 : idx + 6] = np.random.normal(0.05, 0.1) # vx, vy, vz + INITIAL_STATE[idx + 3 : idx + 6] = rng.normal(0.05, 0.1, size=3) # vx, vy, vz # Goals (opposite side of the circle) goals[idx : idx + 3] = -INITIAL_STATE[idx : idx + 3] # x_goal, y_goal, z_goal @@ -145,7 +150,7 @@ def terminal_cost(x: jnp.ndarray, action: jnp.ndarray) -> float: state_constraint_funcs = [ f"(x[{STATE_DIM_PER_ROBOT * i}] - x[{STATE_DIM_PER_ROBOT * j}])**2 + " f"(x[{STATE_DIM_PER_ROBOT * i + 1}] - x[{STATE_DIM_PER_ROBOT * j + 1}])**2 + " - f"(x[{STATE_DIM_PER_ROBOT * i + 2}] - x[{STATE_DIM_PER_ROBOT * j + 2}])**2 - 0.25 " + f"(x[{STATE_DIM_PER_ROBOT * i + 2}] - x[{STATE_DIM_PER_ROBOT * j + 2}])**2 - {D_MIN_SQUARED} " for i in range(NUM_ROBOTS) for j in range(i + 1, NUM_ROBOTS) ] @@ -288,7 +293,7 @@ def terminal_cost(x: jnp.ndarray, action: jnp.ndarray) -> float: os.path.join(TARGET_DIRECTORY, MODEL_NAME, f"animation_{NUM_ROBOTS}_robots.gif") ) - visualize_3d_multi_robot( + saved_to = visualize_3d_multi_robot( states=results.states, desired_states=goals, desired_state_radius=0.3, @@ -299,10 +304,16 @@ def terminal_cost(x: jnp.ndarray, action: jnp.ndarray) -> float: save_animation=True, animation_filename=animation_path, include_min_distance_plot=True, + # Same constant the QP enforces, so the panel reads as satisfied / + # violated and the drawn bubbles touch exactly at the constraint. + threshold=D_MIN, + safety_radius=D_MIN / 2, ellipse_centers=obstacle_centers, ellipse_radii=obstacle_radii, ellipse_rotations=obstacle_rotations, backend="manim-low", # Options: manim-low, manim-medium, manim-high, manim-production ) - print(f"\nAnimation saved to: file://{animation_path}") + # Report what the backend actually wrote, not what we asked for -- these + # differed silently before, leaving a stale file at the advertised path. + print(f"\nAnimation saved to: file://{saved_to}")