Merge pull request #6 from RPegoud/5-fix-add_node-mutation-bug

5 fix add node mutation bug

Author: RPegoud

README.md

@@ -26,8 +26,16 @@ Mutations:
   * [x] Weight reset
   * [x] Add node
   * [x] Add connection
+  * [ ] Mutate activation
 * [ ] Wrap all mutations in a single function
 
+Misc:
+
+* [ ] Add Hydra config for constant attributes
+* [ ] Separate ``max_nodes`` and ``max_connections``
+* [ ] Add bias
+* [ ] Set the minimum sender index to 1 instead of 0
+
 Crossing:
 
 * [ ] Add novelty fields to Network dataclass

main.py

@@ -0,0 +1,40 @@
+import hydra
+import jax
+from omegaconf import DictConfig, OmegaConf
+
+from neat_jax import Mutations, forward, init_network, log_config
+
+INPUT_SIZE = 4
+OUTPUT_SIZE = 3
+
+
+def run_exp(config: dict):
+    config.input_size = INPUT_SIZE
+    config.output_size = OUTPUT_SIZE
+
+    key = jax.random.key(config.params.seed)
+    inputs = jax.random.normal(key, (config.input_size,))
+
+    net, activation_state = init_network(
+        inputs, config.output_size, config.network.max_nodes, key
+    )
+    mutations = Mutations(max_nodes=config.network.max_nodes, **config.mutations)
+    net = mutations.add_node(key, net)
+    activation_state, y = forward(inputs, net, config)
+
+    print(activation_state, y)
+
+
+@hydra.main(
+    config_path="neat_jax/configs",
+    config_name="default_config.yaml",
+    version_base="1.3.2",
+)
+def hydra_entry_point(cfg: DictConfig):
+    OmegaConf.set_struct(cfg, False)
+    log_config(cfg)
+    return run_exp(cfg)
+
+
+if __name__ == "__main__":
+    hydra_entry_point()

neat_jax/__init__.py

@@ -1,13 +1,14 @@
+from .activation_fns import activation_fns_list, get_activation_fn
+from .initialization import activation_state_from_inputs, init_network
 from .mutations import Mutations
 from .neat_dataclasses import ActivationState, Network
 from .nn import (
     forward,
     forward_toggled_nodes,
-    get_activation,
+    get_activation_fn,
     get_active_connections,
     get_required_activations,
-    make_network,
     toggle_receivers,
     update_depth,
 )
-from .utils import plot_network, sample_from_mask
+from .utils import cartesian_product, log_config, plot_network, sample_from_mask

neat_jax/__pycache__/__init__.cpython-310.pyc

@@ -0,0 +1,37 @@
+import jax
+import jax.numpy as jnp
+
+activation_fns_list = [
+    lambda x: jnp.float32(x),  # identity
+    lambda x: 1 / (1 + jnp.exp(-x)),  # sigmoid
+    lambda x: jnp.divide(1, x),  # inverse
+    lambda x: jnp.sinh(x) / jnp.cosh(x),  # hyperbolic cosine
+    lambda x: jnp.float32(jnp.maximum(0, x)),  # relu
+    lambda x: jnp.float32(jnp.abs(x)),  # absolute value
+    lambda x: jnp.sin(x),  # sine
+    lambda x: jnp.exp(jnp.square(-x)),  # gaussian
+    lambda x: jnp.float32(jnp.sign(x)),  # step
+]
+
+
+def get_activation_fn(activation_index: int, x: float) -> jnp.float32:
+    """
+    Given an index, selects an activation function and computes `activation(x)`.
+
+    ```python
+        0: jnp.float32(x) # identity function
+        1: 1 / (1 + jnp.exp(-x)),  # sigmoid
+        2: jnp.divide(1, x),  # inverse
+        3: jnp.sinh(x) / jnp.cosh(x),  # hyperbolic cosine
+        4: jnp.float32(jnp.maximum(0, x)),  # relu
+        5: jnp.float32(jnp.abs(x)),  # absolute value
+        6: jnp.sin(x),  # sine
+        7: jnp.exp(jnp.square(-x)),  # gaussian
+        8: jnp.float32(jnp.sign(x)),  # step
+        ```
+    """
+    return jax.lax.switch(
+        activation_index,
+        activation_fns_list,
+        operand=x,
+    )

neat_jax/__pycache__/activation_fns.cpython-310.pyc

@@ -0,0 +1,14 @@
+defaults:
+  - params: base_params
+  - network: base_net
+  - mutations: mutations
+  - env: gymnax/cartpole
+  - _self_
+
+  - override hydra/hydra_logging: disabled
+  - override hydra/job_logging: disabled
+
+hydra:
+  output_subdir: null
+  run:
+    dir: .

neat_jax/__pycache__/initialization.cpython-310.pyc

@@ -0,0 +1,7 @@
+suit_name: gymnax
+
+scenario:
+  name: CartPole-v1
+  task_name: cartpole
+
+kwargs: {}

neat_jax/__pycache__/mutations.cpython-310.pyc

@@ -0,0 +1,5 @@
+weight_shift_rate: 0.5
+weight_mutation_rate: 0.1
+add_node_rate: 0.03
+add_connection_rate: 0.05
+activation_fn_mutation_rate: 0.1

