utils/logging_utils.py

from typing import Optional
import wandb
import numpy as np
import torch
import os

import matplotlib.pyplot as plt
import cv2
import matplotlib.pyplot as plt
from tqdm import trange, tqdm
import matplotlib.animation as animation
from pathlib import Path
import imageio
        
plt.set_loglevel("warning")

from torchmetrics.functional import mean_squared_error, peak_signal_noise_ratio
from torchmetrics.functional import (
    structural_similarity_index_measure,
    universal_image_quality_index,
)
from algorithms.common.metrics import (
    FrechetVideoDistance,
    LearnedPerceptualImagePatchSimilarity,
    FrechetInceptionDistance,
)


# FIXME: clean up & check this util
def log_video(
    observation_hat,
    observation_gt=None,
    step=0,
    namespace="train",
    prefix="video",
    context_frames=0,
    color=(255, 0, 0),
    logger=None,
    fps=15,
    format="mp4",
    save_local=True,
    local_save_dir=None,
):
    """
    take in video tensors in range [-1, 1] and log into wandb

    :param observation_hat: predicted observation tensor of shape (frame, batch, channel, height, width)
    :param observation_gt: ground-truth observation tensor of shape (frame, batch, channel, height, width)
    :param step: an int indicating the step number
    :param namespace: a string specify a name space this video logging falls under, e.g. train, val
    :param prefix: a string specify a prefix for the video name
    :param context_frames: an int indicating how many frames in observation_hat are ground truth given as context
    :param color: a tuple of 3 numbers specifying the color of the border for ground truth frames
    :param logger: optional logger to use. use global wandb if not specified
    :param fps: frames per second for the video (default: 15)
    :param format: video format, either "mp4" or "gif" (default: "mp4")
    :param save_local: whether to save videos to local disk (default: True)
    :param local_save_dir: directory to save local videos. If None, uses hydra output dir
    """
    import cv2
    import hydra
    from pathlib import Path
    
    # Get local rank for distributed training
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    
    if not logger:
        logger = wandb
    
    n_samples = int(observation_hat.shape[1])

    def _to_uint8_video(video_tensor):
        return (
            video_tensor.detach()
            .clamp(0.0, 1.0)
            .mul(255)
            .to(torch.uint8)
            .cpu()
            .numpy()
        )

    is_rank_zero = local_rank == 0
    save_local = bool(save_local and is_rank_zero)

    # Setup local save directory
    if save_local:
        if local_save_dir is None:
            try:
                hydra_cfg = hydra.core.hydra_config.HydraConfig.get()
                output_dir = Path(hydra_cfg.runtime.output_dir)
            except Exception:
                output_dir = Path.cwd() / "outputs"
            local_save_dir = output_dir / "videos" / namespace
        else:
            local_save_dir = Path(local_save_dir)

        local_save_dir.mkdir(parents=True, exist_ok=True)

        # Save pred videos locally
        pred_dir = local_save_dir / "pred"
        pred_dir.mkdir(parents=True, exist_ok=True)

        # Save gt videos locally if available
        if observation_gt is not None:
            gt_dir = local_save_dir / "gt"
            gt_dir.mkdir(parents=True, exist_ok=True)

    # Save videos
    for i in range(n_samples):
        if not save_local and not (is_rank_zero and logger):
            continue

        video_pred = _to_uint8_video(observation_hat[:, i])  # (T, C, H, W)
        video_gt = None
        if observation_gt is not None:
            video_gt = _to_uint8_video(observation_gt[:, i])

        if save_local:
            # Save prediction video
            if step is not None:
                video_filename_pred = f"{prefix}_{i}_rank{local_rank}_step{step}.{format}"
            else:
                video_filename_pred = f"{prefix}_{i}_rank{local_rank}.{format}"

            video_path_pred = pred_dir / video_filename_pred
            _save_video_to_file(video_pred, str(video_path_pred), fps)

            # Save ground truth video if available
            if video_gt is not None:
                if step is not None:
                    video_filename_gt = f"{prefix}_{i}_rank{local_rank}_step{step}.{format}"
                else:
                    video_filename_gt = f"{prefix}_{i}_rank{local_rank}.{format}"

                video_path_gt = gt_dir / video_filename_gt
                _save_video_to_file(video_gt, str(video_path_gt), fps)

        # Log to wandb (only rank 0 to avoid duplicate logging)
        if is_rank_zero and logger:
            if video_gt is not None:
                video_combined = np.concatenate([video_pred, video_gt], axis=-1)
                logger.log(
                    {
                        f"{namespace}/{prefix}_{i}": wandb.Video(video_combined, fps=fps, format=format),
                        f"trainer/global_step": step,
                    }
                )
            else:
                logger.log(
                    {
                        f"{namespace}/{prefix}_{i}": wandb.Video(video_pred, fps=fps, format=format),
                        f"trainer/global_step": step,
                    }
                )


