Can the network evolved using NEAT be constructed in PyTorch and further trained using backpropagation

[Polarisamoon]

I'm trying to leverage PyTorch to further train the NEAT-generated network via backpropagation for improved performance. Currently, I'm manually building the network in PyTorch based on the node and edge details output by print(pipeline.algorithm.genome.repr(state, *best)). However, this process is cumbersome and becomes impractical when the network has a large number of nodes and edges. Is there a more efficient approach?

[jravishankar]

I have an attempt at an implementation in JAX (though it is likely not perfect).
It is a simple backpropagation routine that updates the weights of each generation (based on this post).

The main backpropagation functionality is contained in backprop_individual,
which ultimately is run using vmap to backpropagate the weights of all genomes in the population.

I was doing this for a classification task, so the loss function evaluations should be changed according to your requirements.
I had some trouble integrating it directly with the problem.evaluate function (which would be a much cleaner solution).

It worked well for simple tasks like XOR and Two Circles, but I had trouble getting it to work for Spiral data.
If anyone has any suggestions for encouraging more complex architectures and avoiding local minima/single-species domination, I would appreciate the advice (not just hyperparameter tuning of course).

`
import json
import os
import warnings

import jax, jax.numpy as jnp
import optax
import datetime, time
import numpy as np

from tensorneat.algorithm import BaseAlgorithm
from tensorneat.problem import BaseProblem
from tensorneat.common import State, StatefulBaseClass

class BackpropPipeline(StatefulBaseClass):
def init(
self,
algorithm: BaseAlgorithm,
problem: BaseProblem,
seed: int = 42,
fitness_target: float = 1,
generation_limit: int = 1000,
is_save: bool = False,
save_dir=None,
show_problem_details: bool = False,
using_multidevice: bool = False,
learning_rate: float = 0.01,
backprop_max_iters: int = 600,
minibatch_size: int = 20
):
assert problem.jitable, "Currently, problem must be jitable"

	self.algorithm = algorithm
	self.problem = problem
	self.seed = seed
	self.fitness_target = fitness_target
	self.generation_limit = generation_limit
	self.pop_size = self.algorithm.pop_size
	self.learning_rate = learning_rate
	self.backprop_max_iters = backprop_max_iters
	self.minibatch_size = minibatch_size

	np.random.seed(self.seed)

	assert (
		algorithm.num_inputs == self.problem.input_shape[-1]
	), f"algorithm input shape is {algorithm.num_inputs} but problem input shape is {self.problem.input_shape}"

	self.best_genome = None
	self.best_fitness = float("-inf")
	self.generation_timestamp = None
	self.is_save = is_save

	if is_save:
		if save_dir is None:
			now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
			self.save_dir = f"./{self.__class__.__name__} {now}"
		else:
			self.save_dir = save_dir
		print(f"save to {self.save_dir}")
		if not os.path.exists(self.save_dir):
			os.makedirs(self.save_dir)
		self.genome_dir = os.path.join(self.save_dir, "genomes")
		if not os.path.exists(self.genome_dir):
			os.makedirs(self.genome_dir)

	self.show_problem_details = show_problem_details

	self.using_multidevice = using_multidevice
	if self.using_multidevice:
		assert jax.device_count() > 1, f"using_multidevice requires more than 1 device, but {jax.device_count()=} devices are available"
		print(f"Using {jax.device_count()} devices!")

def setup(self, state=State()):

	state = state.register(randkey=jax.random.PRNGKey(self.seed))
	state = self.algorithm.setup(state)
	state = self.problem.setup(state)

	if self.is_save:
		# self.save(state=state, path=os.path.join(self.save_dir, "pipeline.pkl"))
		with open(os.path.join(self.save_dir, "config.txt"), "w") as f:
			f.write(json.dumps(self.show_config(), indent=4))
		# create log file
		with open(os.path.join(self.save_dir, "log.txt"), "w") as f:
			f.write("Generation,Max,Min,Mean,Std,Cost Time\n")

	print("initializing finished")
	return state



def backprop_individual(self, state, key, transformed):
	cal_seqs, nodes, conns, _ = transformed
	const_cal_seqs = jax.lax.stop_gradient(cal_seqs)
	const_nodes     = jax.lax.stop_gradient(nodes)

	lr             = self.learning_rate
	n_steps        = self.backprop_max_iters
	minibatch_size = self.minibatch_size
	check_interval = 20

	# RMSProp + clipping
	max_grad_norm = 5.0
	optimizer = optax.chain(
		optax.clip_by_global_norm(max_grad_norm),
		optax.rmsprop(lr, decay=0.001),
	)
	opt_state = optimizer.init(conns)

	# full‑batch loss
	def full_loss_fn(c):
		u_conns = self.algorithm.genome.unflatten_connections(const_nodes, c)
		t       = (const_cal_seqs, const_nodes, c, u_conns)
		preds   = jax.vmap(lambda x: self.algorithm.forward(state, t, x))(self.problem.inputs)
		y       = self.problem.targets
		return -jnp.mean(y * jnp.log(preds + 1e-8)
					   + (1 - y) * jnp.log(1 - preds + 1e-8))

	# save initial performance for early stopping
	init_full_loss    = full_loss_fn(conns)
	best_full_loss    = init_full_loss
	best_conns        = conns
	last_good_conns   = conns
	last_good_opt     = opt_state


	init_carry = (0, True, conns, opt_state, key,
				  best_full_loss, best_conns,
				  last_good_conns, last_good_opt)

	def cond_fun(carry):
		step, active, *_ = carry
		return jnp.logical_and(step < n_steps, active)

	def body_fun(carry):
		(step, active, cur_conns, cur_opt, prng_key,
		 best_loss, best_c, backup_c, backup_opt) = carry

		# minibatch sampling
		prng_key, sk = jax.random.split(prng_key)
		idxs = jax.random.choice(sk, a=self.problem.inputs.shape[0],
								 shape=(minibatch_size,), replace=False)
		xb   = self.problem.inputs[idxs]
		yb   = self.problem.targets[idxs]

		# minibatch loss  + grads
		def mb_loss(c):
			u = self.algorithm.genome.unflatten_connections(const_nodes, c)
			t = (const_cal_seqs, const_nodes, c, u)
			p = jax.vmap(lambda x: self.algorithm.forward(state, t, x))(xb)
			return -jnp.mean(yb * jnp.log(p + 1e-8)
						   + (1 - yb) * jnp.log(1 - p + 1e-8))

		loss, grads = jax.value_and_grad(mb_loss)(cur_conns)

		# optimizer + weight updates
		updates, new_opt = optimizer.update(grads, cur_opt, params=cur_conns)
		new_conns = optax.apply_updates(cur_conns, updates)
		next_step = step + 1

		# check every check_interval steps
		is_check = (next_step % check_interval) == 0
		new_full_loss = jnp.where(is_check, full_loss_fn(new_conns), best_loss)			

		improved = jnp.logical_and(is_check, new_full_loss < best_loss)


		best_loss1 = jnp.where(improved, new_full_loss, best_loss)
		best_c1    = jnp.where(improved, new_conns, best_c)

		backup_c1   = jnp.where(improved, new_conns, backup_c)
		backup_opt1 = jax.lax.cond(improved, lambda _: new_opt,     lambda _: backup_opt, operand=None)


		should_bail = jnp.logical_and(is_check, new_full_loss >= best_loss)


		final_conns = jnp.where(should_bail, backup_c1, new_conns)
		final_opt   = jax.lax.cond(should_bail, lambda _: backup_opt1,lambda _: new_opt,   operand=None)

		new_active = jnp.where(should_bail, False, True)

		return (
			next_step,
			new_active,
			final_conns,
			final_opt,
			prng_key,
			best_loss1,
			best_c1,
			backup_c1,
			backup_opt1
		)

	final_carry = jax.lax.while_loop(cond_fun, body_fun, init_carry)
	(
		_,
		_,
		final_conns,
		_,
		_,
		final_best_loss,
		final_best_conns,
		_,
		_
	) = final_carry

	# rebuild u_conns from the best found weights
            # i don't think this is actually necessary since connections won't disappear from backprop
	final_u_conns = self.algorithm.genome.unflatten_connections(
		const_nodes, final_best_conns
	)
	return const_cal_seqs, const_nodes, final_best_conns, final_u_conns


def step(self, state):
	"""
	returns: 
		state, previous_pop, fitnesses
	state: updated state
	previous_pop: previous population
	fitnesses: fitnesses of previous population
	"""
	randkey_, randkey = jax.random.split(state.randkey)

	pop = self.algorithm.ask(state)


	pop_transformed = jax.vmap(self.algorithm.transform, in_axes=(None, 0))(
		state, pop
	)

	# keys for backprop
	bp_keys = jax.random.split(randkey_, self.pop_size)

	# perform backpropagation on each transformed genome.
	pop_optimized = jax.vmap(self.backprop_individual, in_axes=(None, 0, 0))(
	  state, bp_keys, pop_transformed
	)

	new_pop_nodes, new_pop_conns = pop_optimized[1], pop_optimized[2]

	state = state.update(
	  pop_nodes=new_pop_nodes,
	  pop_conns=new_pop_conns
	 )

	if not self.using_multidevice:
		keys = jax.random.split(randkey_, self.pop_size)
		fitnesses = jax.vmap(self.problem.evaluate, in_axes=(None, 0, None, 0))(
			state, keys, self.algorithm.forward, pop_optimized
		)
	else: # using_multidevice
		num_devices = jax.device_count()
		assert self.pop_size % num_devices == 0, "if you want to use multiple gpus, pop_size must be divisible by jax.device_count()"
		pop_size_per_device = self.pop_size // num_devices

		keys = jax.random.split(randkey_, (num_devices, pop_size_per_device))
		split_pop_transformed = jax.tree_map(
			lambda x: x.reshape(num_devices, pop_size_per_device, *x.shape[1:]),
			pop_transformed
		)

		fitnesses = jax.pmap(
			lambda key_slice, pop_slice: jax.vmap(self.problem.evaluate, in_axes=(None, 0, None, 0))(
				state, key_slice, self.algorithm.forward, pop_slice
			),
			axis_name='devices',
			in_axes=(0, 0)
		)(keys, split_pop_transformed)

		fitnesses = fitnesses.reshape(self.pop_size)

	# replace nan with -inf
	fitnesses = jnp.where(jnp.isnan(fitnesses), -jnp.inf, fitnesses)

	previous_pop = self.algorithm.ask(state)
	state = self.algorithm.tell(state, fitnesses)

	return state.update(randkey=randkey), previous_pop, fitnesses

def auto_run(self, state):
	print("start compile")
	tic = time.time()
	with warnings.catch_warnings():
		warnings.filterwarnings("ignore",
			message=r"The jitted function .* includes a pmap. Using jit-of-pmap can lead to inefficient data movement"
		)
		compiled_step = jax.jit(self.step).lower(state).compile()

	if self.show_problem_details:
		self.compiled_pop_transform_func = (
			jax.jit(jax.vmap(self.algorithm.transform, in_axes=(None, 0)))
			.lower(state, self.algorithm.ask(state))
			.compile()
		)

	print(
		f"compile finished, cost time: {time.time() - tic:.6f}s",
	)

	for _ in range(self.generation_limit):

		self.generation_timestamp = time.time()

		state, previous_pop, fitnesses = compiled_step(state)

		fitnesses = jax.device_get(fitnesses)

		self.analysis(state, previous_pop, fitnesses)

		if max(fitnesses) >= self.fitness_target:
			print("Fitness limit reached!")
			break

	if int(state.generation) >= self.generation_limit:
		print("Generation limit reached!")

	if self.is_save:
		best_genome = jax.device_get(self.best_genome)
		with open(os.path.join(self.genome_dir, f"best_genome.npz"), "wb") as f:
			np.savez(
				f,
				nodes=best_genome[0],
				conns=best_genome[1],
				fitness=self.best_fitness,
			)

	return state, self.best_genome

def analysis(self, state, pop, fitnesses):

	generation = int(state.generation)

	valid_fitnesses = fitnesses[~np.isinf(fitnesses)]
	# avoid there is no valid fitness in the whole population
	if len(valid_fitnesses) == 0:
		max_f, min_f, mean_f, std_f = [float("NaN")] * 4
	else:
		max_f, min_f, mean_f, std_f = (
			max(valid_fitnesses),
			min(valid_fitnesses),
			np.mean(valid_fitnesses),
			np.std(valid_fitnesses),
		)

	new_timestamp = time.time()

	cost_time = new_timestamp - self.generation_timestamp

	max_idx = np.argmax(fitnesses)
	if fitnesses[max_idx] > self.best_fitness:
		self.best_fitness = fitnesses[max_idx]
		self.best_genome = pop[0][max_idx], pop[1][max_idx]

	if self.is_save:
		# save best
		best_genome = jax.device_get((pop[0][max_idx], pop[1][max_idx]))
		file_name = os.path.join(self.genome_dir, f"{generation}.npz")
		with open(file_name, "wb") as f:
			np.savez(
				f,
				nodes=best_genome[0],
				conns=best_genome[1],
				fitness=self.best_fitness,
			)

		# append log
		with open(os.path.join(self.save_dir, "log.txt"), "a") as f:
			f.write(f"{generation},{max_f},{min_f},{mean_f},{std_f},{cost_time}\n")

	print(
		f"Generation: {generation}, Cost time: {cost_time * 1000:.2f}ms\n",
		f"\tfitness: valid cnt: {len(valid_fitnesses)}, max: {max_f:.4f}, min: {min_f:.4f}, mean: {mean_f:.4f}, std: {std_f:.4f}\n",
	)

	self.algorithm.show_details(state, fitnesses)

	if self.show_problem_details:
		pop_transformed = self.compiled_pop_transform_func(
			state, pop #using previous pop instead of requesting the new one from state here
		)
		self.problem.show_details(
			state, state.randkey, self.algorithm.forward, pop_transformed
		)
	# show details for problem

def show(self, state, best, *args, **kwargs):
	transformed = self.algorithm.transform(state, best)
	return self.problem.show(
		state, state.randkey, self.algorithm.forward, transformed, *args, **kwargs
	)

`

[WLS2002]

Thank you @Polarisamoon for raising this issue, and many thanks to @jravishankar for sharing your implementation!

Previously, I was working on converting NEAT-generated networks into PyTorch for further training via backpropagation. However, due to the complexity involved in dynamically building large networks, and being a bit tied up with other responsibilities recently, I had to temporarily put this plan on hold.

I'm really glad to see your JAX-based solution, @jravishankar. Leveraging JAX's native support for automatic differentiation and vmap for efficient population-wide optimization is a very clean and scalable approach. One thing I noticed is that your gradient descent logic is implemented at the pipeline level. In my initial plan, I was hoping to integrate this kind of gradient-based optimization more locally, at the genome level, so that individual genomes can optionally support backpropagation as a modular feature.

Looking ahead, I’d like to explore how we can add native support for JAX-based gradient descent directly into the Genome class. If anyone is interested in collaborating or contributing, feel free to open a PR. I'd be very happy to see this functionality evolve with community input.

Thanks again to everyone involved!