Agent Networks

The networks the algorithms build. Each consists of three components:

1. An observation processor (PyTreeObsSpaceNetwork), which builds a network per observation space in the (possibly nested) observation PyTree. 1D spaces go through architecture_1d, 2D spaces through architecture_2d, and the outputs are concatenated into a single 1D vector.
2. A body (default SimBa), a shared network processing that vector.
3. An output processor (PyTreeActionNetwork / PyTreeQValueNetwork), which builds a head per output space, choosing a discrete or continuous head based on the space, and returns results in the same PyTree structure as the action space.

Constructor arguments are reachable through an algorithm's actor_kwargs / critic_kwargs:

from jaxnasium.algorithms import PPO
from jaxnasium.algorithms.architectures import MLP

algorithm = PPO(
    actor_kwargs={"body": MLP.with_params(hidden_sizes=(128, 128, 128))},
    critic_kwargs={"body": MLP.with_params(hidden_sizes=(128, 128, 128))},
)

jaxnasium.algorithms.ActorNetwork

__init__(obs_space: PyTree[jym.Space], output_space: PyTree[jym.Space], *, key: PRNGKeyArray, body: Callable[..., OutSizedNetwork] = SimBa, obs_architecture_1d: Callable[..., Network] = eqx.nn.Identity, obs_architecture_2d: Callable[..., Network] = _DEFAULT_ARCHITECTURE_2D, discrete_output_layer: Callable[..., Network] = CategoricalLayer, continuous_output_layer: Callable[..., Network] = NormalLayer, assume_independent_output: bool = True)

__call__(x, *, key: PRNGKeyArray | None = None)

jaxnasium.algorithms.ValueNetwork

__init__(obs_space: PyTree[jym.Space], *, key: PRNGKeyArray, body: Callable[..., OutSizedNetwork] = SimBa, obs_architecture_1d: Callable[..., Network] = eqx.nn.Identity, obs_architecture_2d: Callable[..., Network] = _DEFAULT_ARCHITECTURE_2D, output_layer_type: Callable[..., Network] = eqx.Partial(eqx.nn.Linear, out_features='scalar'))

__call__(x, *, key: PRNGKeyArray | None = None)

jaxnasium.algorithms.QValueNetwork

__init__(obs_space: PyTree[jym.Space], output_space: PyTree[jym.Space], *, key: PRNGKeyArray, body: Callable[..., OutSizedNetwork] = SimBa, obs_architecture_1d: Callable[..., Network] = eqx.nn.Identity, obs_architecture_2d: Callable[..., Network] = _DEFAULT_ARCHITECTURE_2D, output_layer_type: Callable[..., Network] = QLayer.with_params(layer_type=eqx.nn.Linear))

__call__(x, action=None, *, key: PRNGKeyArray | None = None) -> Array | PyTree[Array]

AdvantageNetwork is an alias of QValueNetwork.

Initialization

jaxnasium.algorithms.core.set_weight_bias(key: PRNGKeyArray, network: PyTree[eqx.Module], weight_init: Callable[..., jax.nn.initializers.Initializer] | None = DEFAULT_WEIGHT_INIT, bias_init: float | None = DEFAULT_BIAS_INIT)

Sets all eqx.nn.Linear and eqx.nn.Conv layers in a network to a given weight and bias initialization. Defaults to orthogonal weight initialization and zero bias initialization. Setting weight_init or bias_init to None will leave the weights or biases unchanged.