Upload neat\evolution.py with huggingface_hub

neat//evolution.py

@@ -0,0 +1,302 @@
+"""NEAT evolution implementation."""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from typing import List, Dict, Optional, Tuple, Callable
+from .network import Network
+from .genome import Genome
+
+class NEATEvolution:
+    """NEAT evolution implementation with structural mutations."""
+    
+    DEFAULT_CONFIG = {
+        'node_add_prob': 0.2,        # Standard node addition rate
+        'conn_add_prob': 0.3,        # Standard connection addition rate
+        'weight_mutate_prob': 0.8,   # High chance of weight mutation
+        'weight_replace_prob': 0.1,   # Low chance of complete weight replacement
+        'weight_perturb_size': 0.5,  # Standard weight perturbation size
+        'bias_mutate_prob': 0.8,     # High chance of bias mutation
+        'bias_replace_prob': 0.1,     # Low chance of complete bias replacement
+        'bias_perturb_size': 0.5,    # Standard bias perturbation size
+        'complexity_coefficient': 0.0,  # No complexity penalty
+        'species_distance': 2.0,      # Standard species distance
+        'species_elitism': 2,         # Keep top 2 from each species
+        'survival_threshold': 0.3     # Keep 30% of population
+    }
+    
+    def __init__(self, 
+                 n_inputs: int,
+                 n_outputs: int, 
+                 population_size: int,
+                 config: Optional[Dict] = None,
+                 key: Optional[jnp.ndarray] = None):
+        """Initialize NEAT evolution.
+        
+        Args:
+            n_inputs: Number of input nodes (12 for volleyball)
+            n_outputs: Number of output nodes (3 for volleyball)
+            population_size: Size of population
+            config: Optional configuration parameters
+            key: Random key for JAX
+        """
+        self.n_inputs = n_inputs
+        self.n_outputs = n_outputs
+        self.population_size = population_size
+        self.config = {**self.DEFAULT_CONFIG, **(config or {})}
+        
+        # Initialize random key
+        if key is None:
+            self.key = jax.random.PRNGKey(0)
+        else:
+            self.key = key
+            
+        # Initialize population
+        self.population = self._init_population()
+        self.generation = 0
+        self.innovation_number = 0
+        self.species = []
+        
+    def _init_population(self) -> List[Genome]:
+        """Initialize population with minimal networks."""
+        population = []
+        for _ in range(self.population_size):
+            # Split random key
+            self.key, subkey = jax.random.split(self.key)
+            
+            # Create genome with proper input/output sizes
+            genome = Genome(self.n_inputs, self.n_outputs, subkey)
+            
+            # Add random hidden nodes (between 2-6)
+            self.key, subkey = jax.random.split(self.key)
+            n_hidden = int(jax.random.randint(subkey, (), 2, 7))
+            
+            hidden_nodes = []
+            for _ in range(n_hidden):
+                hidden_nodes.append(genome.add_node())
+            
+            # Connect inputs to hidden with 50% probability
+            for i in range(self.n_inputs):
+                for h in hidden_nodes:
+                    self.key, subkey = jax.random.split(self.key)
+                    if jax.random.uniform(subkey) < 0.5:
+                        self.key, subkey = jax.random.split(self.key)
+                        weight = jax.random.normal(subkey) * 0.5
+                        genome.add_connection(i, h, weight)
+            
+            # Connect hidden to outputs with 50% probability
+            output_start = genome.n_nodes - self.n_outputs
+            for h in hidden_nodes:
+                for i in range(self.n_outputs):
+                    self.key, subkey = jax.random.split(self.key)
+                    if jax.random.uniform(subkey) < 0.5:
+                        self.key, subkey = jax.random.split(self.key)
+                        weight = jax.random.normal(subkey) * 0.5
+                        genome.add_connection(h, output_start + i, weight)
+            
+            # Add skip connections with 30% probability
+            for i in range(self.n_inputs):
+                for j in range(self.n_outputs):
+                    self.key, subkey = jax.random.split(self.key)
+                    if jax.random.uniform(subkey) < 0.3:
+                        self.key, subkey = jax.random.split(self.key)
+                        weight = jax.random.normal(subkey) * 0.3
+                        genome.add_connection(i, output_start + j, weight)
+            
+            population.append(genome)
+        return population
+        
+    def ask(self) -> List[Network]:
+        """Get current population as networks."""
+        return [Network(genome) for genome in self.population]
+    
+    def tell(self, fitnesses: List[float]) -> None:
+        """Update population based on fitness scores."""
+        # Sort population by fitness
+        sorted_pop = sorted(zip(self.population, fitnesses), 
+                          key=lambda x: x[1], reverse=True)
+        
+        # For very small populations, keep at least one parent
+        n_parents = max(1, int(self.population_size * self.config['survival_threshold']))
+        parents = [p for p, _ in sorted_pop[:n_parents]]
+        
+        # Ensure we have at least one parent
+        if not parents:
+            # If all fitnesses are equal (including all zeros), keep the first one
+            parents = [sorted_pop[0][0]]
+        
+        # Create new population starting with the best performer
+        new_population = [parents[0]]  # Always keep the best one
+        
+        # Fill rest with mutated offspring
+        while len(new_population) < self.population_size:
+            # Select parent (with replacement)
+            parent = parents[0] if len(parents) == 1 else np.random.choice(parents)
+            child = parent.copy()
+            
+            # Mutate child
+            child = self._mutate_genome(child, self.key)
+            
+            new_population.append(child)
+            
+        self.population = new_population
+        self.generation += 1
+        
+    def _mutate_genome(self, genome: Genome, key: jnp.ndarray) -> Genome:
+        """Mutate a genome.
+        
+        Mutation types:
+        1. Add new nodes (30% chance)
+        2. Add new connections (50% chance) 
+        3. Modify weights (80% chance)
+        4. Modify biases (70% chance)
+        5. Enable/disable connections (20% chance)
+        """
+        # Split random key
+        keys = jax.random.split(key, 6)
+        
+        # Add nodes
+        if jax.random.uniform(keys[0]) < self.config['node_add_prob']:
+            # Add 1-3 nodes with decreasing probability
+            n_nodes = 1
+            while jax.random.uniform(keys[1]) < 0.3 and n_nodes < 4:
+                # Pick random enabled connection
+                enabled_conns = [(src, dst) for (src, dst), enabled in genome.connections.items() if enabled]
+                if enabled_conns:
+                    src, dst = enabled_conns[int(jax.random.randint(keys[2], (), 0, len(enabled_conns)))]
+                    genome.add_node_between(src, dst)
+                n_nodes += 1
+        
+        # Add connections
+        if jax.random.uniform(keys[1]) < self.config['conn_add_prob']:
+            # Add multiple connections with decreasing probability
+            n_conns = 0
+            max_attempts = 20  # Prevent infinite loops
+            attempts = 0
+            
+            while attempts < max_attempts and n_conns < 5:
+                # Pick random nodes
+                src = int(jax.random.randint(keys[2], (), 0, genome.n_nodes))
+                dst = int(jax.random.randint(keys[3], (), 0, genome.n_nodes))
+                
+                # Add connection if valid and not already present
+                if src != dst and (src, dst) not in genome.connections:
+                    weight = jax.random.normal(keys[4]) * 0.5
+                    genome.add_connection(src, dst, weight)
+                    n_conns += 1
+                attempts += 1
+        
+        # Mutate weights
+        if jax.random.uniform(keys[2]) < self.config['weight_mutate_prob']:
+            for conn in list(genome.connections.keys()):
+                if genome.connections[conn]:  # Only mutate enabled connections
+                    if jax.random.uniform(keys[3]) < self.config['weight_replace_prob']:
+                        # Reset weight
+                        genome.weights[conn] = jax.random.normal(keys[4]) * self.config['weight_perturb_size']
+                    else:
+                        # Perturb weight
+                        genome.weights[conn] += jax.random.normal(keys[4]) * self.config['weight_perturb_size']
+        
+        # Mutate biases
+        if jax.random.uniform(keys[3]) < self.config['bias_mutate_prob']:
+            for node in list(genome.biases.keys()):
+                if jax.random.uniform(keys[4]) < self.config['bias_replace_prob']:
+                    # Reset bias
+                    genome.biases[node] = jax.random.normal(keys[5]) * self.config['bias_perturb_size']
+                else:
+                    # Perturb bias
+                    genome.biases[node] += jax.random.normal(keys[5]) * self.config['bias_perturb_size']
+        
+        # Enable/disable connections
+        for conn in list(genome.connections.keys()):
+            if jax.random.uniform(keys[5]) < 0.2:  # 20% chance per connection
+                genome.connections[conn] = not genome.connections[conn]
+        
+        return genome
+    
+    def get_average_nodes(self) -> float:
+        """Get average number of nodes in population."""
+        return np.mean([g.n_nodes for g in self.population])
+    
+    def get_average_connections(self) -> float:
+        """Get average number of connections in population."""
+        return np.mean([len(g.connections) for g in self.population])
+        
+    def get_activation_distribution(self) -> Dict[str, float]:
+        """Get distribution of activation functions in population.
+        
+        Returns:
+            Dictionary mapping activation function names to their frequency
+        """
+        # For now we only use ReLU
+        return {'relu': 1.0}
+    
+    def run_evolution(self, evaluator: Callable[[Network], float], max_generations: int,
+                 fitness_threshold: float, reset_mutations: bool = True,
+                 max_stagnation: int = 15, verbose: bool = True) -> Tuple[Network, float]:
+        """Run the evolution process
+        
+        Args:
+            evaluator: Function that takes a network and returns its fitness
+            max_generations: Maximum number of generations to run
+            fitness_threshold: Target fitness to achieve
+            reset_mutations: Whether to reset mutations when fitness improves
+            max_stagnation: Maximum generations without improvement before stopping
+            verbose: Whether to print progress
+            
+        Returns:
+            Tuple of (best network, best fitness)
+        """
+        best_fitness = float('-inf')
+        best_network = None
+        stagnation_counter = 0
+        
+        for generation in range(max_generations):
+            # Evaluate current population
+            fitnesses = []
+            for genome in self.population:
+                network = genome.to_network()
+                fitness = evaluator(network)
+                genome.fitness = fitness
+                fitnesses.append(fitness)
+                
+                # Update best if improved
+                if fitness > best_fitness:
+                    best_fitness = fitness
+                    best_network = network
+                    stagnation_counter = 0
+                    if reset_mutations:
+                        self.reset_innovation()
+            
+            # Get statistics
+            avg_fitness = sum(fitnesses) / len(fitnesses)
+            generation_best = max(fitnesses)
+            
+            # Print progress
+            if verbose:
+                print(f"\nGeneration {generation}:")
+                print(f"  Best Fitness: {best_fitness:.2f}")
+                print(f"  Generation Best: {generation_best:.2f}")
+                print(f"  Average Nodes: {self.get_average_nodes():.1f}")
+                print(f"  Average Connections: {self.get_average_connections():.1f}")
+            
+            # Check for improvement
+            if generation_best <= best_fitness:
+                stagnation_counter += 1
+            else:
+                stagnation_counter = 0
+            
+            # Create next generation
+            self.tell(fitnesses)
+            
+            # Stop if stagnated too long
+            if stagnation_counter >= max_stagnation:
+                if verbose:
+                    print(f"\nStopping: No improvement for {max_stagnation} generations")
+                break
+                
+        if verbose:
+            print("\nTraining complete!")
+            print(f"Best fitness achieved: {best_fitness:.2f}")
+            
+        return best_network, best_fitness