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

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

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

## Package Contents

### Classes

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

* - {py:obj}`ActionRepeatWrapper <evorl.envs.wrappers.action_wrapper.ActionRepeatWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.action_wrapper.ActionRepeatWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`ActionSquashWrapper <evorl.envs.wrappers.action_wrapper.ActionSquashWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.action_wrapper.ActionSquashWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`EpisodeWrapper <evorl.envs.wrappers.training_wrapper.EpisodeWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.training_wrapper.EpisodeWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`FastVmapAutoResetWrapper <evorl.envs.wrappers.training_wrapper.FastVmapAutoResetWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.training_wrapper.FastVmapAutoResetWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`ObsFlattenWrapper <evorl.envs.wrappers.obs_wrapper.ObsFlattenWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.obs_wrapper.ObsFlattenWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`OneEpisodeWrapper <evorl.envs.wrappers.training_wrapper.OneEpisodeWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.training_wrapper.OneEpisodeWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`RewardScaleWrapper <evorl.envs.wrappers.reward_wrapper.RewardScaleWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.reward_wrapper.RewardScaleWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`SparseRewardWrapper <evorl.envs.wrappers.reward_wrapper.SparseRewardWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.reward_wrapper.SparseRewardWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`VmapAutoResetWrapper <evorl.envs.wrappers.training_wrapper.VmapAutoResetWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.training_wrapper.VmapAutoResetWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`VmapEnvPoolAutoResetWrapper <evorl.envs.wrappers.training_wrapper.VmapEnvPoolAutoResetWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.training_wrapper.VmapEnvPoolAutoResetWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`VmapWrapper <evorl.envs.wrappers.training_wrapper.VmapWrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.training_wrapper.VmapWrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`Wrapper <evorl.envs.wrappers.wrapper.Wrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.wrapper.Wrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### Functions

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

* - {py:obj}`get_wrapper <evorl.envs.wrappers.wrapper.get_wrapper>`
  - ```{autodoc2-docstring} evorl.envs.wrappers.wrapper.get_wrapper
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### API

`````{py:class} ActionRepeatWrapper(env: evorl.envs.env.Env, action_repeat: int)
:canonical: evorl.envs.wrappers.action_wrapper.ActionRepeatWrapper

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

```{autodoc2-docstring} evorl.envs.wrappers.action_wrapper.ActionRepeatWrapper
:parser: autodoc2_docstrings_parser
```

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

````

`````

`````{py:class} ActionSquashWrapper(env: evorl.envs.env.Env)
:canonical: evorl.envs.wrappers.action_wrapper.ActionSquashWrapper

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

```{autodoc2-docstring} evorl.envs.wrappers.action_wrapper.ActionSquashWrapper
:parser: autodoc2_docstrings_parser
```

````{py:property} action_space
:canonical: evorl.envs.wrappers.action_wrapper.ActionSquashWrapper.action_space
:type: evorl.envs.space.Space

````

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

````

`````

`````{py:class} EpisodeWrapper(env: evorl.envs.env.Env, episode_length: int, record_ori_obs: bool = True, discount: float | None = None)
:canonical: evorl.envs.wrappers.training_wrapper.EpisodeWrapper

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

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

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

````

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

````

`````

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

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

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

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

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

````

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

````

`````

`````{py:class} ObsFlattenWrapper(env: evorl.envs.env.Env)
:canonical: evorl.envs.wrappers.obs_wrapper.ObsFlattenWrapper

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

```{autodoc2-docstring} evorl.envs.wrappers.obs_wrapper.ObsFlattenWrapper
:parser: autodoc2_docstrings_parser
```

````{py:property} obs_space
:canonical: evorl.envs.wrappers.obs_wrapper.ObsFlattenWrapper.obs_space
:type: evorl.envs.space.Space

````

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

````

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

````

`````

`````{py:class} OneEpisodeWrapper(env: evorl.envs.env.Env, episode_length: int, record_ori_obs: bool = True, discount: float | None = None)
:canonical: evorl.envs.wrappers.training_wrapper.OneEpisodeWrapper

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

```{autodoc2-docstring} evorl.envs.wrappers.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.training_wrapper.OneEpisodeWrapper.step

````

`````

`````{py:class} RewardScaleWrapper(env: evorl.envs.env.Env, reward_scale: float)
:canonical: evorl.envs.wrappers.reward_wrapper.RewardScaleWrapper

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

```{autodoc2-docstring} evorl.envs.wrappers.reward_wrapper.RewardScaleWrapper
:parser: autodoc2_docstrings_parser
```

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

````

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

````

`````

`````{py:class} SparseRewardWrapper(env: evorl.envs.env.Env, sparse_length: int)
:canonical: evorl.envs.wrappers.reward_wrapper.SparseRewardWrapper

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

```{autodoc2-docstring} evorl.envs.wrappers.reward_wrapper.SparseRewardWrapper
:parser: autodoc2_docstrings_parser
```

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

````

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

````

`````

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

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

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

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

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

````

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

````

`````

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

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

```{autodoc2-docstring} evorl.envs.wrappers.training_wrapper.VmapEnvPoolAutoResetWrapper
:parser: autodoc2_docstrings_parser
```

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

```{autodoc2-docstring} evorl.envs.wrappers.training_wrapper.VmapEnvPoolAutoResetWrapper.reset
:parser: autodoc2_docstrings_parser
```

````

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

````

`````

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

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

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

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

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

````

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

````

`````

`````{py:class} Wrapper(env: evorl.envs.env.Env)
:canonical: evorl.envs.wrappers.wrapper.Wrapper

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

```{autodoc2-docstring} evorl.envs.wrappers.wrapper.Wrapper
:parser: autodoc2_docstrings_parser
```

````{py:property} action_space
:canonical: evorl.envs.wrappers.wrapper.Wrapper.action_space
:type: evorl.envs.env.Space

````

````{py:property} obs_space
:canonical: evorl.envs.wrappers.wrapper.Wrapper.obs_space
:type: evorl.envs.env.Space

````

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

````

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

````

````{py:property} unwrapped
:canonical: evorl.envs.wrappers.wrapper.Wrapper.unwrapped
:type: evorl.envs.env.Env

```{autodoc2-docstring} evorl.envs.wrappers.wrapper.Wrapper.unwrapped
:parser: autodoc2_docstrings_parser
```

````

`````

````{py:function} get_wrapper(env: evorl.envs.env.Env, wrapper_cls: type) -> evorl.envs.wrappers.wrapper.Wrapper | None
:canonical: evorl.envs.wrappers.wrapper.get_wrapper

```{autodoc2-docstring} evorl.envs.wrappers.wrapper.get_wrapper
:parser: autodoc2_docstrings_parser
```
````