old_train.py

import jax
import jax.numpy as jnp
from jax import random
from evojax.task.slimevolley import SlimeVolley
from typing import List, Tuple, Dict
import numpy as np
import time

class NodeGene:
    def __init__(self, id: int, node_type: str, activation: str = 'tanh'):
        self.id = id
        self.type = node_type  # 'input', 'hidden', or 'output'
        self.activation = activation
        
        # Use both id and timestamp for randomization
        timestamp = int(time.time() * 1000)
        key = random.PRNGKey(hash((id, timestamp)) % (2**32))
        self.bias = float(random.normal(key, shape=()) * 0.1)  # Small random bias

class ConnectionGene:
    def __init__(self, source: int, target: int, weight: float = None, enabled: bool = True):
        self.source = source
        self.target = target
        
        # Use source, target, and timestamp for randomization
        timestamp = int(time.time() * 1000)
        key = random.PRNGKey(hash((source, target, timestamp)) % (2**32))
        
        if weight is None:
            key, subkey = random.split(key)
            weight = float(random.normal(subkey, shape=()) * 0.1)  # Small random weight
        self.weight = weight
        self.enabled = enabled
        self.innovation = hash((source, target))

class Genome:
    def __init__(self, n_inputs: int, n_outputs: int):
        # Create input nodes (0 to n_inputs-1)
        self.node_genes = {i: NodeGene(i, 'input') for i in range(n_inputs)}
        
        # Create exactly 3 output nodes for left, right, jump
        n_outputs = 3  # Force exactly 3 outputs
        for i in range(n_outputs):
            self.node_genes[n_inputs + i] = NodeGene(n_inputs + i, 'output')
        
        self.connection_genes: List[ConnectionGene] = []
        
        # Initialize with randomized connections using unique keys
        timestamp = int(time.time() * 1000)
        master_key = random.PRNGKey(hash((n_inputs, n_outputs, timestamp)) % (2**32))
        
        # Add direct connections with random weights
        for i in range(n_inputs):
            for j in range(n_outputs):
                master_key, key = random.split(master_key)
                if random.uniform(key, shape=()) < 0.7:  # 70% chance of connection
                    master_key, key = random.split(master_key)
                    weight = float(random.normal(key, shape=()) * 0.5)  # Larger initial weights
                    self.connection_genes.append(
                        ConnectionGene(i, n_inputs + j, weight=weight)
                    )
        
        # Add hidden nodes with random connections
        master_key, key = random.split(master_key)
        n_hidden = int(random.randint(key, (), 1, 4))  # Random number of hidden nodes
        hidden_start = n_inputs + n_outputs
        
        for i in range(n_hidden):
            node_id = hidden_start + i
            self.node_genes[node_id] = NodeGene(node_id, 'hidden')
            
            # Connect random inputs to this hidden node
            for j in range(n_inputs):
                master_key, key = random.split(master_key)
                if random.uniform(key, shape=()) < 0.5:
                    master_key, key = random.split(master_key)
                    weight = float(random.normal(key, shape=()) * 0.5)
                    self.connection_genes.append(
                        ConnectionGene(j, node_id, weight=weight)
                    )
            
            # Connect this hidden node to random outputs
            for j in range(n_outputs):
                master_key, key = random.split(master_key)
                if random.uniform(key, shape=()) < 0.5:
                    master_key, key = random.split(master_key)
                    weight = float(random.normal(key, shape=()) * 0.5)
                    self.connection_genes.append(
                        ConnectionGene(node_id, n_inputs + j, weight=weight)
                    )

    def mutate(self, config: Dict):
        key = random.PRNGKey(0)
        
        # Mutate connection weights
        for conn in self.connection_genes:
            key, subkey = random.split(key)
            if random.uniform(subkey, shape=()) < config['weight_mutation_rate']:
                key, subkey = random.split(key)
                # Sometimes reset weight completely
                if random.uniform(subkey, shape=()) < 0.1:
                    key, subkey = random.split(key)
                    conn.weight = float(random.normal(subkey, shape=()) * 0.5)
                else:
                    # Otherwise adjust existing weight
                    key, subkey = random.split(key)
                    conn.weight += float(random.normal(subkey) * config['weight_mutation_power'])
        
        # Mutate node biases
        for node in self.node_genes.values():
            key, subkey = random.split(key)
            if random.uniform(subkey, shape=()) < 0.1:  # 10% chance to mutate bias
                key, subkey = random.split(key)
                node.bias += float(random.normal(subkey) * 0.1)
        
        # Add new node
        key, subkey = random.split(key)
        if random.uniform(subkey, shape=()) < config['add_node_rate']:
            if self.connection_genes:
                # Choose random connection to split
                conn = np.random.choice(self.connection_genes)
                new_id = max(self.node_genes.keys()) + 1
                
                # Create new node with random bias
                self.node_genes[new_id] = NodeGene(new_id, 'hidden')
                
                # Create two new connections with some randomization
                key, subkey = random.split(key)
                weight1 = float(random.normal(subkey, shape=()) * 0.5)
                key, subkey = random.split(key)
                weight2 = float(random.normal(subkey, shape=()) * 0.5)
                
                self.connection_genes.append(
                    ConnectionGene(conn.source, new_id, weight=weight1)
                )
                self.connection_genes.append(
                    ConnectionGene(new_id, conn.target, weight=weight2)
                )
                
                # Disable old connection
                conn.enabled = False

        # Add new connection
        key, subkey = random.split(key)
        if random.uniform(subkey, shape=()) < config['add_connection_rate']:
            # Get all possible nodes
            nodes = list(self.node_genes.keys())
            for _ in range(10):  # Try 10 times to find valid connection
                source = np.random.choice(nodes)
                target = np.random.choice(nodes)
                
                # Ensure forward propagation (source id < target id)
                if source < target:
                    # Check if connection already exists
                    if not any(c.source == source and c.target == target 
                             for c in self.connection_genes):
                        key, subkey = random.split(key)
                        weight = float(random.normal(subkey, shape=()) * 0.5)
                        self.connection_genes.append(
                            ConnectionGene(source, target, weight=weight)
                        )
                        break

