轨迹预测指标minADE与minFDE计算

minADE表示最小平均举例误差，minFDE表示最小最终距离误差。

import torch

batch_size, num_preds, num_point, num_dim = 4, 6, 16, 2

pred_trajs = torch.randn(batch_size, num_preds, num_point, num_dim)
gt_traj = torch.randn(batch_size, num_point, num_dim)
valid_mask = torch.ones(batch_size, num_point, dtype=torch.bool)

# minADE
min_ade = ((torch.norm(pred_trajs - gt_traj.unsqueeze(1), p=2, dim=-1) * 
        valid_mask.unsqueeze(1)).sum(dim=-1).min(dim=-1)[0] / valid_mask.sum(dim=-1)).mean()

# minFDE（似乎与 waymo 原始论文中的定义存在差别）
inds_last = (valid_mask * torch.arange(1, valid_mask.size(-1) + 1)).argmax(dim=-1)  # 最后一个有效点的索引
inds_best = torch.norm(
    pred_trajs[torch.arange(batch_size), :, inds_last] -
    gt_traj[torch.arange(batch_size), inds_last].unsqueeze(-2), p=2, dim=-1).argmin(dim=-1)  # 最后一个点最接近 GT 的轨迹索引
min_fde = ((torch.norm(pred_trajs[torch.arange(batch_size), inds_best] - gt_traj, p=2, dim=-1) * 
        valid_mask).sum(dim=-1) / valid_mask.sum(dim=-1)).mean()

print(min_ade)
print(min_fde)