update HyperNEAT;
All example can currently run!
This commit is contained in:
@@ -40,19 +40,22 @@ class HyperNEAT(BaseAlgorithm):
|
||||
output_transform=output_transform,
|
||||
)
|
||||
|
||||
def setup(self, randkey):
|
||||
return State(neat_state=self.neat.setup(randkey))
|
||||
def setup(self, state=State()):
|
||||
state = self.neat.setup(state)
|
||||
state = self.substrate.setup(state)
|
||||
return self.hyper_genome.setup(state)
|
||||
|
||||
def ask(self, state: State):
|
||||
return self.neat.ask(state.neat_state)
|
||||
return self.neat.ask(state)
|
||||
|
||||
def tell(self, state: State, fitness):
|
||||
return state.update(neat_state=self.neat.tell(state.neat_state, fitness))
|
||||
state = self.neat.tell(state, fitness)
|
||||
return state
|
||||
|
||||
def transform(self, individual):
|
||||
transformed = self.neat.transform(individual)
|
||||
query_res = jax.vmap(self.neat.forward, in_axes=(0, None))(
|
||||
self.substrate.query_coors, transformed
|
||||
def transform(self, state, individual):
|
||||
transformed = self.neat.transform(state, individual)
|
||||
query_res = jax.vmap(self.neat.forward, in_axes=(None, 0, None))(
|
||||
state, self.substrate.query_coors, transformed
|
||||
)
|
||||
|
||||
# mute the connection with weight below threshold
|
||||
@@ -74,12 +77,12 @@ class HyperNEAT(BaseAlgorithm):
|
||||
h_nodes, h_conns = self.substrate.make_nodes(
|
||||
query_res
|
||||
), self.substrate.make_conn(query_res)
|
||||
return self.hyper_genome.transform(h_nodes, h_conns)
|
||||
return self.hyper_genome.transform(state, h_nodes, h_conns)
|
||||
|
||||
def forward(self, inputs, transformed):
|
||||
def forward(self, state, inputs, transformed):
|
||||
# add bias
|
||||
inputs_with_bias = jnp.concatenate([inputs, jnp.array([1])])
|
||||
return self.hyper_genome.forward(inputs_with_bias, transformed)
|
||||
return self.hyper_genome.forward(state, inputs_with_bias, transformed)
|
||||
|
||||
@property
|
||||
def num_inputs(self):
|
||||
@@ -94,10 +97,10 @@ class HyperNEAT(BaseAlgorithm):
|
||||
return self.neat.pop_size
|
||||
|
||||
def member_count(self, state: State):
|
||||
return self.neat.member_count(state.neat_state)
|
||||
return self.neat.member_count(state)
|
||||
|
||||
def generation(self, state: State):
|
||||
return self.neat.generation(state.neat_state)
|
||||
return self.neat.generation(state)
|
||||
|
||||
|
||||
class HyperNodeGene(BaseNodeGene):
|
||||
@@ -110,7 +113,7 @@ class HyperNodeGene(BaseNodeGene):
|
||||
self.activation = activation
|
||||
self.aggregation = aggregation
|
||||
|
||||
def forward(self, attrs, inputs, is_output_node=False):
|
||||
def forward(self, state, attrs, inputs, is_output_node=False):
|
||||
return jax.lax.cond(
|
||||
is_output_node,
|
||||
lambda: self.aggregation(inputs), # output node does not need activation
|
||||
@@ -121,6 +124,6 @@ class HyperNodeGene(BaseNodeGene):
|
||||
class HyperNEATConnGene(BaseConnGene):
|
||||
custom_attrs = ["weight"]
|
||||
|
||||
def forward(self, attrs, inputs):
|
||||
def forward(self, state, attrs, inputs):
|
||||
weight = attrs[0]
|
||||
return inputs * weight
|
||||
|
||||
@@ -1,4 +1,10 @@
|
||||
from utils import State
|
||||
|
||||
|
||||
class BaseSubstrate:
|
||||
def setup(self, state=State()):
|
||||
return state
|
||||
|
||||
def make_nodes(self, query_res):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@@ -11,8 +11,8 @@ if __name__ == "__main__":
|
||||
genome=DefaultGenome(
|
||||
num_inputs=27,
|
||||
num_outputs=8,
|
||||
max_nodes=50,
|
||||
max_conns=100,
|
||||
max_nodes=100,
|
||||
max_conns=200,
|
||||
node_gene=DefaultNodeGene(
|
||||
activation_options=(Act.tanh,),
|
||||
activation_default=Act.tanh,
|
||||
@@ -21,6 +21,8 @@ if __name__ == "__main__":
|
||||
),
|
||||
pop_size=1000,
|
||||
species_size=10,
|
||||
compatibility_threshold=3.5,
|
||||
survival_threshold=0.01,
|
||||
),
|
||||
),
|
||||
problem=BraxEnv(
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 25 MiB |
@@ -17,6 +17,7 @@ if __name__ == "__main__":
|
||||
activation_options=(Act.tanh,),
|
||||
activation_default=Act.tanh,
|
||||
),
|
||||
output_transform=Act.tanh
|
||||
),
|
||||
pop_size=1000,
|
||||
species_size=10,
|
||||
|
||||
@@ -17,6 +17,7 @@ if __name__ == "__main__":
|
||||
activation_options=(Act.tanh,),
|
||||
activation_default=Act.tanh,
|
||||
),
|
||||
output_transform=Act.tanh,
|
||||
),
|
||||
pop_size=100,
|
||||
species_size=10,
|
||||
|
||||
@@ -17,6 +17,7 @@ if __name__ == "__main__":
|
||||
activation_options=(Act.tanh,),
|
||||
activation_default=Act.tanh,
|
||||
),
|
||||
output_transform=Act.tanh
|
||||
),
|
||||
pop_size=10000,
|
||||
species_size=10,
|
||||
|
||||
@@ -9,11 +9,9 @@ if __name__ == "__main__":
|
||||
pipeline = Pipeline(
|
||||
algorithm=HyperNEAT(
|
||||
substrate=FullSubstrate(
|
||||
input_coors=[(-1, -1), (0.333, -1), (-0.333, -1), (1, -1)],
|
||||
input_coors=[(-1, -1), (0.333, -1), (-0.333, -1), (1, -1)], # 3(XOR3d inputs) + 1(bias)
|
||||
hidden_coors=[
|
||||
(-1, -0.5),
|
||||
(0.333, -0.5),
|
||||
(-0.333, -0.5),
|
||||
(-1, -0.5), (0.333, -0.5), (-0.333, -0.5),
|
||||
(1, -0.5),
|
||||
(-1, 0),
|
||||
(0.333, 0),
|
||||
@@ -25,14 +23,14 @@ if __name__ == "__main__":
|
||||
(1, 0.5),
|
||||
],
|
||||
output_coors=[
|
||||
(0, 1),
|
||||
(0, 1), # one output
|
||||
],
|
||||
),
|
||||
neat=NEAT(
|
||||
species=DefaultSpecies(
|
||||
genome=DefaultGenome(
|
||||
num_inputs=4, # [-1, -1, -1, 0]
|
||||
num_outputs=1,
|
||||
num_inputs=4, # [*coor1, *coor2]
|
||||
num_outputs=1, # the weight of connection between two coor1 and coor2
|
||||
max_nodes=50,
|
||||
max_conns=100,
|
||||
node_gene=DefaultNodeGene(
|
||||
|
||||
@@ -1,53 +1,74 @@
|
||||
import jax.numpy as jnp
|
||||
import jax
|
||||
|
||||
from config import *
|
||||
from pipeline import Pipeline
|
||||
from algorithm import NEAT
|
||||
from algorithm.neat.gene import NormalGene, NormalGeneConfig
|
||||
from algorithm.hyperneat import HyperNEAT, NormalSubstrateConfig, NormalSubstrate
|
||||
from problem.rl_env import GymNaxConfig, GymNaxEnv
|
||||
|
||||
|
||||
def example_conf():
|
||||
return Config(
|
||||
basic=BasicConfig(seed=42, fitness_target=500, pop_size=10000),
|
||||
neat=NeatConfig(
|
||||
inputs=4,
|
||||
outputs=1,
|
||||
),
|
||||
gene=NormalGeneConfig(
|
||||
activation_default=Act.tanh,
|
||||
activation_options=(Act.tanh,),
|
||||
),
|
||||
hyperneat=HyperNeatConfig(activation=Act.sigmoid, inputs=4, outputs=2),
|
||||
substrate=NormalSubstrateConfig(
|
||||
input_coors=((-1, -1), (-0.5, -1), (0, -1), (0.5, -1), (1, -1)),
|
||||
hidden_coors=(
|
||||
# (-1, -0.5), (-0.5, -0.5), (0, -0.5), (0.5, -0.5),
|
||||
(1, 0),
|
||||
(-1, 0),
|
||||
(-0.5, 0),
|
||||
(0, 0),
|
||||
(0.5, 0),
|
||||
(1, 0),
|
||||
# (1, 0.5), (-1, 0.5), (-0.5, 0.5), (0, 0.5), (0.5, 0.5), (1, 0.5),
|
||||
),
|
||||
output_coors=((-1, 1), (1, 1)),
|
||||
),
|
||||
problem=GymNaxConfig(
|
||||
env_name="CartPole-v1",
|
||||
output_transform=lambda out: jnp.argmax(
|
||||
out
|
||||
), # the action of cartpole is {0, 1}
|
||||
),
|
||||
)
|
||||
from algorithm.neat import *
|
||||
from algorithm.hyperneat import *
|
||||
from utils import Act
|
||||
|
||||
from problem.rl_env import GymNaxEnv
|
||||
|
||||
if __name__ == "__main__":
|
||||
conf = example_conf()
|
||||
pipeline = Pipeline(
|
||||
algorithm=HyperNEAT(
|
||||
substrate=FullSubstrate(
|
||||
input_coors=[
|
||||
(-1, -1),
|
||||
(-0.5, -1),
|
||||
(0, -1),
|
||||
(0.5, -1),
|
||||
(1, -1),
|
||||
], # 4(problem inputs) + 1(bias)
|
||||
hidden_coors=[
|
||||
(-1, -0.5),
|
||||
(0.333, -0.5),
|
||||
(-0.333, -0.5),
|
||||
(1, -0.5),
|
||||
(-1, 0),
|
||||
(0.333, 0),
|
||||
(-0.333, 0),
|
||||
(1, 0),
|
||||
(-1, 0.5),
|
||||
(0.333, 0.5),
|
||||
(-0.333, 0.5),
|
||||
(1, 0.5),
|
||||
],
|
||||
output_coors=[
|
||||
(-1, 1),
|
||||
(1, 1), # one output
|
||||
],
|
||||
),
|
||||
neat=NEAT(
|
||||
species=DefaultSpecies(
|
||||
genome=DefaultGenome(
|
||||
num_inputs=4, # [*coor1, *coor2]
|
||||
num_outputs=1, # the weight of connection between two coor1 and coor2
|
||||
max_nodes=50,
|
||||
max_conns=100,
|
||||
node_gene=DefaultNodeGene(
|
||||
activation_default=Act.tanh,
|
||||
activation_options=(Act.tanh,),
|
||||
),
|
||||
output_transform=Act.tanh, # the activation function for output node in NEAT
|
||||
),
|
||||
pop_size=10000,
|
||||
species_size=10,
|
||||
compatibility_threshold=3.5,
|
||||
survival_threshold=0.03,
|
||||
),
|
||||
),
|
||||
activation=Act.tanh, # the activation function for output node in HyperNEAT
|
||||
activate_time=10,
|
||||
output_transform=jax.numpy.argmax, # action of cartpole is in {0, 1}
|
||||
),
|
||||
problem=GymNaxEnv(
|
||||
env_name="CartPole-v1",
|
||||
),
|
||||
generation_limit=300,
|
||||
fitness_target=500,
|
||||
)
|
||||
|
||||
algorithm = HyperNEAT(conf, NormalGene, NormalSubstrate)
|
||||
pipeline = Pipeline(conf, algorithm, GymNaxEnv)
|
||||
# initialize state
|
||||
state = pipeline.setup()
|
||||
pipeline.pre_compile(state)
|
||||
# print(state)
|
||||
# run until terminate
|
||||
state, best = pipeline.auto_run(state)
|
||||
|
||||
@@ -26,7 +26,7 @@ if __name__ == "__main__":
|
||||
env_name="MountainCar-v0",
|
||||
),
|
||||
generation_limit=10000,
|
||||
fitness_target=0,
|
||||
fitness_target=-86,
|
||||
)
|
||||
|
||||
# initialize state
|
||||
|
||||
@@ -17,6 +17,7 @@ if __name__ == "__main__":
|
||||
activation_options=(Act.tanh,),
|
||||
activation_default=Act.tanh,
|
||||
),
|
||||
output_transform=Act.tanh
|
||||
),
|
||||
pop_size=10000,
|
||||
species_size=10,
|
||||
@@ -26,7 +27,7 @@ if __name__ == "__main__":
|
||||
env_name="MountainCarContinuous-v0",
|
||||
),
|
||||
generation_limit=10000,
|
||||
fitness_target=500,
|
||||
fitness_target=99,
|
||||
)
|
||||
|
||||
# initialize state
|
||||
|
||||
@@ -17,7 +17,7 @@ if __name__ == "__main__":
|
||||
activation_options=(Act.tanh,),
|
||||
activation_default=Act.tanh,
|
||||
),
|
||||
output_transform=lambda out: out
|
||||
output_transform=lambda out: Act.tanh(out)
|
||||
* 2, # the action of pendulum is [-2, 2]
|
||||
),
|
||||
pop_size=10000,
|
||||
@@ -28,7 +28,7 @@ if __name__ == "__main__":
|
||||
env_name="Pendulum-v1",
|
||||
),
|
||||
generation_limit=10000,
|
||||
fitness_target=0,
|
||||
fitness_target=-10,
|
||||
)
|
||||
|
||||
# initialize state
|
||||
|
||||
@@ -23,7 +23,7 @@ if __name__ == "__main__":
|
||||
env_name="Reacher-misc",
|
||||
),
|
||||
generation_limit=10000,
|
||||
fitness_target=500,
|
||||
fitness_target=90,
|
||||
)
|
||||
|
||||
# initialize state
|
||||
|
||||
@@ -5,8 +5,8 @@ from .rl_jit import RLEnv
|
||||
|
||||
|
||||
class BraxEnv(RLEnv):
|
||||
def __init__(self, env_name: str = "ant", backend: str = "generalized"):
|
||||
super().__init__()
|
||||
def __init__(self, max_step=1000, env_name: str = "ant", backend: str = "generalized"):
|
||||
super().__init__(max_step)
|
||||
self.env = envs.create(env_name=env_name, backend=backend)
|
||||
|
||||
def env_step(self, randkey, env_state, action):
|
||||
|
||||
@@ -4,8 +4,8 @@ from .rl_jit import RLEnv
|
||||
|
||||
|
||||
class GymNaxEnv(RLEnv):
|
||||
def __init__(self, env_name):
|
||||
super().__init__()
|
||||
def __init__(self, env_name, max_step=1000):
|
||||
super().__init__(max_step)
|
||||
assert env_name in gymnax.registered_envs, f"Env {env_name} not registered"
|
||||
self.env, self.env_params = gymnax.make(env_name)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user