visualization

visualization

Visualization utilities for synthetic benchmark trajectories.

plot_trajectories_2d

plot_trajectories_2d(trajectories, task_name, output_path=None, mode_ids=None, expert_trajectories=None, expert_mode_ids=None, title=None, num_modes=None, num_styles=None, noise_std=None)

Plot 2D trajectories overlaid on the task layout.

Always returns the figure. Additionally saves to disk when output_path is provided.

Parameters:

Name	Type	Description	Default
trajectories	ndarray	Cartesian trajectories, shape (num_trajectories, num_timesteps, 2), values in [0, 1].	required
task_name	str	SyntheticTaskName.value string.	required
output_path	str | None	Optional PNG path. Saves to disk when provided.	None
mode_ids	ndarray | None	Optional per-trajectory mode index for coloring. Shape (num_trajectories,).	None
expert_trajectories	ndarray | None	Optional faint background trajectories. Shape (num_experts, num_timesteps, 2).	None
expert_mode_ids	ndarray | None	Optional per-expert mode index for coloring.	None
title	str | None	Optional plot title. Defaults to the task display name.	None
num_modes	int | None	Number of modes for variable-mode tasks (radial, corridor_navigation). Forwarded to get_task_layout.	None
num_styles	int | None	Number of styles per corridor for corridor_navigation.	None
noise_std	float | None	Trajectory noise std. Used to size radial obstacles.	None

Returns:

Type	Description
Figure	The matplotlib Figure. Caller is responsible for closing it.

Source code in src/versatil/data/synthetic/visualization.py

def plot_trajectories_2d(
    trajectories: np.ndarray,
    task_name: str,
    output_path: str | None = None,
    mode_ids: np.ndarray | None = None,
    expert_trajectories: np.ndarray | None = None,
    expert_mode_ids: np.ndarray | None = None,
    title: str | None = None,
    num_modes: int | None = None,
    num_styles: int | None = None,
    noise_std: float | None = None,
) -> plt.Figure:
    """Plot 2D trajectories overlaid on the task layout.

    Always returns the figure. Additionally saves to disk when
    ``output_path`` is provided.

    Args:
        trajectories: Cartesian trajectories, shape (num_trajectories,
            num_timesteps, 2), values in [0, 1].
        task_name: SyntheticTaskName.value string.
        output_path: Optional PNG path. Saves to disk when provided.
        mode_ids: Optional per-trajectory mode index for coloring.
            Shape (num_trajectories,).
        expert_trajectories: Optional faint background trajectories.
            Shape (num_experts, num_timesteps, 2).
        expert_mode_ids: Optional per-expert mode index for coloring.
        title: Optional plot title. Defaults to the task display name.
        num_modes: Number of modes for variable-mode tasks (radial,
            corridor_navigation). Forwarded to ``get_task_layout``.
        num_styles: Number of styles per corridor for corridor_navigation.
        noise_std: Trajectory noise std. Used to size radial obstacles.

    Returns:
        The matplotlib Figure. Caller is responsible for closing it.
    """
    _apply_plot_theme()
    layout_kwargs: dict = {"task_name": task_name}
    if num_modes is not None:
        layout_kwargs["num_modes"] = num_modes
    if num_styles is not None:
        layout_kwargs["num_styles"] = num_styles
    if noise_std is not None:
        layout_kwargs["noise_std"] = noise_std
    layout = get_task_layout(**layout_kwargs)
    figure, axes = plt.subplots(figsize=(6, 6), dpi=150)
    _draw_task_background(axes=axes, layout=layout)

    if expert_trajectories is not None:
        for index in range(len(expert_trajectories)):
            mode_index = (
                int(expert_mode_ids[index]) if expert_mode_ids is not None else 0
            )
            color = PLOT_MODE_COLORS[mode_index % len(PLOT_MODE_COLORS)]
            axes.plot(
                expert_trajectories[index, :, 0],
                expert_trajectories[index, :, 1],
                color=color,
                alpha=PLOT_EXPERT_ALPHA,
                linewidth=PLOT_EXPERT_LINEWIDTH,
                zorder=3,
            )

    for index in range(len(trajectories)):
        mode_index = int(mode_ids[index]) if mode_ids is not None else 0
        color = PLOT_MODE_COLORS[mode_index % len(PLOT_MODE_COLORS)]
        axes.plot(
            trajectories[index, :, 0],
            trajectories[index, :, 1],
            color=color,
            alpha=PLOT_TRAJECTORY_ALPHA,
            linewidth=PLOT_TRAJECTORY_LINEWIDTH,
            zorder=4,
        )

    if len(trajectories) > 0:
        final_positions = trajectories[:, -1, :]
        axes.scatter(
            final_positions[:, 0],
            final_positions[:, 1],
            s=PLOT_AGENT_MARKER_SIZE,
            color=PLOT_AGENT_MARKER_COLOR,
            edgecolor="white",
            linewidth=0.6,
            zorder=5,
        )

    axes.legend(
        handles=_build_legend_handles(layout=layout),
        loc="lower right",
        frameon=True,
        framealpha=0.9,
        fontsize=9,
        handlelength=1.8,
        edgecolor=PLOT_BORDER_COLOR,
    )
    plot_title = title if title is not None else TASK_DISPLAY_NAMES[task_name]
    axes.set_title(plot_title, pad=12)
    plt.tight_layout()
    if output_path is not None:
        plt.savefig(
            output_path,
            dpi=150,
            bbox_inches="tight",
            facecolor=PLOT_BACKGROUND_COLOR,
        )
    return figure

