arlbench.core.environments

class arlbench.core.environments.BraxEnv(env_name, n_envs, env_kwargs=None)

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)

Resets the environment.

Return type:

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

step(env_state, action, rng)

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

class arlbench.core.environments.Environment(env_name, env, n_envs, seed=None)

Bases: ABC

An abstract environment class to support various kinds of RL environments.

Sub-classes need to implement the following methods:

• reset()

• step()

Note: Both functions need to be fully jittable to support JAX-based RL algorithms!

As well as the properties:

• action_space

• observation_space

Note: These need to be gymnax spaces, not gymnasium spaces.

abstract action_space()

The action space of the environment (gymnax space).

Returns:

Action space of the environment.

Return type:

gymnax.environments.spaces.Space

property env_name: str

Returns the name/id of the environments.

Returns:

Environment name.

Return type:

str

property n_envs: int

The number of environments.

Returns:

_description_

Return type:

int

abstract observation_space()

The observation space of the environment (gymnax space).

Returns:

Observation space of the environment.

Return type:

gymnax.environments.spaces.Space

abstract reset(rng)

Environment reset() function. Resets the internal environment state.

Parameters:

rng (PRNGKey) – Random number generator key.

Returns:

Returns a tuple containing the environment state

as well as the actual return of the reset() function.

Return type:

tuple[Any, Any]

sample_actions(rng)

Samples a random action for each environment.

Parameters:

rng (PRNGKey) – Random number generator key.

Returns:

Array of sampled actions, one for each environment.

Return type:

jnp.ndarray

abstract step(env_state, action, rng)

Environment step() function. Performs a step

in the environment given an action.

Parameters:

• env_state (Any) – Internal environment state.

• action (Any) – Action to take.

• rng (PRNGKey) – Random number generator key.

Returns:

Returns a tuple containing the environment state

as well as the actual return of the step() function.

Return type:

tuple[Any, Any]

class arlbench.core.environments.EnvpoolEnv(env_name, n_envs, seed, env_kwargs=None)

Bases: Environment

An envpool-based RL environment.

property action_space

Action space of the environment.

property observation_space

Observation space of the environment.

reset(_)

Resets the environment.

step(env_state, action, _)

Steps the environment forward by one step.

class arlbench.core.environments.GymnasiumEnv(env_name, seed, env_kwargs=None)

Bases: Environment

A gymnasium-based RL environment.

property action_space

Action space of the environment.

property observation_space

Observation space of the environment.

reset(rng)

Wraps actual reset for jitting.

Return type:

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

step(env_state, action, rng)

Wraps actual step for jitting.

Return type:

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

class arlbench.core.environments.GymnaxEnv(env_name, n_envs, env_kwargs=None)

Bases: Environment

A gymnax-based RL environment.

property action_space

Action space of the environment.

property observation_space

Observation space of the environment.

reset(rng)

Resets the environment.

sample_action(rng)

Samples a random action from the action space.

step(env_state, action, rng)

Steps the environment forward.

arlbench.core.environments.make_env(env_framework, env_name, cnn_policy=False, n_envs=1, seed=0, env_kwargs=None)

ARLBench equivalent to make_env in gymnasium/gymnax etc. Creates a JAX-compatible RL environment.

Parameters:

• env_framework (str) – Environment framework to use. Must be one of the following: brax, envpool, gymnasium, gymnax, xland

• env_name (str) – Name/id of the environment. Has to match the env_framework.

• cnn_policy (bool, optional) – _description_. Defaults to False.

• n_envs (int, optional) – Number of environments. Defaults to 1.

• seed (int, optional) – Random seed. Defaults to 0.

• env_kwargs (dict[str, Any] | None, optional) – Keyword arguments to pass to the environment. Defaults to None.

Returns:

JAX-compatible RL environment.

Return type:

Environment | Wrapper

Modules

autorl_env	AutoRL Environment module.
brax_env	Brax environment adapter.
envpool_env	Envpool environment adapter.
gymnasium_env	Gymnasium environment adapter.
gymnax_env	Gymnax environment adapter.
make_env(env_framework, env_name[, ...])	ARLBench equivalent to make_env in gymnasium/gymnax etc.
xland_env	XLand environment adapter.

previous

arlbench.core.algorithms.sac.sac

next

arlbench.core.environments.autorl_env