def _save_video_to_file(video_tensor, output_path, fps=15):
    """
    Save a video tensor to file using imageio (better compatibility than cv2).
    
    :param video_tensor: numpy array of shape (T, C, H, W) with values in [0, 255]
    :param output_path: path to save the video
    :param fps: frames per second
    """

    T, C, H, W = video_tensor.shape
    
    # Convert from (T, C, H, W) to (T, H, W, C)
    video_tensor = np.transpose(video_tensor, (0, 2, 3, 1))
    
    # Ensure uint8
    video_tensor = video_tensor.astype(np.uint8)
    
    # Save using imageio with H.264 codec (best compatibility)
    writer = imageio.get_writer(
        output_path, 
        fps=fps,
        codec='libx264',  # H.264 codec - widely supported
        quality=8,  # Good quality (scale 0-10, 10 is best)
        pixelformat='yuv420p',  # Standard pixel format for compatibility
        macro_block_size=1  # Better quality
    )
    
    for frame in video_tensor:
        writer.append_data(frame)
    
    writer.close()
        



def get_validation_metrics_for_videos(
    observation_hat,
    observation_gt,
    lpips_model: Optional[LearnedPerceptualImagePatchSimilarity] = None,
    fid_model: Optional[FrechetInceptionDistance] = None,
    fvd_model: Optional[FrechetVideoDistance] = None,
    lpips_batch_size: int = 100,
):
    """
    :param observation_hat: predicted observation tensor of shape (frame, batch, channel, height, width)
    :param observation_gt: ground-truth observation tensor of shape (frame, batch, channel, height, width)
    :param lpips_model: a LearnedPerceptualImagePatchSimilarity object from algorithm.common.metrics
    :param fid_model: a FrechetInceptionDistance object  from algorithm.common.metrics
    :param fvd_model: a FrechetVideoDistance object  from algorithm.common.metrics
    :param lpips_batch_size: batch size for LPIPS calculation to avoid OOM (default: 100)
    :return: a tuple of metrics
    """
    frame, batch, channel, height, width = observation_hat.shape
    output_dict = {}
    observation_gt = observation_gt.type_as(observation_hat)  # some metrics don't fully support fp16

    if frame < 9:
        fvd_model = None  # FVD requires at least 9 frames

    observation_hat = observation_hat.float()
    observation_gt = observation_gt.float()

    # Clip to [0, 1] range before computing metrics (matching video saving behavior)
    observation_hat_clipped = torch.clamp(observation_hat, 0.0, 1.0)
    observation_gt_clipped = torch.clamp(observation_gt, 0.0, 1.0)

    # Compute video-wise PSNR: frame-wise average per video, then average across videos
    video_psnr_list = []
    for b in range(batch):
        frame_psnr_for_video = []
        for f in range(frame):
            frame_psnr = peak_signal_noise_ratio(observation_hat_clipped[f, b], observation_gt_clipped[f, b], data_range=1.0)
            frame_psnr_for_video.append(frame_psnr)
        video_psnr = torch.stack(frame_psnr_for_video).mean()
        video_psnr_list.append(video_psnr)
    output_dict["psnr"] = torch.stack(video_psnr_list).mean()
    
    observation_hat_clipped = observation_hat_clipped.view(-1, channel, height, width)
    observation_gt_clipped = observation_gt_clipped.view(-1, channel, height, width)

    # Compute MSE on clipped data
    output_dict["mse"] = mean_squared_error(observation_hat_clipped, observation_gt_clipped)
    # output_dict["ssim"] = structural_similarity_index_measure(observation_hat_clipped, observation_gt_clipped, data_range=1.0)
    # output_dict["uiqi"] = universal_image_quality_index(observation_hat_clipped, observation_gt_clipped)

    # LPIPS computation
    if lpips_model is not None:
        # Process LPIPS in batches to avoid OOM
        num_frames = observation_hat_clipped.shape[0]
        
        for i in range(0, num_frames, lpips_batch_size):
            batch_end = min(i + lpips_batch_size, num_frames)
            observation_hat_batch = observation_hat_clipped[i:batch_end]
            observation_gt_batch = observation_gt_clipped[i:batch_end]
            
            lpips_model.update(observation_hat_batch, observation_gt_batch)
            
            # Free GPU memory after each batch
            del observation_hat_batch, observation_gt_batch
            torch.cuda.empty_cache()
        
        lpips = lpips_model.compute().item()
        # Reset the states of non-functional metrics
        output_dict["lpips"] = lpips
        lpips_model.reset()

    # FID computation
    if fid_model is not None:
        observation_hat_uint8 = (observation_hat_clipped * 255).type(torch.uint8)
        observation_gt_uint8 = (observation_gt_clipped * 255).type(torch.uint8)
        fid_model.update(observation_gt_uint8, real=True)
        fid_model.update(observation_hat_uint8, real=False)
        fid = fid_model.compute()
        output_dict["fid"] = fid
        # Reset the states of non-functional metrics
        fid_model.reset()

    return output_dict


