PyTree Layers¤
The pieces that make agent networks work on composite (PyTree) observation and action spaces. The input/output modules build one sub-network per space in the tree; the output layers turn a body's features into a distribution or Q-values for a single space.
Input¤
jaxnasium.algorithms.core.PyTreeObsSpaceNetwork
¤
Builds a separate observation network for each observation space.
Automatically builds a given 1d architecture for 1d observation spaces and a given 2d architecture for 2d observation spaces.
During a forward call this network simply returns a jax.tree.map over all observation spaces and concatenates the outputs of all observation networks as a single 1d vector.
Arguments:
- obs_space: The observation space(s) to build networks for.
- key: A PRNG key for reproducibility.
- architecture_1d: The architecture to use for 1d observation spaces.
- architecture_2d: The architecture to use for 2d observation spaces.
__init__(obs_space: PyTree[SpaceLike], *, key: PRNGKeyArray, architecture_1d: Callable[..., Network] = Identity, architecture_2d: Callable[..., Network] = CNN.with_params(out_channels=(32, 64, 64), kernel_sizes=(3, 3, 2), strides=(1, 1, 1), padding=(0, 0, 0)))
¤
__call__(x, *, key: PRNGKeyArray | None = None)
¤
with_params(*, architecture_1d: Callable[..., Network] = Identity, architecture_2d: Callable[..., Network] = CNN.with_params(out_channels=(32, 64, 64), kernel_sizes=(3, 3, 2), strides=(1, 1, 1), padding=(0, 0, 0))) -> Callable[..., Self]
classmethod
¤
Output¤
jaxnasium.algorithms.core.PyTreeActionNetwork
¤
Output network producing an action distribution per output space.
Builds one head per output-space leaf, selecting a discrete or continuous
head automatically from the space and returning a
distrax.Distribution (policy head).
A single head may itself be multi-dimensional (homogeneous), in which case its sub-layers are stacked and vmapped in the forward call.
During a forward call this maps over all output spaces. When every head
returns a distribution and assume_independent is set, per-space
distributions are wrapped in Independent and, for multiple spaces,
combined into a distrax.Joint. Otherwise the raw PyTree of outputs is
returned unchanged.
__init__(in_features: int, output_space: PyTree[SpaceLike], *, key: PRNGKeyArray, discrete_output_layer: Callable[..., Network] = CategoricalLayer, continuous_output_layer: Callable[..., Network] = NormalLayer, assume_independent: bool = True)
¤
__call__(x, action_mask=None, *, key: PRNGKeyArray | None = None)
¤
with_params(*, discrete_output_layer: Callable[..., Network] = CategoricalLayer, continuous_output_layer: Callable[..., Network] = NormalLayer, assume_independent: bool = True) -> Callable[..., Self]
classmethod
¤
jaxnasium.algorithms.core.PyTreeQValueNetwork
¤
Output network producing raw Q-values per output space (no distribution).
Builds one head per output-space leaf, selecting a discrete or continuous head automatically from the space and returning raw values (Q-value head).
A single head may itself be multi-dimensional (homogeneous), in which case its sub-layers are stacked and vmapped in the forward call.
During a forward call this maps over all output spaces. The raw PyTree of outputs is returned unchanged.
Output layers¤
These are the per-space heads selected by the modules above, and can be swapped through the
discrete_output_layer / continuous_output_layer / output_layer_type arguments of the
agent networks.
jaxnasium.algorithms.core.CategoricalLayer
¤
jaxnasium.algorithms.core.NormalLayer
¤
jaxnasium.algorithms.core.TanhNormalLayer
¤
jaxnasium.algorithms.core.QLayer
¤
Latent -> raw values (no distribution).
Mode is inferred from the output space:
- discrete: one value per action, i.e.
Q(s, .). - continuous: a single scalar
Q(s, a), which requires the action as input.