jaxdem.rl.envWrappers#

Contains wrappers for modifying RL environments.

Functions

`clip_action_env`(env[, min_val, max_val])	Wrap an environment so that its step method clips the action 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 while preserving all current field values.
`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 before calling the original step.

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

Check whether an environment instance is a wrapped environment.

Parameters:: env (Environment) – The environment instance to check.
Returns:: True if the environment is wrapped (i.e., has a _base_env_cls attribute), False otherwise.
Return type:: bool

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

Unwrap an environment to its original base class while preserving 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 (scalar) environment is first broadcast to a batch of n identical copies, so callers do not need the jax.vmap(lambda _: env)(jnp.arange(n)) incantation themselves. Call env.reset(env, keys) afterwards to randomize each copy.

Example

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