def is_grid_env(env_id):
    return "maze2d" in env_id or "diagonal2d" in env_id


def get_maze_grid(env_id):
    # import gym
    # maze_string = gym.make(env_id).str_maze_spec
    if "large" in env_id:
        maze_string = "############\\#OOOO#OOOOO#\\#O##O#O#O#O#\\#OOOOOO#OOO#\\#O####O###O#\\#OO#O#OOOOO#\\##O#O#O#O###\\#OO#OOO#OGO#\\############"
    if "medium" in env_id:
        maze_string = "########\\#OO##OO#\\#OO#OOO#\\##OOO###\\#OO#OOO#\\#O#OO#O#\\#OOO#OG#\\########"
    if "umaze" in env_id:
        maze_string = "#####\\#GOO#\\###O#\\#OOO#\\#####"
    lines = maze_string.split("\\")
    grid = [line[1:-1] for line in lines]
    return grid[1:-1]


def get_random_start_goal(env_id, batch_size):
    maze_grid = get_maze_grid(env_id)
    s2i = {"O": 0, "#": 1, "G": 2}
    maze_grid = [[s2i[s] for s in r] for r in maze_grid]
    maze_grid = np.array(maze_grid)
    x, y = np.nonzero(maze_grid == 0)
    indices = np.random.randint(len(x), size=batch_size)
    start = np.stack([x[indices], y[indices]], -1) + 1
    x, y = np.nonzero(maze_grid == 2)
    goal = np.concatenate([x, y], -1)
    goal = np.tile(goal[None, :], (batch_size, 1)) + 1
    return start, goal


def plot_maze_layout(ax, maze_grid):
    ax.clear()

    if maze_grid is not None:
        for i, row in enumerate(maze_grid):
            for j, cell in enumerate(row):
                if cell == "#":
                    square = plt.Rectangle((i + 0.5, j + 0.5), 1, 1, edgecolor="black", facecolor="black")
                    ax.add_patch(square)

    ax.set_aspect("equal")
    ax.grid(True, color="white", linewidth=4)
    ax.set_axisbelow(True)
    ax.spines["top"].set_linewidth(4)
    ax.spines["right"].set_linewidth(4)
    ax.spines["bottom"].set_linewidth(4)
    ax.spines["left"].set_linewidth(4)
    ax.set_facecolor("lightgray")
    ax.tick_params(
        axis="both",
        which="both",
        bottom=False,
        top=False,
        left=False,
        right=False,
        labelbottom=False,
        labelleft=False,
    )
    ax.set_xticks(np.arange(0.5, len(maze_grid) + 0.5))
    ax.set_yticks(np.arange(0.5, len(maze_grid[0]) + 0.5))
    ax.set_xlim(0.5, len(maze_grid) + 0.5)
    ax.set_ylim(0.5, len(maze_grid[0]) + 0.5)
    ax.grid(True, color="white", which="minor", linewidth=4)


