google / brax

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

How to get observation space, or at least low/high bounds? #360

Open ikamensh opened 1 year ago

ikamensh commented 1 year ago

Hi,

I'm used to having observation_space in gym.Env / VectorEnv classes that tell me how to normalize observations, among other uses. Am I right to think this is missing in brax?

I'm doing

from brax import envs as brax_envs
vecenv = brax_envs.create("halfcheetah", batch_size=training_num)

and I can't find any property along the lines of observation space in the vecenv object

erikfrey commented 1 year ago

Hello there,

You can get these properties if you use the gym wrapper:

from brax.envs.wrappers import gym
env = gym.VectorGymWrapper(vecenv)
print(env.observation_space)
btaba commented 1 year ago

Hi @ikamensh , while the observation_space is populated, they're set to inf by default

https://github.com/google/brax/blob/3b72c952f7c990e8e9e355f736e77f10f408bc14/brax/envs/wrappers/gym.py#L114

If you'd like to make a PR to populate it generically, you can use dof limits

https://github.com/google/brax/blob/3b72c952f7c990e8e9e355f736e77f10f408bc14/brax/base.py#L294