"""Neural network implementation for BackpropNEAT.""" import jax import jax.numpy as jnp import numpy as np from typing import Dict, List, Optional, Tuple, Union from .genome import Genome import copy import random class Network: """Neural network for NEAT implementation. Implements a strictly feed-forward network following original NEAT principles: 1. Start minimal - direct input-output connections only 2. Complexify gradually through structural mutations 3. Protect innovation through speciation 4. No recurrent connections (as per requirements) """ def __init__(self, genome: Genome): """Initialize network from genome.""" # Store genome and sizes self.genome = genome # Verify genome sizes match volleyball requirements if genome.input_size != 12 or genome.output_size != 3: print(f"Warning: Genome size mismatch. Expected 12 inputs, 3 outputs. Got {genome.input_size} inputs, {genome.output_size} outputs") genome.input_size = 12 genome.output_size = 3 self.input_size = 12 # Fixed for volleyball self.output_size = 3 # Fixed for volleyball # Deep copy to avoid shared references self.node_genes = {} self.connection_genes = [] # Create input nodes (0 to 11) for i in range(12): self.node_genes[i] = NodeGene(i, 'input', 'linear') # Create bias node (12) self.node_genes[12] = NodeGene(12, 'bias', 'linear') # Create output nodes (13, 14, 15) for i in range(3): node_id = 13 + i self.node_genes[node_id] = NodeGene(node_id, 'output', 'sigmoid') # Connect to bias with appropriate weight based on action type if i < 2: # Left/Right actions: encourage movement self.connection_genes.append( ConnectionGene(12, node_id, random.uniform(0.0, 1.0), True) ) else: # Jump action: neutral bias self.connection_genes.append( ConnectionGene(12, node_id, random.uniform(-0.5, 0.5), True) ) # Connect to relevant inputs with larger weights if i == 0: # Left action: connect to ball x position and velocity self.connection_genes.append( ConnectionGene(0, node_id, random.uniform(0.5, 1.5), True) # ball x ) self.connection_genes.append( ConnectionGene(2, node_id, random.uniform(0.5, 1.5), True) # ball vx ) elif i == 1: # Right action: connect to ball x position and velocity self.connection_genes.append( ConnectionGene(0, node_id, random.uniform(-1.5, -0.5), True) # ball x ) self.connection_genes.append( ConnectionGene(2, node_id, random.uniform(-1.5, -0.5), True) # ball vx ) else: # Jump action: connect to ball y position and velocity self.connection_genes.append( ConnectionGene(1, node_id, random.uniform(-1.5, -0.5), True) # ball y ) self.connection_genes.append( ConnectionGene(3, node_id, random.uniform(-1.0, 0.0), True) # ball vy ) # Copy existing nodes (if any) for node_id, node in genome.node_genes.items(): if node_id not in self.node_genes: # Skip I/O nodes self.node_genes[node_id] = NodeGene( node_id, node.node_type, node.activation ) # Copy connections if genome.connection_genes: # Clear initial connections if genome has its own self.connection_genes = [] for conn in genome.connection_genes: # Verify connection nodes exist if conn.source not in self.node_genes or conn.target not in self.node_genes: print(f"Warning: Connection {conn.source}->{conn.target} references missing nodes") continue self.connection_genes.append(ConnectionGene( conn.source, conn.target, conn.weight, conn.enabled )) # Verify output connections (13, 14, 15) for output_id in [13, 14, 15]: has_connection = False for conn in self.connection_genes: if conn.enabled and conn.target == output_id: has_connection = True break if not has_connection: print(f"Adding missing connections for output {output_id}") # Connect to bias self.connection_genes.append( ConnectionGene(12, output_id, random.uniform(-1.0, 1.0), True) ) # Connect to random input input_id = random.randint(0, 11) self.connection_genes.append( ConnectionGene(input_id, output_id, random.uniform(-1.0, 1.0), True) ) # Build evaluation order self.node_evals = {} self._build_feed_forward_order() # Verify all outputs are properly connected self._verify_outputs() def _verify_outputs(self): """Verify all outputs have valid connections and evaluations.""" output_ids = {13, 14, 15} # Fixed output IDs # Check node evaluations for output_id in output_ids: if output_id not in self.node_evals: print(f"Adding missing evaluation for output {output_id}") bias_id = 12 self.node_evals[output_id] = { 'inputs': [bias_id], 'weights': [1.0], 'activation': 'sigmoid' } # Add connection if needed if not any(c.target == output_id and c.enabled for c in self.connection_genes): self.connection_genes.append( ConnectionGene(bias_id, output_id, 1.0, True) ) def _create_minimal_connections(self): """Create minimal initial connections for a new network.""" bias_id = 12 output_start = bias_id + 1 # Connect each output to bias and one random input for i in range(self.output_size): output_id = output_start + i # Connect to bias self.connection_genes.append(ConnectionGene( bias_id, output_id, random.uniform(-1.0, 1.0), True )) # Connect to random input input_id = random.randint(0, self.input_size - 1) self.connection_genes.append(ConnectionGene( input_id, output_id, random.uniform(-1.0, 1.0), True )) def _build_feed_forward_order(self): """Build evaluation order ensuring feed-forward only topology.""" try: # Fixed node sets for volleyball input_nodes = set(range(12)) # 0-11 bias_node = {12} # Bias node output_nodes = {13, 14, 15} # Output nodes # Create adjacency lists connections = {} for conn in self.connection_genes: if not conn.enabled: continue if conn.source not in connections: connections[conn.source] = [] connections[conn.source].append(conn.target) # Start with inputs and bias evaluated evaluated = input_nodes | bias_node eval_order = [] # Helper function to check if a node can be evaluated def can_evaluate(node_id): if node_id in connections: return all(dep in evaluated for dep in connections[node_id]) return True # Keep trying to evaluate nodes until we can't anymore while True: ready_nodes = set() for node_id in self.node_genes: if node_id not in evaluated and can_evaluate(node_id): ready_nodes.add(node_id) if not ready_nodes: break # Add nodes to evaluation order for node_id in sorted(ready_nodes): incoming = [] incoming_weights = [] for conn in self.connection_genes: if conn.enabled and conn.target == node_id: incoming.append(conn.source) incoming_weights.append(conn.weight) if incoming: # Only add if node has inputs self.node_evals[node_id] = { 'inputs': incoming, 'weights': incoming_weights, 'activation': self.node_genes[node_id].activation } eval_order.append(node_id) evaluated.add(node_id) # Ensure all outputs have evaluations for output_id in output_nodes: if output_id not in self.node_evals: print(f"Adding default evaluation for output {output_id}") # Connect to bias by default self.node_evals[output_id] = { 'inputs': [12], # Bias node 'weights': [1.0], 'activation': 'sigmoid' } # Add connection if needed if not any(c.target == output_id and c.enabled for c in self.connection_genes): self.connection_genes.append( ConnectionGene(12, output_id, 1.0, True) ) except Exception as e: print(f"Error in feed-forward build: {e}") # Create minimal fallback evaluations self.node_evals = {} for i in range(3): # 3 outputs output_id = 13 + i self.node_evals[output_id] = { 'inputs': [12], # Bias node 'weights': [1.0], 'activation': 'sigmoid' } def forward(self, inputs: jnp.ndarray) -> jnp.ndarray: """Forward pass through the network.""" try: # Only use first 8 inputs like original network inputs = inputs[:8] # Handle input shape original_shape = inputs.shape if len(inputs.shape) == 1: inputs = inputs.reshape(1, -1) batch_size = inputs.shape[0] # Get max node ID for activation array max_node_id = max(node.id for node in self.node_genes.values()) # Initialize activations array activations = jnp.zeros((batch_size, max_node_id + 1)) # Set input values (0-7) for i in range(8): if i < len(inputs): activations = activations.at[:, i].set(inputs[:, i]) else: activations = activations.at[:, i].set(0.0) # Initialize recurrent nodes (8-11) with previous outputs # For now just use zeros, in the future we could store previous outputs for i in range(8, 12): activations = activations.at[:, i].set(0.0) # Evaluate nodes in order (hidden then output) for node_id, eval_info in self.node_evals.items(): try: # Skip input and recurrent nodes if node_id < 12: continue # Get weighted sum of inputs act = jnp.zeros(batch_size) for conn_source, conn_weight in zip(eval_info['inputs'], eval_info['weights']): act += activations[:, conn_source] * conn_weight # Apply activation function if eval_info['activation'] == 'tanh': act = jnp.tanh(act) elif eval_info['activation'] == 'sigmoid': act = jax.nn.sigmoid(act) elif eval_info['activation'] == 'relu': act = jax.nn.relu(act) # Apply threshold like original network for output nodes if node_id >= 20: # Output nodes act = jnp.where(act > 0.75, 1.0, 0.0) activations = activations.at[:, node_id].set(act) except Exception as e: print(f"Error at node {node_id}: {e}") # Get output node activations output = activations[:, -3:] # Update recurrent nodes for next time step # (In a real implementation, we'd need to store these) for i in range(8, 12): act = jnp.zeros(batch_size) for conn_source, conn_weight in zip(eval_info['inputs'], eval_info['weights']): if conn_source >= 20: # Only use output nodes act += activations[:, conn_source] * conn_weight activations = activations.at[:, i].set(jnp.tanh(act)) # Return to original shape if len(original_shape) == 1: output = output.reshape(-1) return output except Exception as e: print(f"Error in forward pass: {e}") return jnp.zeros(3) def predict(self, inputs: jnp.ndarray) -> jnp.ndarray: """Make a prediction for the given inputs. Args: inputs: Input array of shape (input_size,) or (batch_size, input_size) Returns: Predictions of shape (3,) for single input or (batch_size, 3) for batch """ outputs = self.forward(inputs) # Ensure correct output shape for volleyball (always 3 outputs) if len(outputs.shape) == 1: # Single input case - ensure shape (3,) if outputs.shape[0] != 3: print(f"Adjusting output shape from {outputs.shape} to (3,)") return jnp.pad(outputs, (0, max(0, 3 - outputs.shape[0]))) return outputs else: # Batch case - ensure shape (batch_size, 3) if outputs.shape[1] != 3: print(f"Adjusting output shape from {outputs.shape} to (batch_size, 3)") return jnp.pad(outputs, ((0, 0), (0, max(0, 3 - outputs.shape[1])))) return outputs def clone(self) -> 'Network': """Create a copy of this network with a cloned genome.""" return Network(self.genome.clone()) def mutate(self, config: Dict): """Mutate the network's genome.""" self.genome.mutate(config) # Rebuild evaluation order after mutation self._build_feed_forward_order() def to_genome(self) -> Genome: """Convert network back to genome representation.""" genome = Genome(self.input_size, self.output_size) genome.node_genes = copy.deepcopy(self.node_genes) genome.connection_genes = copy.deepcopy(self.connection_genes) return genome class BaseNetwork: """Base Network class for NEAT.""" def __init__(self, n_inputs: int, n_outputs: int): self.input_size = n_inputs self.output_size = n_outputs self.fitness = float('-inf') # Initialize weights and biases with JAX key = jax.random.PRNGKey(0) # Use larger initial weights to encourage exploration self.weights = jax.random.normal(key, (n_outputs, n_inputs)) * 0.5 # Add small positive bias to encourage some initial movement self.bias = jnp.ones(n_outputs) * 0.1 def forward(self, x: jnp.ndarray) -> jnp.ndarray: """Forward pass through the network.""" if x.ndim > 1: # Batched input h = jnp.dot(x, self.weights.T) + self.bias[None, :] else: # Single input h = jnp.dot(x, self.weights.T) + self.bias return jnp.tanh(h) def get_params(self) -> Tuple[jnp.ndarray, jnp.ndarray]: """Get network parameters.""" return self.weights, self.bias def set_params(self, params: Tuple[jnp.ndarray, jnp.ndarray]): """Set network parameters.""" self.weights, self.bias = params def get_weights_numpy(self) -> np.ndarray: """Get weights as numpy array for visualization.""" return np.array(self.weights) class NodeGene: """Node gene containing node information.""" def __init__(self, node_id: int, node_type: str, activation: str = 'tanh'): """Initialize node gene. Args: node_id: Node ID node_type: Type of node ('input', 'hidden', or 'output') activation: Activation function ('tanh', 'sigmoid', or 'relu') """ self.id = node_id self.type = node_type self.activation = activation # Initialize with larger random bias for hidden/output nodes if node_type in ['hidden', 'output']: key = jax.random.PRNGKey(node_id) # Use node_id as seed for reproducibility self.bias = jax.random.normal(key, ()) * 0.5 # Increased from 0.1 else: self.bias = 0.0 # No bias for input nodes class ConnectionGene: """Gene representing a connection between nodes.""" def __init__(self, source: int, target: int, weight: float = None, enabled: bool = True): self.source = source self.target = target # Initialize with larger weights if not provided if weight is None: key = jax.random.PRNGKey(hash((source, target)) % 2**32) self.weight = jax.random.uniform(key, (), minval=-2.0, maxval=2.0) else: self.weight = weight self.enabled = enabled self.innovation = None # Will be set by NEAT