Upload old_train.py with huggingface_hub

old_train.py

@@ -0,0 +1,377 @@
+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()