class Network:
    def __init__(self, genome: Genome):
        self.genome = genome
        # Sort nodes by ID to ensure consistent ordering
        self.input_nodes = sorted([n for n in genome.node_genes.values() if n.type == 'input'], key=lambda x: x.id)
        self.hidden_nodes = sorted([n for n in genome.node_genes.values() if n.type == 'hidden'], key=lambda x: x.id)
        self.output_nodes = sorted([n for n in genome.node_genes.values() if n.type == 'output'], key=lambda x: x.id)
        
        # Verify we have exactly 3 output nodes
        assert len(self.output_nodes) == 3, f"Expected 3 output nodes, got {len(self.output_nodes)}"
        
    def forward(self, x: jnp.ndarray) -> jnp.ndarray:
        # Ensure input is 2D with shape (batch_size, input_dim)
        if len(x.shape) == 1:
            x = jnp.expand_dims(x, 0)
        
        batch_size = x.shape[0]
        
        # Initialize node values
        values = {}
        for node in self.genome.node_genes.values():
            values[node.id] = jnp.zeros((batch_size,))
            values[node.id] = values[node.id] + node.bias
        
        # Set input values
        for i, node in enumerate(self.input_nodes):
            values[node.id] = x[:, i]
        
        # Process nodes in order
        for node in self.hidden_nodes + self.output_nodes:
            # Sum incoming connections
            total = jnp.zeros((batch_size,))
            total = total + node.bias
            
            for conn in self.genome.connection_genes:
                if conn.enabled and conn.target == node.id:
                    total = total + values[conn.source] * conn.weight
            
            # Apply activation
            values[node.id] = jnp.tanh(total)
        
        # Get output values and ensure shape (batch_size, 3)
        outputs = []
        for node in self.output_nodes:
            outputs.append(values[node.id])
        
        # Stack along new axis to get (batch_size, 3)
        return jnp.stack(outputs, axis=-1)

