Error occurred while evaluating genome 2: JAX encountered invalid PRNG key data: expected key_data.ndim >= 1; got ndim=0

[kehsan]

I have this code for creating PRNGkey:

import jax
import neat
import jax.random as random
from evojax.task.slimevolley import SlimeVolley
import time

jax.config.update("jax_enable_x64", True)

def fitness_fn(genomes, config):
    # Generate a new PRNG key for each fitness evaluation
    # This ensures that each genome gets a unique PRNG key
    prng_key = random.PRNGKey(int(time.time()))

    for genome_id, genome in genomes:
        try:
            net = neat.nn.FeedForwardNetwork.create(genome, config)
            env = SlimeVolley(test=False)

            # Split the PRNG key to get a unique key for this genome
            #prng_key, genome_key = random.split(prng_key)
            
            prng_key, genome_key = random.split(prng_key, 2)
            genome_key = genome_key.astype(jnp.uint64)
                
            print('genome_key:', genome_key)
            
            state = env.reset(genome_key)
            print(f"State after reset: {state}")

            total_reward = 0
            while not state.done:
                action = net.activate(state.obs)
                state, reward, done = env.step(action)
                total_reward += reward

            genome.fitness = total_reward
        except Exception as e:
            print(f"Error occurred while evaluating genome {genome_id}: {e}")
            genome.fitness = 0

config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction,
                    neat.DefaultSpeciesSet, neat.DefaultStagnation,
                    r'C:\Users\kehsa\Sakana\project1\evojax\config.txt')


p = neat.Population(config)
p.run(fitness_fn, 50)

I get this output:

genome_key: [1086632469 1657850374]
Error occurred while evaluating genome 1: JAX encountered invalid PRNG key data: expected key_data.ndim >= 1; got ndim=0

It breaks on: state=env.reset(genome_key)

I don't fully understand what the error is and trying different methods to reshape did not work. My Jax version is 0.4.26

[jakevdp]

This line looks problematic:

genome_key = genome_key.astype(jnp.uint64)

A key is a pair of uint32 values, I'm not sure what your intent is in converting it to uint64, but I can't imagine a scenario where it is correct to convert the dtype of your PRNG key.

I can't check whether this is the case because I'm unable to run your code without access to your config file. If you want further help, I'd suggest putting together a minimal reproducible example so that others can reproduce the behavior you're seeing.

[kehsan]

The line ' genome_key = genome_key.astype(jnp.uint64)' was llm suggestion and it makes no difference whether it is there or not, still the same error. My config file is attached. I also tried ' genome_key = jrandom.PRNGKey(genome_id) '
same error, jrandom is jax.random
config.txt

[jakevdp]

Thank you for linking to the config file, but that doesn't make your example minimally reproducible. Can you please edit your script so that it's self-contained, and I can paste it into an interpreter without modification to see the output you're seeing? Note this may take a few minutes of work on your part, but it helps in avoiding any guesswork in trying to help troubleshoot your problem.

[kehsan]

This is the entire code I am running on single cell on a notebook and the config.txt is the only variable from outside the code, I run nothing before running this cell. I cleared the kernel and just ran this cell again

import jax
import neat
import jax.random as random
from evojax.task.slimevolley import SlimeVolley
import time

#jax.config.update("jax_enable_x64", True)

def fitness_fn(genomes, config):
   # Generate a new PRNG key for each fitness evaluation
   # This ensures that each genome gets a unique PRNG key
   prng_key = random.PRNGKey(int(time.time()))

   for genome_id, genome in genomes:
       try:
           net = neat.nn.FeedForwardNetwork.create(genome, config)
           env = SlimeVolley(test=False)

           # Split the PRNG key to get a unique key for this genome
           #prng_key, genome_key = random.split(prng_key)
           
           prng_key, genome_key = random.split(prng_key, 2)
           #genome_key = genome_key.astype(jnp.uint64)
               
           print('genome_key:', genome_key)
           
           state = env.reset(genome_key)
           print(f"State after reset: {state}")

           total_reward = 0
           while not state.done:
               action = net.activate(state.obs)
               state, reward, done = env.step(action)
               total_reward += reward

           genome.fitness = total_reward
       except Exception as e:
           print(f"Error occurred while evaluating genome {genome_id}: {e}")
           genome.fitness = 0

config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction,
                   neat.DefaultSpeciesSet, neat.DefaultStagnation,
                   r'C:\Users\kehsa\Sakana\project1\evojax\config.txt')

p = neat.Population(config)
p.run(fitness_fn, 50) 

and got these errors:
/// genome_key: [1055397681 2317186347]
Error occurred while evaluating genome 1: JAX encountered invalid PRNG key data: expected key_data.ndim >= 1; got ndim=0
genome_key: [1941212796 2626979923]
........ ///

[jakevdp]

Is this any different from the code you pasted above? It still accesses a local config.txt file that is not going to be in the path of another user trying to run it. Can you update the code so that it is self-contained? Note that this may involve figuring out a different strategy to configure your runtime. The neat package docs may be helpful there: I'm unfamiliar with the package so I cannot offer any advice on that account.

The benefit of doing this work is that we may actually be able to answer your question!

[kehsan]

