Skip to content

zymera.nets

nets

Composable agent building blocks (the parts you wire into a policy).

Raw-JAX and functional so blocks compose freely and stay jit/vmap-safe. Start minimal: this seeds the module with one foundational block (an MLP). Add encoders, belief nets, attention, aggregators, low-level controllers here as experiments need them; graduate proven blocks in with a test.

mlp_init 

mlp_init(key, sizes)

Glorot-init params for an MLP with the given layer sizes.

Source code in zymera/nets.py 
16
17
18
19
20
21
22
23
24
25
def mlp_init(key: jax.Array, sizes: Sequence[int]):
    """Glorot-init params for an MLP with the given layer sizes."""
    params = []
    keys = jax.random.split(key, len(sizes) - 1)
    for k, d_in, d_out in zip(keys, sizes[:-1], sizes[1:]):
        scale = jnp.sqrt(2.0 / (d_in + d_out))
        w = jax.random.normal(k, (d_in, d_out)) * scale
        b = jnp.zeros((d_out,))
        params.append((w, b))
    return params

mlp_apply 

mlp_apply(params, x, activation=jax.nn.relu)

Apply the MLP; activation on every layer except the last.

Source code in zymera/nets.py 
28
29
30
31
32
33
34
def mlp_apply(params, x: jax.Array, activation=jax.nn.relu) -> jax.Array:
    """Apply the MLP; activation on every layer except the last."""
    for i, (w, b) in enumerate(params):
        x = x @ w + b
        if i < len(params) - 1:
            x = activation(x)
    return x