jaxdem.rl.envWrappers#

Wrappers that modify RL environments.

Functions

`clip_action_env`(env[, min_val, max_val])	Wrap an environment so that its step method clips the action to [min_val, max_val] before calling the original step.
`is_wrapped`(env)	Check whether an environment instance is a wrapped environment.
`unwrap`(env)	Unwrap an environment to its original base class.
`vectorise_env`(env[, n])	Promote an environment instance to a parallel version by applying jax.vmap(...) to its static methods.

jaxdem.rl.envWrappers.clip_action_env(env: Environment, min_val: float = -1.0, max_val: float = 1.0) → Environment[source]#: Wrap an environment so that its step method clips the action to [min_val, max_val] before calling the original step.

jaxdem.rl.envWrappers.is_wrapped(env: Environment) → bool[source]#

Check whether an environment instance is a wrapped environment.

jaxdem.rl.envWrappers.unwrap(env: Environment) → Environment[source]#

Unwrap an environment to its original base class. Keeps all current field values.

Parameters:: env (Environment) – The wrapped environment instance.
Returns:: A new instance of the original base environment class with the same field values as the wrapped instance.
Return type:: Environment

jaxdem.rl.envWrappers.vectorise_env(env: Environment, n: int | None = None) → Environment[source]#

Promote an environment instance to a parallel version by applying jax.vmap(…) to its static methods.

Parameters:

env (Environment) – The environment to vectorize. May already carry a leading batch dimension (e.g. produced with jax.vmap).
n (int, optional) – When given, the wrapper first broadcasts the (scalar) environment to a batch of n identical copies, so callers do not need to write jax.vmap(lambda _: env)(jnp.arange(n)) themselves. Then call env.reset(env, keys) to randomize each copy.

Example

>>> env = vectorise_env(env, n=32)
>>> env = env.reset(env, jax.random.split(key, 32))