# {py:mod}`evorl.envs.wrappers.ma_training_wrapper`

```{py:module} evorl.envs.wrappers.ma_training_wrapper
```

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

## Module Contents

### Classes

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

* - {py:obj}`EpisodeWrapper <evorl.envs.wrappers.ma_training_wrapper.EpisodeWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.EpisodeWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`FastVmapAutoResetWrapper <evorl.envs.wrappers.ma_training_wrapper.FastVmapAutoResetWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.FastVmapAutoResetWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`OneEpisodeWrapper <evorl.envs.wrappers.ma_training_wrapper.OneEpisodeWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.OneEpisodeWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`VmapAutoResetWrapper <evorl.envs.wrappers.ma_training_wrapper.VmapAutoResetWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.VmapAutoResetWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`VmapWrapper <evorl.envs.wrappers.ma_training_wrapper.VmapWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.VmapWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### API

`````{py:class} EpisodeWrapper(env: evorl.envs.env.Env, episode_length: int)
:canonical: evorl.envs.wrappers.ma_training_wrapper.EpisodeWrapper

Bases: {py:obj}`evorl.envs.wrappers.wrapper.Wrapper`

```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.EpisodeWrapper
:parser: autodoc2_docstrings_parser
```

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

````

````{py:method} step(state: evorl.envs.env.EnvState, action: jax.Array) -> evorl.envs.env.EnvState
:canonical: evorl.envs.wrappers.ma_training_wrapper.EpisodeWrapper.step

````

`````

`````{py:class} FastVmapAutoResetWrapper(env: evorl.envs.env.Env, num_envs: int = 1)
:canonical: evorl.envs.wrappers.ma_training_wrapper.FastVmapAutoResetWrapper

Bases: {py:obj}`evorl.envs.wrappers.wrapper.Wrapper`

```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.FastVmapAutoResetWrapper
:parser: autodoc2_docstrings_parser
```

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

````

````{py:method} step(state: evorl.envs.env.EnvState, action: jax.Array) -> evorl.envs.env.EnvState
:canonical: evorl.envs.wrappers.ma_training_wrapper.FastVmapAutoResetWrapper.step

````

`````

`````{py:class} OneEpisodeWrapper(env: evorl.envs.env.Env, episode_length: int)
:canonical: evorl.envs.wrappers.ma_training_wrapper.OneEpisodeWrapper

Bases: {py:obj}`evorl.envs.wrappers.ma_training_wrapper.EpisodeWrapper`

```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.OneEpisodeWrapper
:parser: autodoc2_docstrings_parser
```

````{py:method} step(state: evorl.envs.env.EnvState, action: jax.Array) -> evorl.envs.env.EnvState
:canonical: evorl.envs.wrappers.ma_training_wrapper.OneEpisodeWrapper.step

````

`````

`````{py:class} VmapAutoResetWrapper(env: evorl.envs.env.Env, num_envs: int = 1)
:canonical: evorl.envs.wrappers.ma_training_wrapper.VmapAutoResetWrapper

Bases: {py:obj}`evorl.envs.wrappers.wrapper.Wrapper`

```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.VmapAutoResetWrapper
:parser: autodoc2_docstrings_parser
```

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

````

````{py:method} step(state: evorl.envs.env.EnvState, action: jax.Array) -> evorl.envs.env.EnvState
:canonical: evorl.envs.wrappers.ma_training_wrapper.VmapAutoResetWrapper.step

````

`````

`````{py:class} VmapWrapper(env: evorl.envs.env.Env, num_envs: int = 1, vmap_step: bool = False)
:canonical: evorl.envs.wrappers.ma_training_wrapper.VmapWrapper

Bases: {py:obj}`evorl.envs.wrappers.wrapper.Wrapper`

```{autodoc2-docstring} evorl.envs.wrappers.ma_training_wrapper.VmapWrapper
:parser: autodoc2_docstrings_parser
```

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

````

````{py:method} step(state: evorl.envs.env.EnvState, action: jax.Array) -> evorl.envs.env.EnvState
:canonical: evorl.envs.wrappers.ma_training_wrapper.VmapWrapper.step

````

`````