I put the config value in a cell as such :
config_dict = {
"NEAT": {
"fitness_criterion": "max",
"fitness_threshold": 100,
"pop_size": 150,
"reset_on_extinction": False,
},
"DefaultGenome": {
"num_inputs": 12,
.... other values

and tried this format

Merge dictionaries

merged_config = {**config_dict["NEAT"], **config_dict["DefaultGenome"], **config_dict["DefaultReproduction"], **config_dict["DefaultSpeciesSet"], **config_dict["DefaultStagnation"]}

Create NEAT config

neat_config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction, neat.DefaultSpeciesSet, neat.DefaultStagnation, **merged_config)

p = neat.Population(neat_config)
p.run(fitness_fn, 50)

but getting this error:
init() got an unexpected keyword argument 'fitness_criterion'

difficulty feeding config.txt to neat.config() !!

[jakevdp]

OK, I tried this myself but pip install neat failed. It looks like the latest release is from 2010 (https://pypi.org/project/neat/#history), so it's not surprising that it can't be installed on a modern Python stack.

Perhaps you can come up with a reproduction that avoids neat entirely?

[kehsan]

It is an assignment that I have to use neat. I installed neat from neat-python wheel I believe https://pypi.org/project/neat-python/#files under download files. It is uploaded 2017. "pip install neat-python" may also work.

I tried to recreate the config.txt file on a cell and pass it as dictionary to neat_config as such ///
import jax
import neat
import jax.random as random
from evojax.task.slimevolley import SlimeVolley
import time

def fitness_fn(genomes, config):
# Generate a new PRNG key for each fitness evaluation
# This ensures that each genome gets a unique PRNG key
prng_key = random.PRNGKey(int(time.time()))

for genome_id, genome in genomes:
    try:
        net = neat.nn.FeedForwardNetwork.create(genome, config)
        env = SlimeVolley(test=False)

        # Split the PRNG key to get a unique key for this genome
        #prng_key, genome_key = random.split(prng_key)
        
        prng_key, genome_key = random.split(prng_key, 2)
            
        print('genome_key:', genome_key)
        
        state = env.reset(genome_key)
        print(f"State after reset: {state}")

        total_reward = 0
        while not state.done:
            action = net.activate(state.obs)
            state, reward, done = env.step(action)
            total_reward += reward

        genome.fitness = total_reward
    except Exception as e:
        print(f"Error occurred while evaluating genome {genome_id}: {e}")
        genome.fitness = 0

Merge dictionaries

#merged_config = {**config_dict["[NEAT]"], **config_dict["[DefaultGenome]"], **config_dict["[DefaultReproduction]"], **config_dict["[DefaultSpeciesSet]"], **config_dict["[DefaultStagnation]"]}

merged_config = ""
for key, value in config_dict.items():
merged_config += f"[{key.upper()}]\n"
merged_config += "\n".join([f"{k.replace('_', ' ').capitalize()} = {v}" for k, v in value.items()])
merged_config += "\n\n" # Add extra newline between sections

#print(merged_config)

Create NEAT config

neat_config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction, neat.DefaultSpeciesSet, neat.DefaultStagnation, merged_config)

p = neat.Population(neat_config)
p.run(fitness_fn, 50)///

The code dictionary has the exact format of the config.txt but the neat.config is still looking for a local file and does not seem to accept other formats, as I get this error, I am not sure if there is a way around it.

Cell In[3], line 55
49 merged_config += "\n\n" # Add extra newline between sections
51 #print(merged_config)
52
53
54 # Create NEAT config
---> 55 neat_config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction, neat.DefaultSpeciesSet, neat.DefaultStagnation, merged_config)
57 p = neat.Population(neat_config)
58 p.run(fitness_fn, 50)

File ~\miniconda3\envs\sakana\lib\site-packages\neat\config.py:153, in Config.init(self, genome_type, reproduction_type, species_set_type, stagnation_type, filename)
150 self.stagnation_type = stagnation_type
152 if not os.path.isfile(filename):
--> 153 raise Exception('No such config file: ' + os.path.abspath(filename))
155 parameters = ConfigParser()
156 with open(filename) as f:

Exception: No such config file: C:\Users\kehsa\sakana\project1\evojax[NEAT]

[jakevdp]

Please use triple back ticks for code blocks. I’ve already edited your previous posts - look at those to see how to do it. You can also see the github help pages, e.g. https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks

[jakevdp]

pip install neat-python may work

Thank you – this is important information and it would have been helpful if you'd included it in the original question! Otherwise I see import neat and have to do my best to guess what code you're importing. The key to getting answers to code questions on the internet is to make it as easy as possible for the answerer to understand exactly what code you're running.

[kehsan]

I printed my genome_key after split and seems to be the right type
Genome Key after Split: [1528553899 553433336]
Type of genome_key: <class 'jaxlib.xla_extension.ArrayImpl'>
I wonder if my env.reset(genome_key) that i use from :
from evojax.task.slimevolley import SlimeVolley
i attach it here
slimevolley.txt
It is suppose ot be slimevolley.py but i could not upload that so i saved and uploaded as slimevolley.txt

[kehsan]

@jakevdp, I guess you have no extra ideas of how I can figure this syntax error ///Error occurred while evaluating genome 1: JAX encountered invalid PRNG key data: expected key_data.ndim >= 1; got ndim=0 ///. Thanks anyways.

[jakevdp]

I'm still waiting for a minimal reproduction – if you provide one, I may be able to help. Without that, all I can say is that somewhere in your code one of the keys has the wrong shape and you should make sure it has the correct shape instead.

[kehsan]

@jakevdp I am not sure exactly how the minimal reproduction from your point of view is. If you have installed neat-python and other related import libraries from the code given and have config.txt file also downloaded you should be able to reproduce the error. If I am wrong please let me know at which point. The github pages I am following to reproduce slimevolley using neat algorithm and JAX are:

https://github.com/google/evojax
git clone https://github.com/hardmaru/slimevolleygym.git

[jakevdp]

Sorry - I've lost track of this. Each of your responses has been a new thread, rather than a reply to the existing thread, so now it's rather difficult to follow the conversation.

Could you edit your original question with the minimal reproducible example of the code you're running, along with all instructions needed to run it (please use triple back-ticks for code formatting), and then reply in this thread to let me know when you've done so? Thanks!