torch_world.metrics¶
Functions on plain tensors, so they apply to any model's features or predictions. torch_world.evaluation holds the protocols that run a model and score it.
Representation quality¶
linear_probe(
train_x: Tensor,
train_y: Tensor,
test_x: Tensor,
test_y: Tensor,
ridge: float = 0.001,
iterations: int = 100,
) -> Probe
Fit a linear map from frozen features (N, D) to targets and score it on the test split.
Integer targets (N,) are classes: multinomial logistic regression fit with full-batch L-BFGS, scored by
accuracy. Float targets (N,) or (N, K) are regressed in closed form, scored by R². Features are
standardized with the train statistics, and ridge is an L2 penalty on the weights. This is the linear
evaluation of I-JEPA and LeJEPA, and the physical probing of LeWM.
knn_accuracy(
train_x: Tensor,
train_y: Tensor,
test_x: Tensor,
test_y: Tensor,
k: int = 20,
temperature: float = 0.07,
center: bool = False,
) -> Probe
Weighted k-nearest-neighbor classification on cosine similarity: each of the k nearest train features votes
for its class with weight exp(similarity / temperature). Needs no training, as reported by DINO and
LeJEPA. center subtracts the mean train feature first, so a component shared by every input does not
dominate the cosine similarity.
Wu et al., "Unsupervised Feature Learning via Non-Parametric Instance Discrimination", CVPR 2018. Paper: https://arxiv.org/abs/1805.01978
attentive_probe(
train_x: Tensor,
train_y: Tensor,
test_x: Tensor,
test_y: Tensor,
epochs: int = 20,
batch_size: int = 256,
lr: float = 0.001,
weight_decay: float = 0.01,
num_heads: int = 8,
generator: Generator | None = None,
) -> Probe
Train an AttentiveProbe on frozen tokens (N, T, D) with class labels (N,) and report test accuracy,
the evaluation of V-JEPA and V-JEPA 2. Unlike a linear probe on pooled features, it can read which tokens
matter.
Bases: Module
A learned query token cross-attends to frozen tokens (B, N, D), followed by an MLP and a linear classifier.
Bardes et al., "Revisiting Feature Prediction for Learning Visual Representations from Video" (V-JEPA), 2024. Paper: https://arxiv.org/abs/2404.08471 Code: https://github.com/facebookresearch/jepa
dataclass
¶
Test score of a probe (accuracy for classes, R² for regression) and its test predictions.
Coefficient of determination of (N, K) predictions, averaged over the K targets.
Effective rank of embeddings (N, D): the exponential of the entropy of their normalized singular values.
Ranges from 1 (collapsed) to min(N, D), needs no labels, and tracks downstream accuracy across
hyperparameters of one method. As in the paper, the embeddings are used as they are; center subtracts their
mean first, for features dominated by one component shared by every input (pooled tokens of mostly-empty
images), which would otherwise read as rank 1.
Garrido et al., "RankMe: Assessing the Downstream Performance of Pretrained Self-Supervised Representations by Their Rank", ICML 2023. Paper: https://arxiv.org/abs/2210.02885
Effective rank of the linear discriminant analysis matrix of embeddings (N, Q, D): N surrogate classes
(for example images) with Q >= 2 samples each (for example augmented views).
The between-class covariance is whitened by the within-class covariance, so only directions that separate classes count. Designed for JEPAs, where it predicts linear-probe accuracy better than RankMe.
Thilak et al., "LiDAR: Sensing Linear Probing Performance in Joint Embedding SSL Architectures", ICLR 2024. Paper: https://arxiv.org/abs/2312.04000
Prediction quality¶
Mean L_p error per step between predicted and encoded latents (B, T, ...); returns (T,).
Open-loop prediction error grows with the horizon; its curve compares world models on the same encoder. V-JEPA 2-AC
trains and evaluates with p = 1; p = 2 gives the mean squared error of DINO-WM and LeWM. Spikes of the
one-step error measure surprise, as in LeWM's violation-of-expectation tests.
ssim(
pred: Tensor,
target: Tensor,
data_range: float = 1.0,
window: int = 11,
sigma: float = 1.5,
) -> Tensor
Structural similarity of images (..., C, H, W) with a Gaussian window, one value per image, as DINO-WM reports
for decoded predictions.
Wang et al., "Image Quality Assessment: From Error Visibility to Structural Similarity", IEEE TIP 2004. Paper: https://doi.org/10.1109/TIP.2003.819861
Peak signal-to-noise ratio in dB of images (..., C, H, W), one value per image.
Symmetric Chamfer distance between point sets (B, N, D) and (B, M, D): the mean squared distance from
each point to its nearest neighbor in the other set, summed over both directions. DINO-WM scores its deformable
tasks (rope, granular) with it.
Robotics and control¶
success_rate(
success: Tensor, confidence: float = 0.95
) -> Rate
Fraction of successful episodes (N,) with its Wilson score interval, which stays inside [0, 1] and is
reliable at the few dozen episodes robot evaluations afford (16 episodes at 50% give [28%, 72%]).
Wilson, "Probable Inference, the Law of Succession, and Statistical Inference", JASA 1927.
dataclass
¶
A success rate over n episodes with its Wilson score interval [low, high].
Index of the first successful step of each episode (N, T), or -1 for episodes that never succeed.
Fraction of the initial distance to the goal closed by the end of each episode (N, T): 1 at the goal,
0 when no closer than at the start, negative when farther.
Success weighted by path length: the mean of success * shortest / max(path, shortest) over episodes, with
the straight line from the start position to the goal position (N, D) as the shortest path. A success by
a detour twice too long counts one half.
Anderson et al., "On Evaluation of Embodied Navigation Agents", 2018. Paper: https://arxiv.org/abs/1807.06757
Mean size of the change between consecutive actions (N, T, A) per episode; lower is smoother. Jerky
commands that flip direction every step, typical of sampling-based planners, score high.
Retrieval¶
recall_at_k(
similarity: Tensor,
k: int | tuple[int, ...] = (1, 5, 10),
targets: Tensor | None = None,
)
Retrieval recall@k for a query-by-candidate similarity (Q, C): the fraction of queries whose correct
candidate (targets (Q,), by default the diagonal) ranks among the top k. Returns a float, or a dict for
several k, as VL-JEPA reports for text-to-video and video-to-text retrieval (transpose for the other way).
Fraction of rows of logits (N, C) whose k highest scores include the label, as in zero-shot
classification.