def plot_start_goal(ax, start_goal: None):
    def draw_star(center, radius, num_points=5, color="black"):
        angles = np.linspace(0.0, 2 * np.pi, num_points, endpoint=False) + 5 * np.pi / (2 * num_points)
        inner_radius = radius / 2.0

        points = []
        for angle in angles:
            points.extend(
                [
                    center[0] + radius * np.cos(angle),
                    center[1] + radius * np.sin(angle),
                    center[0] + inner_radius * np.cos(angle + np.pi / num_points),
                    center[1] + inner_radius * np.sin(angle + np.pi / num_points),
                ]
            )

        star = plt.Polygon(np.array(points).reshape(-1, 2), color=color)
        ax.add_patch(star)

    start_x, start_y = start_goal[0]
    start_outer_circle = plt.Circle((start_x, start_y), 0.16, facecolor="white", edgecolor="black")
    ax.add_patch(start_outer_circle)
    start_inner_circle = plt.Circle((start_x, start_y), 0.08, color="black")
    ax.add_patch(start_inner_circle)

    goal_x, goal_y = start_goal[1]
    goal_outer_circle = plt.Circle((goal_x, goal_y), 0.16, facecolor="white", edgecolor="black")
    ax.add_patch(goal_outer_circle)
    draw_star((goal_x, goal_y), radius=0.08)


def make_trajectory_images(env_id, trajectory, batch_size, start, goal, plot_end_points=True):
    images = []
    for batch_idx in range(batch_size):
        fig, ax = plt.subplots()
        if is_grid_env(env_id):
            maze_grid = get_maze_grid(env_id)
        else:
            maze_grid = None
        plot_maze_layout(ax, maze_grid)
        ax.scatter(trajectory[:, batch_idx, 0], trajectory[:, batch_idx, 1], c=np.arange(len(trajectory)), cmap="Reds"),
        if plot_end_points:
            start_goal = (start[batch_idx], goal[batch_idx])
            plot_start_goal(ax, start_goal)
        # plt.title(f"sample_{batch_idx}")
        fig.tight_layout()
        fig.canvas.draw()
        img_shape = fig.canvas.get_width_height()[::-1] + (4,)
        img = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8).copy().reshape(img_shape)
        images.append(img)

        plt.close()
    return images


def make_convergence_animation(
    env_id,
    plan_history,
    trajectory,
    start,
    goal,
    open_loop_horizon,
    namespace,
    interval=100,
    plot_end_points=True,
    batch_idx=0,
):
    # - plan_history: contains for each time step all the MPC predicted plans for each pyramid noise level.
    #                 Structured as a list of length (episode_len // open_loop_horizon), where each
    #                 element corresponds to a control_time_step and stores a list of length pyramid_height,
    #                 where each element is a plan at a different pyramid noise level and stored as a tensor of
    #                 shape (episode_len // open_loop_horizon - control_time_step,
    #                        batch_size, x_stacked_shape)

    # select index and prune history
    start, goal = start[batch_idx], goal[batch_idx]
    trajectory = trajectory[:, batch_idx]
    plan_history = [[pm[:, batch_idx] for pm in pt] for pt in plan_history]
    trajectory, plan_history = prune_history(plan_history, trajectory, goal, open_loop_horizon)

    # animate the convergence of the first plan
    fig, ax = plt.subplots()
    if "large" in env_id:
        fig.set_size_inches(3.5, 5)
    else:
        fig.set_size_inches(3, 3)
    ax.set_axis_off()
    fig.subplots_adjust(left=0, bottom=0, right=1, top=1)

    if is_grid_env(env_id):
        maze_grid = get_maze_grid(env_id)
    else:
        maze_grid = None

    def update(frame):
        plot_maze_layout(ax, maze_grid)

        plan_history_m = plan_history[0][frame]
        plan_history_m = plan_history_m.numpy()
        ax.scatter(
            plan_history_m[:, 0],
            plan_history_m[:, 1],
            c=np.arange(len(plan_history_m))[::-1],
            cmap="Reds",
        )

        if plot_end_points:
            plot_start_goal(ax, (start, goal))

    frames = tqdm(range(len(plan_history[0])), desc="Making convergence animation")
    ani = animation.FuncAnimation(fig, update, frames=frames, interval=interval)
    prefix = wandb.run.id if wandb.run is not None else env_id
    filename = f"/tmp/{prefix}_{namespace}_convergence.mp4"
    ani.save(filename, writer="ffmpeg", fps=5)
    return filename


