Closed RobertTLange closed 3 years ago
Similar to how distributions work in distrax, I want to change the API to work with a jittable environment class. E.g.
distrax
env = gymnax.make('env_name') obs, state = env.reset(key) obs, state, reward, done, info = env.step(key, state, action)
Hence, the environment parameters are "absorbed" in the class instance. This should not be too difficult as long as we are careful about the pytree.
pytree
References:
Addressed in dd3e7288fb9fcc094442e48b7050de90d58d0d82.
Similar to how distributions work in
distrax
, I want to change the API to work with a jittable environment class. E.g.Hence, the environment parameters are "absorbed" in the class instance. This should not be too difficult as long as we are careful about the
pytree
.References:
distrax
Jittable abstract class: https://github.com/deepmind/distrax/blob/master/distrax/_src/utils/jittable.pydistrax
Distribution abstract class: https://github.com/deepmind/distrax/blob/master/distrax/_src/distributions/distribution.pydistrax
Example Categorical distribution: https://github.com/deepmind/distrax/blob/master/distrax/_src/distributions/categorical.py