save_rollouts_gif

save_rollouts_gif(trajectories, task_name, output_path, mode_ids=None, image_size=DEFAULT_IMAGE_SIZE, frames_per_second=30)

Save an animated GIF showing all rollout trajectories evolving in parallel.

At each frame, every rollout's current position is drawn on the same canvas along with the trail up to that timestep. Colors encode mode_id.

Parameters:

Name	Type	Description	Default
trajectories	ndarray	Rollout trajectories, shape (num_rollouts, num_timesteps, 2), values in [0, 1].	required
task_name	str	SyntheticTaskName.value string.	required
output_path	str	Destination GIF path.	required
mode_ids	ndarray | None	Optional per-rollout mode index for agent coloring. Shape (num_rollouts,).	None
image_size	int	Side length of each rendered frame in pixels.	DEFAULT_IMAGE_SIZE
frames_per_second	int	GIF playback rate.	30

Source code in src/versatil/data/synthetic/visualization.py

def save_rollouts_gif(
    trajectories: np.ndarray,
    task_name: str,
    output_path: str,
    mode_ids: np.ndarray | None = None,
    image_size: int = DEFAULT_IMAGE_SIZE,
    frames_per_second: int = 30,
) -> None:
    """Save an animated GIF showing all rollout trajectories evolving in parallel.

    At each frame, every rollout's current position is drawn on the same
    canvas along with the trail up to that timestep. Colors encode mode_id.

    Args:
        trajectories: Rollout trajectories, shape (num_rollouts,
            num_timesteps, 2), values in [0, 1].
        task_name: SyntheticTaskName.value string.
        output_path: Destination GIF path.
        mode_ids: Optional per-rollout mode index for agent coloring.
            Shape (num_rollouts,).
        image_size: Side length of each rendered frame in pixels.
        frames_per_second: GIF playback rate.
    """
    if task_name not in TASK_DISPLAY_NAMES:
        raise ValueError(f"Unknown synthetic task: {task_name}")
    layout = get_task_layout(task_name=task_name)
    num_timesteps = trajectories.shape[1]
    duration_milliseconds = int(round(1000 / frames_per_second))
    frames: list[Image.Image] = []
    for timestep in range(num_timesteps):
        positions_at_timestep = trajectories[:, timestep, :]
        trails_up_to_timestep = trajectories[:, : timestep + 1, :]
        frame_array = _render_multi_agent_frame(
            positions=positions_at_timestep,
            trails=trails_up_to_timestep,
            mode_ids=mode_ids,
            layout=layout,
            image_size=image_size,
        )
        frames.append(Image.fromarray(frame_array))
    frames[0].save(
        output_path,
        save_all=True,
        append_images=frames[1:],
        duration=duration_milliseconds,
        loop=0,
        optimize=False,
    )