def prune_history(plan_history, trajectory, goal, open_loop_horizon):
    dist = np.linalg.norm(
        trajectory[:, :2] - np.array(goal)[None],
        axis=-1,
    )
    reached = dist < 0.2
    if reached.any():
        cap_idx = np.argmax(reached)
        trajectory = trajectory[: cap_idx + open_loop_horizon + 1]
        plan_history = plan_history[: cap_idx // open_loop_horizon + 2]

    pruned_plan_history = []
    for plans in plan_history:
        pruned_plan_history.append([])
        for m in range(len(plans)):
            plan = plans[m]
            pruned_plan_history[-1].append(plan)
        plan = pruned_plan_history[-1][-1]
        dist = np.linalg.norm(plan.numpy()[:, :2] - np.array(goal)[None], axis=-1)
        reached = dist < 0.2
        if reached.any():
            cap_idx = np.argmax(reached) + 1
            pruned_plan_history[-1] = [p[:cap_idx] for p in pruned_plan_history[-1]]
    return trajectory, pruned_plan_history


def make_mpc_animation(
    env_id,
    plan_history,
    trajectory,
    start,
    goal,
    open_loop_horizon,
    namespace,
    interval=100,
    plot_end_points=True,
    batch_idx=0,
):
    # - plan_history: contains for each time step all the MPC predicted plans for each pyramid noise level.
    #                 Structured as a list of length (episode_len // open_loop_horizon), where each
    #                 element corresponds to a control_time_step and stores a list of length pyramid_height,
    #                 where each element is a plan at a different pyramid noise level and stored as a tensor of
    #                 shape (episode_len // open_loop_horizon - control_time_step,
    #                        batch_size, x_stacked_shape)

    # select index and prune history
    start, goal = start[batch_idx], goal[batch_idx]
    trajectory = trajectory[:, batch_idx]
    plan_history = [[pm[:, batch_idx] for pm in pt] for pt in plan_history]
    trajectory, plan_history = prune_history(plan_history, trajectory, goal, open_loop_horizon)

    # animate the convergence of the plans
    fig, ax = plt.subplots()
    if "large" in env_id:
        fig.set_size_inches(3.5, 5)
    else:
        fig.set_size_inches(3, 3)
    ax.set_axis_off()
    fig.subplots_adjust(left=0, bottom=0, right=1, top=1)
    trajectory_colors = np.linspace(0, 1, len(trajectory))

    if is_grid_env(env_id):
        maze_grid = get_maze_grid(env_id)
    else:
        maze_grid = None

    def update(frame):
        control_time_step = 0
        while frame >= 0:
            frame -= len(plan_history[control_time_step])
            control_time_step += 1
        control_time_step -= 1
        m = frame + len(plan_history[control_time_step])
        num_steps_taken = 1 + open_loop_horizon * control_time_step
        plot_maze_layout(ax, maze_grid)

        plan_history_m = plan_history[control_time_step][m]
        plan_history_m = plan_history_m.numpy()
        ax.scatter(
            trajectory[:num_steps_taken, 0],
            trajectory[:num_steps_taken, 1],
            c=trajectory_colors[:num_steps_taken],
            cmap="Blues",
        )
        ax.scatter(
            plan_history_m[:, 0],
            plan_history_m[:, 1],
            c=np.arange(len(plan_history_m))[::-1],
            cmap="Reds",
        )

        if plot_end_points:
            plot_start_goal(ax, (start, goal))

    num_frames = sum([len(p) for p in plan_history])
    frames = tqdm(range(num_frames), desc="Making MPC animation")
    ani = animation.FuncAnimation(fig, update, frames=frames, interval=interval)
    prefix = wandb.run.id if wandb.run is not None else env_id
    filename = f"/tmp/{prefix}_{namespace}_mpc.mp4"
    ani.save(filename, writer="ffmpeg", fps=5)

    return filename