Skip to content

Architectures¤

from jaxnasium.algorithms import SAC
from jaxnasium.algorithms.architectures import BroNet

algorithm = SAC(critic_kwargs={"body": BroNet.with_params(depth=2, width_size=256)})

Any Equinox module following the same convention works too; these are just the bundled ones.

jaxnasium.algorithms.architectures.SimBa ¤

Create a Simba neural network with the given hidden dimensions and output space. https://arxiv.org/pdf/2410.09754

Operates on 1D inputs. Observation normalization (RSNorm) is left to the agent (enable normalize_observations).

depth is the number of blocks.

__init__(in_features: int, *, key: PRNGKeyArray, depth: int = 1, width_size: int = 256) ¤
with_params(*, depth: int = 1, width_size: int = 256) -> Callable[..., Self] classmethod ¤

jaxnasium.algorithms.architectures.BroNet ¤

Create a BroNet neural network with the given hidden dimensions and output space. https://arxiv.org/html/2405.16158v1

Operates on 1D inputs.

__init__(in_features: int, *, key: PRNGKeyArray, depth: int = 1, width_size: int = 256) ¤
with_params(*, depth: int = 1, width_size: int = 256) -> Callable[..., Self] classmethod ¤

jaxnasium.algorithms.architectures.MLP ¤

Simple MLP architecture. Final hidden size is the output features. Final outputs are non-linear (unless the activation is set to identity). Attach additional layers if a different final activation is desired.

__init__(in_features: int, *, key: PRNGKeyArray, hidden_sizes: Sequence[int] = (128, 128), activation: Callable = jax.nn.relu) ¤
with_params(*, hidden_sizes: Sequence[int] = (128, 128), activation: Callable = jax.nn.relu) -> Callable[..., Self] classmethod ¤

jaxnasium.algorithms.architectures.CNN ¤

Standard CNN architecture similar to the DQN Nature paper.

Operates on 2D inputs. Assumes channels first format (C, H, W).

Flattens the output of the CNN into a 1d vector.

__init__(input_shape: tuple[int, int, int], *, key: PRNGKeyArray, out_channels: Sequence[int] = (32, 64, 64), kernel_sizes: Sequence[int] = (3, 3, 3), strides: Sequence[int] = (1, 1, 1), padding: Sequence[int] = (1, 1, 1), channels_axis: Literal['first', 'last'] = 'first', activation: Callable = jax.nn.relu) ¤
with_params(*, out_channels: Sequence[int] = (32, 64, 64), kernel_sizes: Sequence[int] = (3, 3, 3), strides: Sequence[int] = (1, 1, 1), padding: Sequence[int] = (1, 1, 1), channels_axis: Literal['first', 'last'] | None = None, activation: Callable = jax.nn.relu) -> Callable[..., Self] classmethod ¤

jaxnasium.algorithms.architectures.Identity ¤

Identity with out_features set up