jaxdem.rl.env_wrappers

Wrappers that modify RL environments.

Functions

clip_action_env(env[, min_val, max_val])	Clip actions consistently for checkpointing and physics steps.
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.env_wrappers.vectorise_env(env: Environment, n: int | None = None) → Environment

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))

jaxdem.rl.env_wrappers.clip_action_env(env: Environment, min_val: float = -1.0, max_val: float = 1.0) → Environment

Clip actions consistently for checkpointing and physics steps.

Both checkpoint(env, action) and step(env, action) receive the action clipped to [min_val, max_val]. This keeps action-dependent historical baselines consistent with the applied dynamics. Inactive slots stay zero even when zero lies outside the clipping interval. This wrapper requires floating-point continuous actions; integer categories are rejected.

jaxdem.rl.env_wrappers.is_wrapped(env: Environment) → bool

Check whether an environment instance is a wrapped environment.

Parameters:

env (Environment) – The environment instance to check.

Returns:

True if the environment is wrapped, False otherwise.

Return type:

bool

jaxdem.rl.env_wrappers.unwrap(env: Environment) → Environment

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