evorl.envs.brax

Module Contents

Classes

BraxAdapter

Adapter for Brax environments.

Functions

create_brax_env

Create Brax environment.

create_wrapped_brax_env

Create wrapped Brax environment for training.

API

class evorl.envs.brax.BraxAdapter(env: brax.envs.Env)[source]

Bases: evorl.envs.env.EnvAdapter

Adapter for Brax environments.

property action_space: evorl.envs.space.Space
property obs_space: evorl.envs.space.Space
reset(key: chex.PRNGKey) → evorl.envs.env.EnvState[source]
step(state: evorl.envs.env.EnvState, action: evorl.types.Action) → evorl.envs.env.EnvState[source]
evorl.envs.brax.create_brax_env(env_name: str, **kwargs) → evorl.envs.brax.BraxAdapter[source]

Create Brax environment.

Parameters:

  • env_name – Environment name.

  • kwargs – Arguments passing into Brax.

Returns:

Brax env.

evorl.envs.brax.create_wrapped_brax_env(env_name: str, episode_length: int = 1000, parallel: int = 1, autoreset_mode: evorl.envs.wrappers.training_wrapper.AutoresetMode = AutoresetMode.NORMAL, discount: float | None = 1.0, record_ori_obs: bool = False, **kwargs) → evorl.envs.env.Env[source]

Create wrapped Brax environment for training.

Parameters:

  • env_name – Environment name.

  • episode_length – Max episode length.

  • parallel – Number of parallel environments.

  • autoreset_mode – Autoreset mode.

  • discount – Discount factor.

  • record_ori_obs – Whether record original observation in AutoresetMode.NORMAL and AutoresetMode.FAST mode.

  • kwargs – Other arguments passing into Brax.

Returns:

Wrapped Brax env.