search

RobertTLange / gymnax

RL Environments in JAX 🌍
Apache License 2.0
577 stars 54 forks source link

GymnaxtoBrax Wrapper imports failing #65

Open limbryan opened 9 months ago

limbryan commented 9 months ago

Hey!

https://github.com/RobertTLange/gymnax/blob/9155a960c8fc951c58c562e5c207069da4741bf5/gymnax/wrappers/brax.py#L9

This line seems to be throwing an error when using the wrapper. Looking at the other wrappers, should probably be something like this? 

from gymnax.environments.environment import Environment, EnvState, EnvParams