neat_jax/__pycache__/neat_dataclasses.cpython-310.pyc

@@ -0,0 +1,2 @@
+max_nodes: 30
+max_connections: 45

neat_jax/__pycache__/nn.cpython-310.pyc

@@ -0,0 +1,3 @@
+seed: 0 # RNG seed.
+population_size: 1024
+total_generations: 1e7

neat_jax/__pycache__/utils.cpython-310.pyc

@@ -0,0 +1,125 @@
+import chex
+import jax
+import jax.numpy as jnp
+
+from .neat_dataclasses import ActivationState, Network
+from .utils.utils import cartesian_product
+
+
+def get_initial_activations(inputs: chex.Array, senders: chex.Array) -> chex.Array:
+    """Initializes ActivationState.values based on an array of inputs."""
+    default_values = jnp.zeros_like(senders)
+    keys = jnp.arange(len(inputs))
+    matches = jnp.searchsorted(keys, senders)
+    mask = (matches < len(keys)) & (keys[matches] == senders)
+    return jnp.where(mask, inputs[matches], default_values)
+
+
+def create_toggled_mask(
+    node_types: chex.Array, senders: chex.Array, input_size: int
+) -> chex.Array:
+    """Initializes the ActivationState.toggled mask based on input nodes"""
+    input_nodes_indices = jnp.where(node_types == 0, size=input_size)[0]
+    sender_mask = jnp.isin(senders, input_nodes_indices)
+    return jnp.int32(sender_mask)
+
+
+def activation_state_from_inputs(
+    inputs: chex.Array,
+    senders: chex.Array,
+    node_types: chex.Array,
+    input_size: int,
+    max_nodes: int,
+) -> "ActivationState":
+    """
+    Resets the ActivationState in prevision of a forward pass or a depth scan.
+
+    Args:
+        inputs (chex.Array): The activation values of the network's input nodes
+        max_nodes (int): The maximum capacity of the network
+
+    Returns:
+        ActivationState: The reset ActivationState with:
+
+            - ``values``: initialized based on inputs
+            - ``toggled``: input neurons toggled
+            - ``activation_counts``: set to zero
+            - ``has_fired``: set to zero
+            - ``outdated_depths``: True
+    """
+
+    values = get_initial_activations(inputs, senders)
+    toggled = create_toggled_mask(node_types, senders, input_size)
+
+    return ActivationState(
+        values=values,
+        toggled=toggled,
+        activation_counts=jnp.zeros(max_nodes, dtype=jnp.int32),
+        has_fired=jnp.zeros(max_nodes, dtype=jnp.int32),
+        node_depths=jnp.zeros(max_nodes, dtype=jnp.int32),
+        outdated_depths=True,
+    )
+
+
+def init_network(
+    inputs: chex.Array,
+    output_size: int,
+    max_nodes: int,
+    key: chex.PRNGKey,
+    scale_weights: float = 0.1,
+) -> tuple[Network, ActivationState]:
+    """Creates a Network and ActivationState from an input array."""
+
+    input_size = len(inputs)
+    n_initial_connections = input_size * output_size
+    sender_receiver_pairs = cartesian_product(
+        jnp.arange(input_size),
+        jnp.arange(input_size, input_size + output_size),
+        size=max_nodes,
+        fill_value=-max_nodes,
+    )
+    senders = sender_receiver_pairs[:, 0]
+    receivers = sender_receiver_pairs[:, 1]
+
+    weights_init = (
+        jax.random.normal(key, (n_initial_connections,), dtype=jnp.float32)
+        * scale_weights
+    )
+    weights = jnp.zeros(max_nodes).at[:n_initial_connections].set(weights_init)
+
+    node_types = jnp.concatenate(
+        [
+            jnp.zeros(input_size, dtype=jnp.int32),  # input nodes = 0
+            jnp.full(output_size, 2, dtype=jnp.int32),  # output nodes = 2
+            jnp.full(max_nodes - input_size - output_size, 3, dtype=jnp.int32),
+        ]  # disabled nodes = 3
+    )
+    activation_fns = jnp.zeros(max_nodes, dtype=jnp.int32)  # no activation by default
+
+    activations = get_initial_activations(inputs, senders)
+    toggled_nodes = create_toggled_mask(node_types, senders, input_size)
+    activation_counts = jnp.zeros(max_nodes, dtype=jnp.int32)
+    activation_counts = jnp.zeros(max_nodes, dtype=jnp.int32)
+    has_fired = jnp.zeros(max_nodes, dtype=jnp.int32)
+
+    return (
+        Network(
+            node_indices=jnp.arange(max_nodes, dtype=jnp.int32),
+            node_types=node_types,
+            weights=weights,
+            activation_fns=activation_fns,
+            senders=senders,
+            receivers=receivers,
+            input_size=input_size,
+            output_size=output_size,
+            max_nodes=max_nodes,
+        ),
+        ActivationState(
+            values=activations,
+            toggled=toggled_nodes,
+            activation_counts=activation_counts,
+            has_fired=has_fired,
+            node_depths=jnp.zeros(max_nodes, dtype=jnp.int32),
+            outdated_depths=True,
+        ),
+    )