# {py:mod}`evorl.evaluators.mo_brax_evaluator`

```{py:module} evorl.evaluators.mo_brax_evaluator
```

```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator
:parser: autodoc2_docstrings_parser
:allowtitles:
```

## Module Contents

### Classes

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

* - {py:obj}`BraxEvaluator <evorl.evaluators.mo_brax_evaluator.BraxEvaluator>`
  - ```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.BraxEvaluator
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### Functions

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

* - {py:obj}`eval_env_step <evorl.evaluators.mo_brax_evaluator.eval_env_step>`
  - ```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.eval_env_step
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`eval_metrics <evorl.evaluators.mo_brax_evaluator.eval_metrics>`
  - ```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.eval_metrics
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`eval_rollout_episode <evorl.evaluators.mo_brax_evaluator.eval_rollout_episode>`
  - ```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.eval_rollout_episode
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`fast_eval_metrics <evorl.evaluators.mo_brax_evaluator.fast_eval_metrics>`
  - ```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.fast_eval_metrics
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### API

`````{py:class} BraxEvaluator
:canonical: evorl.evaluators.mo_brax_evaluator.BraxEvaluator

Bases: {py:obj}`evorl.evaluators.evaluator.Evaluator`

```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.BraxEvaluator
:parser: autodoc2_docstrings_parser
```

````{py:method} evaluate(agent_state: chex.ArrayTree, key: chex.PRNGKey, num_episodes: int) -> chex.ArrayTree
:canonical: evorl.evaluators.mo_brax_evaluator.BraxEvaluator.evaluate

````

````{py:attribute} metric_names
:canonical: evorl.evaluators.mo_brax_evaluator.BraxEvaluator.metric_names
:type: tuple[str]
:value: >
   'pytree_field(...)'

```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.BraxEvaluator.metric_names
:parser: autodoc2_docstrings_parser
```

````

`````

````{py:function} eval_env_step(env_fn: evorl.envs.EnvStepFn, action_fn: evorl.agent.AgentActionFn, env_state: evorl.envs.EnvState, agent_state: evorl.agent.AgentState, key: chex.PRNGKey, metric_names: tuple[str] = ()) -> tuple[evorl.rollout.SampleBatch, evorl.envs.EnvState]
:canonical: evorl.evaluators.mo_brax_evaluator.eval_env_step

```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.eval_env_step
:parser: autodoc2_docstrings_parser
```
````

````{py:function} eval_metrics(env_fn: collections.abc.Callable[[evorl.envs.EnvState, evorl.types.Action], evorl.envs.EnvState], action_fn: collections.abc.Callable[[evorl.agent.AgentState, evorl.rollout.SampleBatch, chex.PRNGKey], tuple[evorl.types.Action, evorl.types.PolicyExtraInfo]], env_state: evorl.envs.EnvState, agent_state: evorl.agent.AgentState, key: chex.PRNGKey, rollout_length: int, discount: float, metric_names: tuple[str] = ()) -> tuple[evorl.types.PyTreeDict, evorl.envs.EnvState]
:canonical: evorl.evaluators.mo_brax_evaluator.eval_metrics

```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.eval_metrics
:parser: autodoc2_docstrings_parser
```
````

````{py:function} eval_rollout_episode(env_fn: collections.abc.Callable[[evorl.envs.EnvState, evorl.types.Action], evorl.envs.EnvState], action_fn: collections.abc.Callable[[evorl.agent.AgentState, evorl.rollout.SampleBatch, chex.PRNGKey], tuple[evorl.types.Action, evorl.types.PolicyExtraInfo]], env_state: evorl.envs.EnvState, agent_state: evorl.agent.AgentState, key: chex.PRNGKey, rollout_length: int, metric_names: tuple[str] = ()) -> tuple[evorl.rollout.SampleBatch, evorl.envs.EnvState]
:canonical: evorl.evaluators.mo_brax_evaluator.eval_rollout_episode

```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.eval_rollout_episode
:parser: autodoc2_docstrings_parser
```
````

````{py:function} fast_eval_metrics(env_fn: collections.abc.Callable[[evorl.envs.EnvState, evorl.types.Action], evorl.envs.EnvState], action_fn: collections.abc.Callable[[evorl.agent.AgentState, evorl.rollout.SampleBatch, chex.PRNGKey], tuple[evorl.types.Action, evorl.types.PolicyExtraInfo]], env_state: evorl.envs.EnvState, agent_state: evorl.agent.AgentState, key: chex.PRNGKey, rollout_length: int, metric_names: tuple[str] = ()) -> tuple[evorl.types.PyTreeDict, evorl.envs.EnvState]
:canonical: evorl.evaluators.mo_brax_evaluator.fast_eval_metrics

```{autodoc2-docstring} evorl.evaluators.mo_brax_evaluator.fast_eval_metrics
:parser: autodoc2_docstrings_parser
```
````