tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

odify genome for the official release

Browse Source
This commit is contained in:
root
2024-07-10 11:24:11 +08:00
parent 075460f896
commit ee8ec84202
83 changed files with 588 additions and 611 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
examples/brax/ant.py Normal file
View File

@@ -0,0 +1,39 @@
from pipeline import Pipeline
from algorithm.neat import *
from problem.rl_env import BraxEnv
from tensorneat.common import Act
if __name__ == "__main__":
pipeline = Pipeline(
algorithm=NEAT(
species=DefaultSpecies(
genome=DefaultGenome(
num_inputs=27,
num_outputs=8,
max_nodes=100,
max_conns=200,
node_gene=DefaultNodeGene(
activation_options=(Act.tanh,),
activation_default=Act.tanh,
),
output_transform=Act.tanh,
),
pop_size=1000,
species_size=10,
compatibility_threshold=3.5,
survival_threshold=0.01,
),
),
problem=BraxEnv(
env_name="ant",
),
generation_limit=10000,
fitness_target=5000,
)
# initialize state
state = pipeline.setup()
# print(state)
# run until terminate
state, best = pipeline.auto_run(state)

48
examples/brax/half_cheetah.py Normal file
View File

@@ -0,0 +1,48 @@
import jax
from pipeline import Pipeline
from algorithm.neat import *
from problem.rl_env import BraxEnv
from tensorneat.common import Act
def sample_policy(randkey, obs):
return jax.random.uniform(randkey, (6,), minval=-1, maxval=1)
if __name__ == "__main__":
pipeline = Pipeline(
algorithm=NEAT(
species=DefaultSpecies(
genome=DefaultGenome(
num_inputs=17,
num_outputs=6,
max_nodes=50,
max_conns=100,
node_gene=DefaultNodeGene(
activation_options=(Act.tanh,),
activation_default=Act.tanh,
),
output_transform=Act.tanh,
),
pop_size=1000,
species_size=10,
),
),
problem=BraxEnv(
env_name="halfcheetah",
max_step=1000,
obs_normalization=True,
sample_episodes=1000,
sample_policy=sample_policy,
),
generation_limit=10000,
fitness_target=5000,
)
# initialize state
state = pipeline.setup()
# print(state)
# run until terminate
state, best = pipeline.auto_run(state)

37
examples/brax/reacher.py Normal file
View File

@@ -0,0 +1,37 @@
from pipeline import Pipeline
from algorithm.neat import *
from problem.rl_env import BraxEnv
from tensorneat.common import Act
if __name__ == "__main__":
pipeline = Pipeline(
algorithm=NEAT(
species=DefaultSpecies(
genome=DefaultGenome(
num_inputs=11,
num_outputs=2,
max_nodes=50,
max_conns=100,
node_gene=DefaultNodeGene(
activation_options=(Act.tanh,),
activation_default=Act.tanh,
),
output_transform=Act.tanh,
),
pop_size=100,
species_size=10,
),
),
problem=BraxEnv(
env_name="reacher",
),
generation_limit=10000,
fitness_target=5000,
)
# initialize state
state = pipeline.setup()
# print(state)
# run until terminate
state, best = pipeline.auto_run(state)

19
examples/brax/show_test.py Normal file
View File

@@ -0,0 +1,19 @@
import jax
from problem.rl_env import BraxEnv
def random_policy(randkey, forward_func, obs):
return jax.random.uniform(randkey, (6,), minval=-1, maxval=1)
if __name__ == "__main__":
problem = BraxEnv(env_name="walker2d", max_step=1000, action_policy=random_policy)
state = problem.setup()
randkey = jax.random.key(0)
problem.show(
state,
randkey,
act_func=lambda state, params, obs: obs,
params=None,
save_path="walker2d_random_policy",
)

60
examples/brax/walker.py Normal file
View File

@@ -0,0 +1,60 @@
from pipeline import Pipeline
from algorithm.neat import *
from problem.rl_env import BraxEnv
from tensorneat.common import Act
import jax, jax.numpy as jnp
def split_right_left(randkey, forward_func, obs):
right_obs_keys = jnp.array([2, 3, 4, 11, 12, 13])
left_obs_keys = jnp.array([5, 6, 7, 14, 15, 16])
right_action_keys = jnp.array([0, 1, 2])
left_action_keys = jnp.array([3, 4, 5])
right_foot_obs = obs
left_foot_obs = obs
left_foot_obs = left_foot_obs.at[right_obs_keys].set(obs[left_obs_keys])
left_foot_obs = left_foot_obs.at[left_obs_keys].set(obs[right_obs_keys])
right_action, left_action = jax.vmap(forward_func)(jnp.stack([right_foot_obs, left_foot_obs]))
# print(right_action.shape)
# print(left_action.shape)
return jnp.concatenate([right_action, left_action])
if __name__ == "__main__":
pipeline = Pipeline(
algorithm=NEAT(
species=DefaultSpecies(
genome=DefaultGenome(
num_inputs=17,
num_outputs=3,
max_nodes=50,
max_conns=100,
node_gene=DefaultNodeGene(
activation_options=(Act.tanh,),
activation_default=Act.tanh,
),
output_transform=Act.tanh,
),
pop_size=1000,
species_size=10,
),
),
problem=BraxEnv(
env_name="walker2d",
max_step=1000,
action_policy=split_right_left
),
generation_limit=10000,
fitness_target=5000,
)
# initialize state
state = pipeline.setup()
# print(state)
# run until terminate
state, best = pipeline.auto_run(state)