arlbench.core.environments.make_env

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