Add WOSAC interaction + map metrics. Switch from np -> torch. #138
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Notes: - Need to add safeguards to load each map only once - Might be slow if we increase num_agents per scenario, next step will be torch. I added some tests to see the distance and ttc computations are correct, and metrics_sanity_check looks okay. I'll keep making some plots to validate it.
Ref in original code:
message BernoulliEstimate {
// Additive smoothing to apply to the underlying 2-bins histogram, to avoid
// infinite values for empty bins.
optional float additive_smoothing_pseudocount = 4 [default = 0.001];
}
First step: extract the map from the vecenv
A bunch of little tests in test_map_metric_features.py to ensure this do what it is supposed to do. python -m pufferlib.ocean.benchmark.test_map_metrics Next steps should be straightforward. Will need to check at some point if doing this on numpy isnt too slow
This works, and passes all the tests, I would still want to make additionnal checks with the renderer because we never know. With this, we have the whole set of WOSAC metrics (except for traffic lights), and we might also have the same issue as the original WOSAC code: it is slow. Next step would be to transition from numpy to torch.
…en WOSAC sees an offorad or a collision python pufferlib/ocean/benchmark/visual_sanity_check.py
Is fixed now |
…etrics. It makes the computation way faster, and all the tests pass. I didn't switch kinematics to torch because it was already fast, but I might make the change for consistency.
Author
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 |
… wbd/wosac_map_metrics
daphne-cornelisse
changed the title
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
Greptile OverviewGreptile SummaryThis 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:
Code organization:
Testing: Confidence Score: 4/5
Important Files ChangedFile Analysis
Sequence DiagramsequenceDiagram
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
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+196
to
+204
| # NOTE: I know this is ugly, I will fix it another day | ||
| speed = metrics.compute_kinematic_features( | ||
| x=center_x.cpu().numpy(), | ||
| y=center_y.cpu().numpy(), | ||
| heading=heading.cpu().numpy(), | ||
| seconds_per_step=seconds_per_step, | ||
| )[0] | ||
| if not isinstance(speed, torch.Tensor): | ||
| speed = torch.as_tensor(speed, device=center_x.device, dtype=center_x.dtype) |
There was a problem hiding this 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
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.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
PR to add all the WOSAC metrics (except for traffic lights)