Jumpy's random functions are a bit weird for numpy as it uses a np.array as the key but needs to convert this to the RandomState for each function call.
Imo, as Gymnasium and other projects use np.random.Generator then it makes it difficult to easily change to using jumpy rather than numpy.
Therefore, the proposal is to create a jp.random.Generator object that is equivalent to np.random.Generator for numpy and for jumpy, behind the scenes we store the key and split the prng and return the results.
To consider is if the object inherits from np.random.Generator with functions being overrides with a check if numpy or jax is used.
Jumpy's random functions are a bit weird for numpy as it uses a
np.arrayas the key but needs to convert this to theRandomStatefor each function call.Imo, as Gymnasium and other projects use
np.random.Generatorthen it makes it difficult to easily change to using jumpy rather than numpy.Therefore, the proposal is to create a
jp.random.Generatorobject that is equivalent tonp.random.Generatorfor numpy and for jumpy, behind the scenes we store the key and split theprngand return the results.To consider is if the object inherits from
np.random.Generatorwith functions being overrides with a check if numpy or jax is used.