# {py:mod}`evox.problems.neuroevolution.brax`

```{py:module} evox.problems.neuroevolution.brax
```

```{autodoc2-docstring} evox.problems.neuroevolution.brax
:allowtitles:
```

## Module Contents

### Classes

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

* - {py:obj}`BraxProblem <evox.problems.neuroevolution.brax.BraxProblem>`
  - ```{autodoc2-docstring} evox.problems.neuroevolution.brax.BraxProblem
    :summary:
    ```
````

### API

`````{py:class} BraxProblem(policy: torch.nn.Module, env_name: str, max_episode_length: int, num_episodes: int, seed: int = None, pop_size: int | None = None, rotate_key: bool = True, reduce_fn: typing.Callable[[torch.Tensor, int], torch.Tensor] = torch.mean, backend: str | None = None, device: torch.device | None = None, compile_policy: bool = True)
:canonical: evox.problems.neuroevolution.brax.BraxProblem

Bases: {py:obj}`evox.core.Problem`

```{autodoc2-docstring} evox.problems.neuroevolution.brax.BraxProblem
```

```{rubric} Initialization
```

```{autodoc2-docstring} evox.problems.neuroevolution.brax.BraxProblem.__init__
```

````{py:method} _evaluate_brax_record(key: torch.Tensor, model_state: typing.Dict[str, torch.Tensor]) -> typing.Tuple[typing.Dict[str, torch.Tensor], torch.Tensor, typing.List[typing.Any]]
:canonical: evox.problems.neuroevolution.brax.BraxProblem._evaluate_brax_record

```{autodoc2-docstring} evox.problems.neuroevolution.brax.BraxProblem._evaluate_brax_record
```

````

````{py:method} evaluate(pop_params: typing.Dict[str, torch.nn.Parameter]) -> torch.Tensor
:canonical: evox.problems.neuroevolution.brax.BraxProblem.evaluate

```{autodoc2-docstring} evox.problems.neuroevolution.brax.BraxProblem.evaluate
```

````

````{py:method} visualize(weights: typing.Dict[str, torch.nn.Parameter], seed: int | None = None, output_type: typing.Literal[HTML, rbg_array] = 'HTML', /, **kwargs) -> str | typing.Sequence[numpy.ndarray] | numpy.ndarray
:canonical: evox.problems.neuroevolution.brax.BraxProblem.visualize

```{autodoc2-docstring} evox.problems.neuroevolution.brax.BraxProblem.visualize
```

````

`````