MOME Container Implementation does not support non-array Genotypes

[Lookatator]

Hi,

It seems the current implementation of the MOME Container does not work if the Genotype is a pytree more general than an array. For example the following line fails to run for Neural Network Genotypes:

QDax/qdax/core/containers/mome_repertoire.py

Line 314 in 45c4c2a

pareto_front_genotypes=cell_genotype.squeeze(axis=0),

We have a working fix in the lab; which we will clean before submitting it in a PR :)

[limbryan]

I think this has been fixed and addressed now with #143 right @felixchalumeau? Where we change the explicit type in the repertoire to account for non-array genotypes and arbitrary pytrees (

QDax/qdax/core/containers/mapelites_repertoire.py

Line 390 in 773e646

lambda x: jnp.zeros(shape=(num_centroids,) + x.shape, dtype=x.dtype),

)

[Lookatator]

Mmh from what I see, the issue still seems to be on the main and develop branches.

What I meant is cell_genotype.squeeze(axis=0) fails to run if cell_genotype is not an array.

Supposedly, the line should be:

jax.tree_util.tree_map(lambda x: x.squeeze(axis=0), cell_genotype)

And the line below: new_batch_of_genotypes=jnp.expand_dims(genotype, axis=0) should also be updated I guess.

@hannah-jan what do you think?