Upload neat\genome.py with huggingface_hub

neat//genome.py

@@ -0,0 +1,454 @@
+"""NEAT Genome implementation.
+
+This module implements the core NEAT genome structure and operations.
+Each genome represents a neural network with nodes (neurons) and connections (synapses).
+The genome can be mutated to evolve the network structure and weights over time.
+"""
+
+from dataclasses import dataclass
+import jax.numpy as jnp
+import jax.random as jrandom
+from typing import Dict, List, Tuple, Optional
+import time
+import random
+import numpy as np
+
+@dataclass
+class NodeGene:
+    """Node gene containing activation function and type.
+    
+    Attributes:
+        node_id: Unique identifier for this node
+        node_type: Type of node ('input', 'hidden', 'recurrent', or 'output')
+        activation: Activation function ('tanh', 'relu', 'sigmoid', or 'linear')
+    """
+    node_id: int
+    node_type: str  # 'input', 'hidden', 'recurrent', or 'output'
+    activation: str  # 'tanh', 'relu', 'sigmoid', or 'linear'
+
+@dataclass
+class ConnectionGene:
+    """Connection gene containing connection properties.
+    
+    Attributes:
+        source: ID of source node
+        target: ID of target node
+        weight: Connection weight
+        enabled: Whether connection is enabled
+        innovation: Unique innovation number for this connection
+    """
+    source: int
+    target: int
+    weight: float
+    enabled: bool = True
+    innovation: int = 0
+
+class Genome:
+    """NEAT Genome implementation.
+    
+    A genome represents a neural network as a collection of node and connection genes.
+    The network topology can be modified through mutation operations.
+    
+    Attributes:
+        input_size: Number of input nodes
+        output_size: Number of output nodes
+        node_genes: Dictionary mapping node IDs to NodeGene objects
+        connection_genes: List of ConnectionGene objects
+        key: Random key for reproducible randomness
+        innovation_number: Counter for assigning unique innovation numbers
+    """
+    
+    def __init__(self, input_size: int, output_size: int):
+        """Initialize genome with specified number of inputs and outputs.
+        
+        Args:
+            input_size: Number of input nodes
+            output_size: Number of output nodes (must be 3 for volleyball)
+        """
+        self.input_size = input_size
+        self.output_size = output_size
+        self.node_genes: Dict[int, NodeGene] = {}
+        self.connection_genes: List[ConnectionGene] = []
+        
+        # Initialize random key
+        timestamp = int(time.time() * 1000)
+        self.key = jrandom.PRNGKey(hash((input_size, output_size, timestamp)) % (2**32))
+        
+        # Counter for assigning unique innovation numbers
+        self.innovation_number = 0
+        
+        # Initialize minimal network structure
+        self._init_minimal()
+    
+    def _init_minimal(self):
+        """Initialize minimal feed-forward network structure.
+        
+        Network structure:
+        - Input nodes [0-7]: Game state inputs
+        - Hidden layer 1 [8-15]: First processing layer (8 nodes)
+        - Hidden layer 2 [16-23]: Second processing layer (8 nodes)
+        - Output nodes [24-26]: Action outputs (left, right, jump)
+        
+        Using larger initial weights for faster learning:
+        - Input->Hidden1: N(0, 2.0) for strong initial responses
+        - Hidden1->Hidden2: N(0, 2.0) for feature processing
+        - Hidden2->Output: N(0, 4.0) for decisive actions
+        """
+        # Create input nodes (0-7)
+        for i in range(8):  # Only 8 inputs used
+            self.node_genes[i] = NodeGene(
+                node_id=i,
+                node_type='input',
+                activation='linear'  # Input nodes are always linear
+            )
+        
+        # Create first hidden layer (8-15)
+        hidden1_size = 8
+        hidden1_start = 8  # Right after inputs
+        for i in range(hidden1_size):
+            node_id = hidden1_start + i
+            self.node_genes[node_id] = NodeGene(
+                node_id=node_id,
+                node_type='hidden',
+                activation='relu'  # ReLU for faster learning
+            )
+            
+            # Connect all inputs to this hidden node
+            for input_id in range(8):
+                weight = float(jrandom.normal(self.key) * 2.0)
+                self.connection_genes.append(ConnectionGene(
+                    source=input_id,
+                    target=node_id,
+                    weight=weight,
+                    enabled=True,
+                    innovation=self.innovation_number
+                ))
+                self.innovation_number += 1
+        
+        # Create second hidden layer (16-23)
+        hidden2_size = 8
+        hidden2_start = hidden1_start + hidden1_size
+        for i in range(hidden2_size):
+            node_id = hidden2_start + i
+            self.node_genes[node_id] = NodeGene(
+                node_id=node_id,
+                node_type='hidden',
+                activation='relu'  # ReLU for faster learning
+            )
+            
+            # Connect all hidden1 nodes to this hidden2 node
+            for h1_id in range(hidden1_start, hidden1_start + hidden1_size):
+                weight = float(jrandom.normal(self.key) * 2.0)
+                self.connection_genes.append(ConnectionGene(
+                    source=h1_id,
+                    target=node_id,
+                    weight=weight,
+                    enabled=True,
+                    innovation=self.innovation_number
+                ))
+                self.innovation_number += 1
+        
+        # Create output nodes (24-26)
+        output_start = hidden2_start + hidden2_size
+        for i in range(self.output_size):
+            node_id = output_start + i
+            self.node_genes[node_id] = NodeGene(
+                node_id=node_id,
+                node_type='output',
+                activation='tanh'  # tanh for [-1,1] outputs
+            )
+            
+            # Connect all hidden2 nodes to this output
+            for h2_id in range(hidden2_start, hidden2_start + hidden2_size):
+                weight = float(jrandom.normal(self.key) * 4.0)  # Larger weights for outputs
+                self.connection_genes.append(ConnectionGene(
+                    source=h2_id,
+                    target=node_id,
+                    weight=weight,
+                    enabled=True,
+                    innovation=self.innovation_number
+                ))
+                self.innovation_number += 1
+    
+    def mutate(self, config: Dict):
+        """Mutate the genome by modifying weights and network structure.
+        
+        Args:
+            config: Dictionary containing mutation parameters:
+                - weight_mutation_rate: Probability of mutating each weight
+                - weight_mutation_power: Standard deviation for weight mutations
+                - add_node_rate: Probability of adding a new node
+                - add_connection_rate: Probability of adding a new connection
+        """
+        # Mutate connection weights
+        for conn in self.connection_genes:
+            if jrandom.uniform(self.key) < config['weight_mutation_rate']:
+                # Get new random key
+                self.key, subkey = jrandom.split(self.key)
+                # Add random value from normal distribution
+                conn.weight += float(jrandom.normal(subkey) * config['weight_mutation_power'])
+        
+        # Add new nodes (disabled for now since we're using fixed topology)
+        if config['add_node_rate'] > 0:
+            if jrandom.uniform(self.key) < config['add_node_rate']:
+                self._add_node()
+        
+        # Add new connections (disabled for now)
+        if config['add_connection_rate'] > 0:
+            if jrandom.uniform(self.key) < config['add_connection_rate']:
+                self._add_connection()
+    
+    def _add_node(self):
+        """Add a new node by splitting an existing connection."""
+        if not self.connection_genes:
+            return
+            
+        # Choose a random connection to split
+        conn_to_split = np.random.choice(self.connection_genes)
+        conn_to_split.enabled = False
+        
+        # Create new node
+        new_node_id = max(self.node_genes.keys()) + 1
+        self.node_genes[new_node_id] = NodeGene(
+            node_id=new_node_id,
+            node_type='hidden',
+            activation='relu'
+        )
+        
+        # Create two new connections
+        self.connection_genes.extend([
+            ConnectionGene(
+                source=conn_to_split.source,
+                target=new_node_id,
+                weight=1.0,
+                enabled=True,
+                innovation=self.innovation_number
+            ),
+            ConnectionGene(
+                source=new_node_id,
+                target=conn_to_split.target,
+                weight=conn_to_split.weight,
+                enabled=True,
+                innovation=self.innovation_number + 1
+            )
+        ])
+        self.innovation_number += 2
+    
+    def _add_connection(self):
+        """Add a new connection between two unconnected nodes."""
+        # Get list of all possible connections
+        existing_connections = {(c.source, c.target) for c in self.connection_genes}
+        possible_connections = []
+        
+        for source in self.node_genes:
+            for target in self.node_genes:
+                # Skip if connection already exists
+                if (source, target) in existing_connections:
+                    continue
+                    
+                # Skip if would create cycle (except recurrent)
+                if self.node_genes[source].node_type != 'recurrent' and \
+                   self.would_create_cycle(source, target):
+                    continue
+                    
+                possible_connections.append((source, target))
+        
+        if possible_connections:
+            # Choose random connection
+            source, target = random.choice(possible_connections)
+            
+            # Create new connection
+            weight = float(jrandom.normal(self.key) * 1.0)
+            self.connection_genes.append(ConnectionGene(
+                source=source,
+                target=target,
+                weight=weight,
+                enabled=True,
+                innovation=self.innovation_number
+            ))
+            self.innovation_number += 1
+    
+    def would_create_cycle(self, source: int, target: int) -> bool:
+        """Check if adding connection would create cycle in network.
+        
+        Args:
+            source: Source node ID
+            target: Target node ID
+            
+        Returns:
+            True if connection would create cycle, False otherwise
+        """
+        # Skip cycle detection for recurrent connections
+        if self.node_genes[source].node_type == 'recurrent' or \
+           self.node_genes[target].node_type == 'recurrent':
+            return False
+            
+        # Do depth-first search from target to see if we can reach source
+        visited = set()
+        
+        def dfs(node: int) -> bool:
+            if node == source:
+                return True
+            if node in visited:
+                return False
+                
+            visited.add(node)
+            for conn in self.connection_genes:
+                if conn.source == node and conn.enabled:
+                    if dfs(conn.target):
+                        return True
+            return False
+        
+        return dfs(target)
+    
+    def add_node_between(self, source: int, target: int):
+        """Add a new node between two nodes, splitting an existing connection.
+        
+        Args:
+            source: Source node ID
+            target: Target node ID
+        """
+        # Find and disable the existing connection
+        for conn in self.connection_genes:
+            if conn.source == source and conn.target == target and conn.enabled:
+                conn.enabled = False
+                
+                # Create new node
+                new_id = max(self.node_genes.keys()) + 1
+                self.node_genes[new_id] = NodeGene(
+                    node_id=new_id,
+                    node_type='hidden',
+                    activation='relu'
+                )
+                
+                # Create two new connections
+                self.connection_genes.extend([
+                    ConnectionGene(
+                        source=source,
+                        target=new_id,
+                        weight=1.0,
+                        enabled=True,
+                        innovation=self.innovation_number
+                    ),
+                    ConnectionGene(
+                        source=new_id,
+                        target=target,
+                        weight=conn.weight,
+                        enabled=True,
+                        innovation=self.innovation_number + 1
+                    )
+                ])
+                self.innovation_number += 2
+                break
+    
+    def add_connection(self, source: int, target: int, weight: Optional[float] = None) -> bool:
+        """Add a new connection between two nodes.
+        
+        Args:
+            source: Source node ID
+            target: Target node ID
+            weight: Optional connection weight. If None, a random weight is generated.
+            
+        Returns:
+            True if connection was added, False if invalid or already exists
+        """
+        # Check if connection already exists
+        if any(c.source == source and c.target == target for c in self.connection_genes):
+            return False
+            
+        # Validate nodes exist
+        if source not in self.node_genes or target not in self.node_genes:
+            return False
+            
+        # Ensure feed-forward (no cycles)
+        if source >= target:  # Simple way to ensure feed-forward
+            return False
+            
+        # Generate random weight if not provided
+        if weight is None:
+            weight = float(jrandom.normal(self.key) * 1.0)
+            
+        # Add new connection
+        self.connection_genes.append(ConnectionGene(
+            source=source,
+            target=target,
+            weight=weight,
+            enabled=True,
+            innovation=self.innovation_number
+        ))
+        self.innovation_number += 1
+        return True
+    
+    def crossover(self, other: 'Genome', key: jnp.ndarray) -> 'Genome':
+        """Perform crossover between two genomes.
+        
+        Args:
+            other: Other parent genome
+            key: JAX PRNG key
+            
+        Returns:
+            Child genome
+        """
+        # Create child genome
+        child = Genome(self.input_size, self.output_size)
+        
+        # Inherit node genes
+        for node_id in self.node_genes:
+            if node_id in other.node_genes:
+                # Inherit randomly from either parent
+                if jrandom.uniform(key) < 0.5:
+                    child.node_genes[node_id] = self.node_genes[node_id]
+                else:
+                    child.node_genes[node_id] = other.node_genes[node_id]
+            else:
+                # Inherit from fitter parent
+                child.node_genes[node_id] = self.node_genes[node_id]
+        
+        # Inherit connection genes
+        for conn in self.connection_genes:
+            if conn.innovation in [c.innovation for c in other.connection_genes]:
+                # Inherit randomly from either parent
+                other_conn = next(c for c in other.connection_genes if c.innovation == conn.innovation)
+                if jrandom.uniform(key) < 0.5:
+                    child.connection_genes.append(ConnectionGene(
+                        source=conn.source,
+                        target=conn.target,
+                        weight=conn.weight,
+                        enabled=conn.enabled,
+                        innovation=conn.innovation
+                    ))
+                else:
+                    child.connection_genes.append(ConnectionGene(
+                        source=other_conn.source,
+                        target=other_conn.target,
+                        weight=other_conn.weight,
+                        enabled=other_conn.enabled,
+                        innovation=other_conn.innovation
+                    ))
+            else:
+                # Inherit from fitter parent
+                child.connection_genes.append(ConnectionGene(
+                    source=conn.source,
+                    target=conn.target,
+                    weight=conn.weight,
+                    enabled=conn.enabled,
+                    innovation=conn.innovation
+                ))
+        
+        return child
+        
+    def clone(self) -> 'Genome':
+        """Create a copy of this genome.
+        
+        Returns:
+            Copy of genome
+        """
+        clone = Genome(self.input_size, self.output_size)
+        clone.node_genes = self.node_genes.copy()
+        clone.connection_genes = [ConnectionGene(**conn.__dict__) for conn in self.connection_genes]
+        return clone
+
+    @property
+    def n_nodes(self) -> int:
+        """Get total number of nodes in the genome."""
+        return len(self.node_genes)