# {py:mod}`evorl.envs.brax`

```{py:module} evorl.envs.brax
```

```{autodoc2-docstring} evorl.envs.brax
:parser: autodoc2_docstrings_parser
:allowtitles:
```

## Module Contents

### Classes

````{list-table}
:class: autosummary longtable
:align: left

* - {py:obj}`BraxAdapter <evorl.envs.brax.BraxAdapter>`
  - ```{autodoc2-docstring} evorl.envs.brax.BraxAdapter
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### Functions

````{list-table}
:class: autosummary longtable
:align: left

* - {py:obj}`create_brax_env <evorl.envs.brax.create_brax_env>`
  - ```{autodoc2-docstring} evorl.envs.brax.create_brax_env
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`create_wrapped_brax_env <evorl.envs.brax.create_wrapped_brax_env>`
  - ```{autodoc2-docstring} evorl.envs.brax.create_wrapped_brax_env
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### API

`````{py:class} BraxAdapter(env: brax.envs.Env)
:canonical: evorl.envs.brax.BraxAdapter

Bases: {py:obj}`evorl.envs.env.EnvAdapter`

```{autodoc2-docstring} evorl.envs.brax.BraxAdapter
:parser: autodoc2_docstrings_parser
```

````{py:property} action_space
:canonical: evorl.envs.brax.BraxAdapter.action_space
:type: evorl.envs.space.Space

````

````{py:property} obs_space
:canonical: evorl.envs.brax.BraxAdapter.obs_space
:type: evorl.envs.space.Space

````

````{py:method} reset(key: chex.PRNGKey) -> evorl.envs.env.EnvState
:canonical: evorl.envs.brax.BraxAdapter.reset

````

````{py:method} step(state: evorl.envs.env.EnvState, action: evorl.types.Action) -> evorl.envs.env.EnvState
:canonical: evorl.envs.brax.BraxAdapter.step

````

`````

````{py:function} create_brax_env(env_name: str, **kwargs) -> evorl.envs.brax.BraxAdapter
:canonical: evorl.envs.brax.create_brax_env

```{autodoc2-docstring} evorl.envs.brax.create_brax_env
:parser: autodoc2_docstrings_parser
```
````

````{py:function} 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
:canonical: evorl.envs.brax.create_wrapped_brax_env

```{autodoc2-docstring} evorl.envs.brax.create_wrapped_brax_env
:parser: autodoc2_docstrings_parser
```
````