def evaluate_network(network: Network, env: SlimeVolley, n_episodes: int = 10) -> float:
    total_reward = 0.0
    
    # Generate a unique key for this evaluation
    timestamp = int(time.time() * 1000)
    network_id = id(network)
    master_key = random.PRNGKey(hash((network_id, timestamp)) % (2**32))
    
    for episode in range(n_episodes):
        # Reset environment with proper key shape
        master_key, reset_key = random.split(master_key)
        state = env.reset(reset_key[None, :])  # Add batch dimension
        done = False
        episode_reward = 0.0
        steps = 0
        
        while not done and steps < 1000:  # Add step limit
            # Get observation and normalize
            obs = state.obs[None, :] / 10.0  # Add batch dimension and scale inputs
            
            # Get action from network (shape: batch_size, 3)
            raw_action = network.forward(obs)
            
            # Convert to binary actions using thresholds
            thresholds = jnp.array([0.3, 0.3, 0.4])  # left, right, jump
            binary_action = (raw_action > thresholds).astype(jnp.float32)
            
            # Prevent simultaneous left/right using logical operations
            both_active = jnp.logical_and(binary_action[:, 0] > 0, binary_action[:, 1] > 0)
            prefer_left = raw_action[:, 0] > raw_action[:, 1]
            
            # Update binary action based on preference
            binary_action = binary_action.at[:, 0].set(
                jnp.where(both_active, prefer_left.astype(jnp.float32), binary_action[:, 0])
            )
            binary_action = binary_action.at[:, 1].set(
                jnp.where(both_active, (~prefer_left).astype(jnp.float32), binary_action[:, 1])
            )
            
            # Step environment
            master_key, step_key = random.split(master_key)
            next_state, reward, done = env.step(state, binary_action)  # Already batched
            
            # Process reward and done flag
            if isinstance(reward, jnp.ndarray):
                reward = float(jnp.reshape(reward, (-1,))[0])  # Get first element if batched
            if isinstance(done, jnp.ndarray):
                done = bool(jnp.reshape(done, (-1,))[0])  # Convert to Python bool
            
            # Add small reward for movement to encourage exploration
            any_movement = jnp.any(binary_action[:, :2] > 0)
            movement_reward = 0.1 if bool(any_movement) else 0.0
            
            # Add small reward for keeping ball in play
            ball_height = float(jnp.reshape(next_state.obs[1], (-1,))[0]) if hasattr(next_state.obs, '__getitem__') else 0.0
            height_reward = 0.1 if ball_height > 0.5 else 0.0
            
            # Add reward for ball position and velocity
            ball_x = float(jnp.reshape(next_state.obs[4], (-1,))[0])  # Ball x position
            ball_vx = float(jnp.reshape(next_state.obs[6], (-1,))[0])  # Ball x velocity
            position_reward = 0.2 if ball_x > 0 else 0.0  # Reward for keeping ball on opponent's side
            velocity_reward = 0.1 if ball_vx > 0 else 0.0  # Reward for hitting ball towards opponent
            
            # Calculate step reward with more emphasis on game outcome
            step_reward = reward * 2.0  # Double the importance of winning/losing
            bonus_reward = movement_reward + height_reward + position_reward + velocity_reward
            total_step_reward = step_reward + bonus_reward * 0.5  # Scale down bonus rewards
            
            episode_reward += total_step_reward
            state = next_state
            steps += 1
            
            # Early termination bonus
            if done and reward > 0:  # Won the point
                episode_reward += 10.0
            
        total_reward += episode_reward
        
    return total_reward / n_episodes

def main():
    # Initialize environment
    env = SlimeVolley()
    
    # NEAT configuration
    config = {
        'population_size': 50,  # Smaller population for faster iteration
        'weight_mutation_rate': 0.8,
        'weight_mutation_power': 0.3,  # Increased for more exploration
        'add_node_rate': 0.3,
        'add_connection_rate': 0.5,
    }
    
    # Create initial population
    population = [
        Network(Genome(n_inputs=12, n_outputs=3))
        for _ in range(config['population_size'])
    ]
    
    best_fitness = float('-inf')
    generations_without_improvement = 0
    
    # Evolution loop
    for generation in range(500):  # More generations
        print(f"\nGeneration {generation}")
        print("-" * 20)
        
        # Evaluate population
        fitnesses = []
        for i, net in enumerate(population):
            fitness = evaluate_network(net, env)
            fitnesses.append(fitness)
            print(f"Network {i}: Fitness = {fitness:.2f}")
            
            if fitness > best_fitness:
                best_fitness = fitness
                generations_without_improvement = 0
                print(f"New best fitness: {best_fitness:.2f}")
            
        # Check for improvement
        generations_without_improvement += 1
        if generations_without_improvement > 20:
            print("No improvement for 20 generations, increasing mutation rates")
            config['weight_mutation_rate'] = min(1.0, config['weight_mutation_rate'] * 1.2)
            config['weight_mutation_power'] = min(0.5, config['weight_mutation_power'] * 1.2)
            generations_without_improvement = 0
        
        # Print progress
        avg_fitness = sum(fitnesses) / len(fitnesses)
        print(f"\nBest fitness: {best_fitness:.2f}")
        print(f"Average fitness: {avg_fitness:.2f}")
        
        # Selection and reproduction
        new_population = []
        sorted_indices = np.argsort(fitnesses)[::-1]  # Best to worst
        
        # Keep best networks
        n_elite = 5  # Fewer elites
        new_population.extend([population[i] for i in sorted_indices[:n_elite]])
        print(f"Keeping top {n_elite} networks")
        
        # Create offspring from best networks
        while len(new_population) < config['population_size']:
            # Tournament selection
            tournament_size = 5
            tournament = np.random.choice(sorted_indices[:20], tournament_size, replace=False)
            parent_idx = tournament[np.argmax([fitnesses[i] for i in tournament])]
            parent = population[parent_idx]
            
            # Create offspring
            child_genome = Genome(12, 3)
            child_genome.node_genes = parent.genome.node_genes.copy()
            child_genome.connection_genes = parent.genome.connection_genes.copy()
            
            # Mutate child
            child_genome.mutate(config)
            
            # Add to new population
            new_population.append(Network(child_genome))
            
        population = new_population
        print(f"Created {len(population)} networks for next generation")

if __name__ == '__main__':
    main()