Add WOSAC interaction + map metrics. Switch from np -> torch.

[WaelDLZ]

PR to add all the WOSAC metrics (except for traffic lights)

• Interaction and Map-based metrics are implemented in torch now

[daphne-cornelisse]

Still need to see how we manage the Tracks_to_predict flag, if we want the interactive metrics to make sense we should control as much agents as we can. Or we should also take into account agents that were valid but not controlled in the distance computation (in interaction_features.py)

Is fixed now

[WaelDLZ]

This PR would make the WOSAC metrics complete, but also kinda slow (same issue as in the original WOSAC code).

Is also fixed, using torch make things way faster if you have access to a GPU.

Test on a A100 80Gb, with 1024 controlled agents (and ~110 evaluated agents): rollouts take ~22s, compute metrics takes about 5s. (it was taking >5min on numpy)

With 4096 controlled agents (~430 evaluated agents): rollouts take ~90s and compute metrics ~20s

It could be made faster in future work by processing all the scenarios in parallel

[greptile-apps]

Greptile Overview

Greptile Summary

This PR implements interaction-based and map-based metrics for the WOSAC (Waymo Open Sim Agents Challenge) evaluation framework, completing the metric coverage except for traffic lights.

Major changes:

• Added collision detection and time-to-collision computation using PyTorch-based Minkowski sums and signed distance calculations
• Implemented signed distance to road edges with proper handling of cyclic polylines and complex road geometries (donut roads, acute corners)
• Migrated distance computations from NumPy to PyTorch for GPU acceleration
• Added scenario-level likelihood estimation for boolean features (collision, offroad) using Bernoulli distributions
• Extended evaluator to compute 5 new metrics: distance to nearest object, time to collision, collision indication, distance to road edge, and offroad indication
• Renamed control_tracks_to_predict to control_wosac for clarity
• Added C++ bindings to retrieve agent dimensions and road edge polylines

Code organization:

• Refactored codebase into modular components: geometry_utils.py, kinematic_features.py, interaction_features.py, map_metric_features.py
• Comprehensive test coverage with visual validation for map metrics

Testing:
The PR includes thorough unit tests covering edge cases like invalid objects, cyclic polylines, and various geometric configurations.

Confidence Score: 4/5

• safe to merge with minor performance consideration
• the implementation is mathematically sound with comprehensive tests covering edge cases. geometry utilities properly handle complex scenarios like cyclic polylines and acute corners. the main concern is a CPU-GPU data transfer inefficiency in TTC computation that should be optimized but doesn't affect correctness
• pufferlib/ocean/benchmark/interaction_features.py - contains inefficient CPU-GPU round-trip in compute_time_to_collision function

Important Files Changed

File Analysis

Filename	Score	Overview
pufferlib/ocean/benchmark/metrics.py	4/5	added interaction and map metric computation functions (compute_interaction_features, compute_map_features), moved helper functions to separate modules, switched from NumPy to PyTorch for distance calculations
pufferlib/ocean/benchmark/interaction_features.py	3/5	new file implementing collision detection and time-to-collision computation using PyTorch with Minkowski sums; contains CPU-GPU data transfer inefficiency in compute_time_to_collision (line 196-204)
pufferlib/ocean/benchmark/map_metric_features.py	4/5	new file implementing signed distance to road edges with proper handling of cyclic polylines and donut-shaped roads using PyTorch
pufferlib/ocean/benchmark/geometry_utils.py	5/5	new file with PyTorch implementations of 2D geometry operations (box corners, Minkowski sums, signed distances to convex polygons)
pufferlib/ocean/benchmark/evaluator.py	4/5	extended WOSAC evaluator to compute interaction (collision, TTC, distance to nearest object) and map-based (offroad, distance to road edge) metrics with proper aggregation
pufferlib/ocean/drive/drive.py	4/5	renamed control mode from control_tracks_to_predict to control_wosac, added methods to retrieve agent dimensions and road edge polylines for metrics computation
pufferlib/ocean/drive/drive.h	4/5	refactored should_control_agent to use switch statement, added c_get_road_edge_counts and c_get_road_edge_polylines functions, updated control mode naming

Sequence Diagram

sequenceDiagram
    participant Evaluator as WOSACEvaluator
    participant Metrics as metrics.py
    participant Interaction as interaction_features.py
    participant Map as map_metric_features.py
    participant Geometry as geometry_utils.py
    participant Estimators as estimators.py
    
    Evaluator->>Metrics: compute_kinematic_features(x, y, heading)
    Metrics->>Metrics: compute speeds & accelerations
    Metrics-->>Evaluator: kinematic features
    
    Evaluator->>Metrics: compute_interaction_features(x, y, heading, scenario_ids, dimensions, eval_mask)
    loop For each scenario
        Metrics->>Interaction: compute_distance_to_nearest_object()
        Interaction->>Geometry: get_2d_box_corners()
        Interaction->>Geometry: minkowski_sum_of_box_and_box_points()
        Interaction->>Geometry: signed_distance_from_point_to_convex_polygon()
        Interaction-->>Metrics: distances
        Metrics->>Interaction: compute_time_to_collision()
        Interaction->>Metrics: compute_kinematic_features() [CPU transfer]
        Interaction-->>Metrics: TTC values
    end
    Metrics-->>Evaluator: distances, collisions, TTC
    
    Evaluator->>Metrics: compute_map_features(x, y, heading, scenario_ids, dimensions, polylines)
    loop For each scenario
        Metrics->>Map: compute_distance_to_road_edge()
        Map->>Geometry: get_2d_box_corners()
        Map->>Map: _compute_signed_distance_to_polylines()
        Map-->>Metrics: distances
    end
    Metrics-->>Evaluator: road edge distances, offroad indicators
    
    Evaluator->>Estimators: log_likelihood_estimate_timeseries()
    Estimators-->>Evaluator: kinematic likelihoods
    
    Evaluator->>Estimators: log_likelihood_estimate_scenario_level()
    Estimators-->>Evaluator: collision & offroad likelihoods

Loading

[greptile-apps]

_{22 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

style: inefficient CPU-GPU data transfer within TTC computation

transfers all position data to CPU, computes speed in NumPy, then transfers back to GPU. consider implementing speed calculation directly in PyTorch to avoid this round-trip

Prompt To Fix With AI

This is a comment left during a code review.
Path: pufferlib/ocean/benchmark/interaction_features.py
Line: 196:204

Comment:
**style:** inefficient CPU-GPU data transfer within TTC computation

transfers all position data to CPU, computes speed in NumPy, then transfers back to GPU. consider implementing speed calculation directly in PyTorch to avoid this round-trip

How can I resolve this? If you propose a fix, please make it concise.