Update Wall-X to 1.1.0 (#104)

wall_x/utils/metrics.py

@@ -0,0 +1,95 @@
+from typing import List
+
+import torch
+
+
+def get_action_accuracy(
+    gt: torch.FloatTensor,  # [Batch_Size, Horizon, Action_Dim]
+    pred: torch.FloatTensor,
+    thresholds: List[float] = [0.1, 0.2],
+) -> torch.FloatTensor:
+    device = gt.device
+    diff = torch.abs(gt - pred).reshape(-1, gt.shape[-1])
+
+    # get the percentage of diff lower than threshold for all action dimensions
+    accuracies = torch.zeros(len(thresholds), device=device)
+    for idx, threshold in enumerate(thresholds):
+        accuracy = torch.mean(
+            (torch.mean((diff < threshold).float(), dim=1) >= 1.0).float()
+        )
+        accuracies[idx] = accuracy
+    return accuracies
+
+
+def dtw_distance(seq1, seq2):
+    """
+    Compute the Dynamic Time Warping distance between two sequences.
+
+    ``seq1`` and ``seq2`` must have shape ``(T, D)``, where ``T`` is the
+    number of time steps and ``D`` is the feature dimension.
+    """
+    n, m = seq1.shape[0], seq2.shape[0]
+
+    seq1_d, seq2_d = seq1.double(), seq2.double()
+    dtw_matrix = torch.full(
+        (n + 1, m + 1), float("inf"), dtype=torch.float64, device=seq1.device
+    )
+    dtw_matrix[0, 0] = 0
+
+    for i in range(1, n + 1):
+        for j in range(1, m + 1):
+            cost = torch.dist(seq1_d[i - 1], seq2_d[j - 1])
+            dtw_matrix[i, j] = cost + torch.min(
+                torch.stack(
+                    [
+                        dtw_matrix[i - 1, j],
+                        dtw_matrix[i, j - 1],
+                        dtw_matrix[i - 1, j - 1],
+                    ]
+                )
+            )
+
+    return dtw_matrix[n, m]
+
+
+def frechet_distance(seq1, seq2):
+    """
+    Compute the discrete Frechet distance between two sequences.
+
+    ``seq1`` and ``seq2`` must have shape ``(T, D)``, where ``T`` is the
+    number of time steps and ``D`` is the feature dimension.
+    """
+    n, m = seq1.shape[0], seq2.shape[0]
+
+    seq1_d, seq2_d = seq1.double(), seq2.double()
+    frechet_matrix = torch.full(
+        (n, m), float("inf"), dtype=torch.float64, device=seq1.device
+    )
+    frechet_matrix[0, 0] = torch.dist(seq1_d[0], seq2_d[0])
+
+    for j in range(1, m):
+        frechet_matrix[0, j] = torch.max(
+            frechet_matrix[0, j - 1], torch.dist(seq1_d[0], seq2_d[j])
+        )
+
+    for i in range(1, n):
+        frechet_matrix[i, 0] = torch.max(
+            frechet_matrix[i - 1, 0], torch.dist(seq1_d[i], seq2_d[0])
+        )
+
+    for i in range(1, n):
+        for j in range(1, m):
+            frechet_matrix[i, j] = torch.max(
+                torch.min(
+                    torch.stack(
+                        [
+                            frechet_matrix[i - 1, j],
+                            frechet_matrix[i, j - 1],
+                            frechet_matrix[i - 1, j - 1],
+                        ]
+                    )
+                ),
+                torch.dist(seq1_d[i], seq2_d[j]),
+            )
+
+    return frechet_matrix[n - 1, m - 1]