Wrappers

AutoReset Wrapper

Bases: UnderspecifiedEnv

This is a wrapper around an UnderspecifiedEnv, allowing for the environment to be automatically reset upon completion of an episode. This behaviour is similar to the default Gymnax interface. The user can specify a callable sample_level that takes in a PRNGKey and returns a level.

Warning

To maintain compliance with UnderspecifiedEnv interface, user can reset to an arbitrary level. This includes levels outside the support of sample_level(). Consequently, the tagged rng is defaulted to jax.random.PRNGKey(0). If your code relies on this, careful attention may be required.

---

AutoResetFiniteWrapper

Bases: UnderspecifiedEnv

On episode termination, the environment is reset to a level specified by some fixed, discrete distribution with finite support.

Parameters:

Name	Type	Description	Default
env	UnderspecifiedEnv		required
levels	_type_	The set of levels to sample from.	required
p	_type_	probabilities for each level, defaults to uniform. Defaults to None.	None

Warning

To maintain compliance with the UnderspecifiedEnv interface, one can still reset to an arbitrary level. This includes levels outside the specified distribution. An optional flag is provided such that if the user wishes to check if a level is in distribution, the wrapped state will be marked with the first index index corresponding to the reset level in distribution.

---

AutoReplay Wrapper

Bases: UnderspecifiedEnv

This wrapper replay the same level over and over again by resetting to the same level after each episode. This is useful for training/rolling out multiple times on the same level.