arlbench.core.environments.brax_env

Brax environment adapter.

Classes

BraxEnv(env_name, n_envs[, env_kwargs])

A brax-based RL environment.
class arlbench.core.environments.brax_env.BraxEnv(env_name, n_envs, env_kwargs=None)[source]

Bases: Environment

A brax-based RL environment.

property action_space: Space

The action space of the environment.

property observation_space: Space

The observation space of the environment.

reset(rng)[source]

Resets the environment.

Return type:

tuple[State, Union[Array, ndarray, bool_, number]]

step(env_state, action, rng)[source]

Steps the environment forward.

Return type:

tuple[State, tuple[Union[Array, ndarray, bool_, number], Union[Array, ndarray, bool_, number], Union[Array, ndarray, bool_, number], dict]]

previous

arlbench.core.environments.autorl_env

next

arlbench.core.environments.envpool_env