arlbench.core.environments package

Submodules

arlbench.core.environments.autorl_env module

AutoRL Environment module.

class arlbench.core.environments.autorl_env.Environment(env_name, env, n_envs, seed=None)[source]

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()[source]

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()[source]

The observation space of the environment (gymnax space).

Returns:

Observation space of the environment.

Return type:

gymnax.environments.spaces.Space

abstract reset(rng)[source]

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)[source]

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)[source]
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]

arlbench.core.environments.brax_env module

Brax environment adapter.

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

arlbench.core.environments.envpool_env module

Envpool environment adapter.

class arlbench.core.environments.envpool_env.EnvpoolEnv(env_name, n_envs, seed, env_kwargs=None)[source]

Bases: Environment

An envpool-based RL environment.

property action_space

Action space of the environment.

property observation_space

Observation space of the environment.

reset(_)[source]

Resets the environment.

step(env_state, action, _)[source]

Steps the environment forward by one step.

arlbench.core.environments.envpool_env.numpy_to_jax(x)[source]

Converts numpy arrays to jax numpy arrays.

Return type:

Array

arlbench.core.environments.gymnasium_env module

Gymnasium environment adapter.

class arlbench.core.environments.gymnasium_env.GymnasiumEnv(env_name, seed, env_kwargs=None)[source]

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)[source]

Wraps actual reset for jitting.

Return type:

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

step(env_state, action, rng)[source]

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

arlbench.core.environments.gymnax_env module

Gymnax environment adapter.

class arlbench.core.environments.gymnax_env.GymnaxEnv(env_name, n_envs, env_kwargs=None)[source]

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)[source]

Resets the environment.

sample_action(rng)[source]

Samples a random action from the action space.

step(env_state, action, rng)[source]

Steps the environment forward.

arlbench.core.environments.make_env module

Environment creation function for ARLBench.

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

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

arlbench.core.environments.xland_env module

XLand environment adapter.

class arlbench.core.environments.xland_env.XLandEnv(env_name, n_envs, env_kwargs=None, cnn_policy=False)[source]

Bases: Environment

A xland-based RL environment.

property action_space

Action space of the environment.

property observation_space

Observation space of the environment.

reset(rng)[source]

Resets the environment.

sample_action(rng)[source]

Samples a random action from the action space.

step(env_state, action, rng)[source]

Steps the environment forward by one step.

Module contents

class arlbench.core.environments.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]]

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

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()[source]

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()[source]

The observation space of the environment (gymnax space).

Returns:

Observation space of the environment.

Return type:

gymnax.environments.spaces.Space

abstract reset(rng)[source]

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)[source]

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)[source]
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)[source]

Bases: Environment

An envpool-based RL environment.

property action_space

Action space of the environment.

property observation_space

Observation space of the environment.

reset(_)[source]

Resets the environment.

step(env_state, action, _)[source]

Steps the environment forward by one step.

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

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)[source]

Wraps actual reset for jitting.

Return type:

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

step(env_state, action, rng)[source]

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)[source]

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)[source]

Resets the environment.

sample_action(rng)[source]

Samples a random action from the action space.

step(env_state, action, rng)[source]

Steps the environment forward.

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

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