brax
brax copied to clipboard
Add _get_obs to fast environment
Hello,
fast env is designed as a minimalistic env for fast unit testing. However, it does expose a _get_obs method like the other brax environments. This breaks tests that users could create (at least I did) that rely on this method, while those tests would work for other brax environments (at the expense of a much slower test time).
I propose adding a _get_obs method to fast env to more accurately reflect the structure of brax environments.
I believe the current test suite covers this feature, please correct me if I am wrong.