Skip to content

Conversation

@thowell
Copy link
Collaborator

@thowell thowell commented May 3, 2025

@thowell thowell linked an issue May 3, 2025 that may be closed by this pull request
15 tasks
@thowell thowell marked this pull request as ready for review May 13, 2025 15:33
Copy link
Collaborator

@erikfrey erikfrey left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One nit, and one question about whether this is JAX compatible

geomgroup = vec6(0, 0, 0, 0, 0, 0)

if bodyexclude is None:
bodyexclude = wp.zeros(nrays, dtype=int)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this JAX-safe? I thought we can't allocate arrays during cuda graph capture in JAX - no?

Copy link
Collaborator Author

@thowell thowell May 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure that ray.ray will work as is with cuda graph capture and be compatible with JAX.

the rangefinder sensor utilizes ray._ray and the Model field sensor_rangefinder_bodyid for the input to bodyexclude and should work with cuda graph capture.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#277 is updated to add fields ray_dist and ray_geomid to Data and ray.ray utilizes these fields instead of allocating arrays.

@erikfrey
Copy link
Collaborator

@thowell I see more changes to ray.ray in #277 - do you plan to make ray.ray JAX compatible in that PR, or this one? Let me know and I'll review/approve accordingly.

@thowell thowell merged commit c751ebe into google-deepmind:main May 30, 2025
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Sensors: sensor_pos (collection 1)

2 participants