search

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.
Apache License 2.0
2.14k stars 234 forks source link

What replaced the old "pmap.is_synchornized" called in new brax versions? #494

Closed Mi-Przystupa closed 1 day ago

Mi-Przystupa commented 1 week ago

Hello,

I'm modifying a repository that depends heavily on Brax to work, but the version they use is terribly out of date so I am forced to modify things to align with the latest brax version.

This is the repo: https://github.com/frt03/mxt_bench/

Specifically, I would like to know what function or equivalent behavior does the same as 

brax.training.pmap.is_synchronized(...)

The code I'm referring to is here: https://github.com/frt03/mxt_bench/blob/8c05bd2b834a6b6d6066b8ae74dc317306ffb499/mxt_bench/algo/ppo_mlp.py#L468

My main goal is just to get it to work, so knowing what the new equivalent set of functions to address this would help.

Thank you!

erikfrey commented 1 day ago

Hi @Mi-Przystupa - the function name was changed to is_replicated: https://github.com/google/brax/blob/main/brax/training/pmap.py#L44