Jittable `Environment` class

[Robert Tjarko Lange]

Similar to how distributions work in distrax, I want to change the API to work with a jittable environment class. E.g.

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.

References:

• distrax Jittable abstract class: https://github.com/deepmind/distrax/blob/master/distrax/_src/utils/jittable.py
• distrax Distribution abstract class: https://github.com/deepmind/distrax/blob/master/distrax/_src/distributions/distribution.py
• distrax Example Categorical distribution: https://github.com/deepmind/distrax/blob/master/distrax/_src/distributions/categorical.py

[Robert Tjarko Lange]

Addressed in dd3e728.