[Proposal] Add Optional State Variable in initial method of FuncEnv

[realquantumcookie]

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 NotImplementedError

I propose to change this to

def initial(self, rng : Any, state : Optional[StateType] = None) -> StateType

Otherwise, 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) -> StateType

Motivation

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

[Kallinteris-Andreas]

Could you please provide your runtime performance measurements from profiling?

I believe just in time compilations should optimize that way