Description
Proposal
Hello Gymnasium Contributors,
It is very nice to see that the gymnasium library is experimenting with pure function environments!
The current signature for the initial(...) method looks like
def initial(self, rng: Any) -> StateType:
"""Initial state."""
raise NotImplementedErrorI propose to change this to
def initial(self, rng : Any, state : Optional[StateType] = None) -> StateTypeOtherwise, we can also add a reset(self, rng, state) method for resetting (not initializing) the environment (I actually prefer this, since this allows us to write pure environment resets that can be JITted)
def reset(self, rng : Any, state : StateType) -> StateTypeMotivation
I'm currently writing a gymnasium environment with Mujoco / MJX backends and are trying to use the new FuncEnv experimental API.
However in a lot of mujoco / MJX environments it is very computationally expensive to generate mjData or mjx.mjData from an mjModel => We can just simply reset some qpos to reset the environment. So this modification allows that to happen.
Pitch
No response
Alternatives
No response
Additional context
Related: #833
Checklist
- I have checked that there is no similar issue in the repo