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)	Promote an environment instance to a parallel version by applying jax.vmap(...) to its static methods.

jaxdem.rl.envWrappers.vectorise_env(env: Environment) → Environment

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

Wrap an environment so that its step method clips the action before calling the original step.

jaxdem.rl.envWrappers.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 (i.e., has a _base_env_cls attribute), False otherwise.

Return type:

bool

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

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