refactor folder locations
This commit is contained in:
root
0
test/__init__.py
Normal file
0
test/__init__.py
Normal file
54
test/crossover_mutation.py
Normal file
54
test/crossover_mutation.py
Normal file
@@ -0,0 +1,54 @@
|
||||
import jax, jax.numpy as jnp
|
||||
from tensorneat.common import Act
|
||||
from algorithm.neat import *
|
||||
import numpy as np
|
||||
|
||||
|
||||
def main():
|
||||
algorithm = NEAT(
|
||||
species=DefaultSpecies(
|
||||
genome=DefaultGenome(
|
||||
num_inputs=3,
|
||||
num_outputs=1,
|
||||
max_nodes=100,
|
||||
max_conns=100,
|
||||
),
|
||||
pop_size=1000,
|
||||
species_size=10,
|
||||
compatibility_threshold=3.5,
|
||||
),
|
||||
mutation=DefaultMutation(
|
||||
conn_add=0.4,
|
||||
conn_delete=0,
|
||||
node_add=0.9,
|
||||
node_delete=0,
|
||||
),
|
||||
)
|
||||
|
||||
state = algorithm.setup(jax.random.key(0))
|
||||
pop_nodes, pop_conns = algorithm.species.ask(state.species)
|
||||
|
||||
batch_transform = jax.vmap(algorithm.genome.transform)
|
||||
batch_forward = jax.vmap(algorithm.forward, in_axes=(None, 0))
|
||||
|
||||
for _ in range(50):
|
||||
winner, losser = jax.random.randint(state.randkey, (2, 1000), 0, 1000)
|
||||
elite_mask = jnp.zeros((1000,), dtype=jnp.bool_)
|
||||
elite_mask = elite_mask.at[:5].set(1)
|
||||
|
||||
state = algorithm.create_next_generation(
|
||||
jax.random.key(0), state, winner, losser, elite_mask
|
||||
)
|
||||
pop_nodes, pop_conns = algorithm.species.ask(state.species)
|
||||
|
||||
transforms = batch_transform(pop_nodes, pop_conns)
|
||||
outputs = batch_forward(jnp.array([1, 0, 1]), transforms)
|
||||
|
||||
try:
|
||||
assert not jnp.any(jnp.isnan(outputs))
|
||||
except:
|
||||
print(_)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
42
test/nan_fitness.py
Normal file
42
test/nan_fitness.py
Normal file
@@ -0,0 +1,42 @@
|
||||
import jax, jax.numpy as jnp
|
||||
from tensorneat.common import Act
|
||||
from algorithm.neat import *
|
||||
import numpy as np
|
||||
|
||||
|
||||
def main():
|
||||
node_path = "../examples/brax/nan_node.npy"
|
||||
conn_path = "../examples/brax/nan_conn.npy"
|
||||
nodes = np.load(node_path)
|
||||
conns = np.load(conn_path)
|
||||
nodes, conns = jax.device_put([nodes, conns])
|
||||
|
||||
genome = DefaultGenome(
|
||||
num_inputs=8,
|
||||
num_outputs=2,
|
||||
max_nodes=20,
|
||||
max_conns=20,
|
||||
node_gene=DefaultNodeGene(
|
||||
activation_options=(Act.tanh,),
|
||||
activation_default=Act.tanh,
|
||||
),
|
||||
)
|
||||
|
||||
transformed = genome.transform(nodes, conns)
|
||||
seq, nodes, conns = transformed
|
||||
print(seq)
|
||||
|
||||
exit(0)
|
||||
# print(*transformed, sep='\n')
|
||||
|
||||
key = jax.random.key(0)
|
||||
dummy_input = jnp.zeros((8,))
|
||||
output = genome.forward(dummy_input, transformed)
|
||||
print(output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
a = jnp.array([1, 3, 5, 6, 8])
|
||||
b = jnp.array([1, 2, 3])
|
||||
print(jnp.isin(a, b))
|
||||
# main()
|
||||
308
test/test_b_spline_conn.ipynb
Normal file
308
test/test_b_spline_conn.ipynb
Normal file
File diff suppressed because one or more lines are too long
112
test/test_compile_time.ipynb
Normal file
112
test/test_compile_time.ipynb
Normal file
@@ -0,0 +1,112 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"id": "initial_id",
|
||||
"metadata": {
|
||||
"collapsed": true,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-02T08:29:04.093990Z",
|
||||
"start_time": "2024-06-02T08:29:04.085992900Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import jax\n",
|
||||
"from jax import vmap, jit, numpy as jnp"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 26,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def func(x, y):\n",
|
||||
" return x + y\n",
|
||||
"\n",
|
||||
"def loop2():\n",
|
||||
" s = 0\n",
|
||||
" for i in range(1000):\n",
|
||||
" x = jnp.full((10000, 1), i)\n",
|
||||
" y = jnp.full((10000, 1), i + 1)\n",
|
||||
" s = (vmap(func)(x, y)).sum()\n",
|
||||
" return s\n",
|
||||
"\n",
|
||||
"def loop3():\n",
|
||||
" s = 0\n",
|
||||
" vmap_func = vmap(func)\n",
|
||||
" for i in range(1000):\n",
|
||||
" x = jnp.full((10000, 1), i)\n",
|
||||
" y = jnp.full((10000, 1), i + 1)\n",
|
||||
" s = (vmap_func(x, y)).sum()\n",
|
||||
" return s"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-02T08:31:13.023886300Z",
|
||||
"start_time": "2024-06-02T08:31:13.003026800Z"
|
||||
}
|
||||
},
|
||||
"id": "39f803029127aaa8"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 27,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array(19990000, dtype=int32)"
|
||||
},
|
||||
"execution_count": 27,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"compile_loop = jit(loop3).lower().compile()\n",
|
||||
"compile_loop()"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-06-02T08:31:14.526380100Z",
|
||||
"start_time": "2024-06-02T08:31:13.870916800Z"
|
||||
}
|
||||
},
|
||||
"id": "ab9f83d0a313f51d"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"outputs": [],
|
||||
"source": [],
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
},
|
||||
"id": "c1bd963e51aa5fd4"
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 2
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython2",
|
||||
"version": "2.7.6"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
283
test/test_efficient_b_spline.ipynb
Normal file
283
test/test_efficient_b_spline.ipynb
Normal file
File diff suppressed because one or more lines are too long
233
test/test_flatten.ipynb
Normal file
233
test/test_flatten.ipynb
Normal file
@@ -0,0 +1,233 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "initial_id",
|
||||
"metadata": {
|
||||
"collapsed": true,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T09:01:41.824974900Z",
|
||||
"start_time": "2024-05-31T09:01:39.138674100Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from algorithm.neat.genome import DefaultGenome\n",
|
||||
"from utils.tools import flatten_conns, unflatten_conns\n",
|
||||
"import jax, jax.numpy as jnp\n",
|
||||
"from jax import vmap"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "((5, 5), (5, 4))"
|
||||
},
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"genome = DefaultGenome(num_inputs=3, num_outputs=1, max_nodes=5, max_conns=5)\n",
|
||||
"state = genome.setup()\n",
|
||||
"key = jax.random.PRNGKey(0)\n",
|
||||
"nodes, conns = genome.initialize(state, key)\n",
|
||||
"nodes.shape, conns.shape"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T09:01:45.179170400Z",
|
||||
"start_time": "2024-05-31T09:01:41.832976100Z"
|
||||
}
|
||||
},
|
||||
"id": "89fb5cd0e77a028d"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(Array([0, 1, 2, 4, 3], dtype=int32, weak_type=True),\n Array([[ 0. , -1.013169 , 1. , 0. , 0. ],\n [ 1. , -0.3775248 , 1. , 0. , 0. ],\n [ 2. , 0.7407059 , 1. , 0. , 0. ],\n [ 3. , -0.66817343, 1. , 0. , 0. ],\n [ 4. , 0.5336131 , 1. , 0. , 0. ]], dtype=float32, weak_type=True),\n Array([[[ nan, nan, nan, nan,\n 0.13149254],\n [ nan, nan, nan, nan,\n 0.02001922],\n [ nan, nan, nan, nan,\n -0.79229796],\n [ nan, nan, nan, nan,\n nan],\n [ nan, nan, nan, -0.57102853,\n nan]]], dtype=float32, weak_type=True))"
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"transformed = genome.transform(state, nodes, conns)\n",
|
||||
"transformed"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T09:01:45.729969500Z",
|
||||
"start_time": "2024-05-31T09:01:45.178173400Z"
|
||||
}
|
||||
},
|
||||
"id": "aaa88227bbf29936"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(Array([[ 0. , -1.013169 , 1. , 0. , 0. ],\n [ 1. , -0.3775248 , 1. , 0. , 0. ],\n [ 2. , 0.7407059 , 1. , 0. , 0. ],\n [ 3. , -0.66817343, 1. , 0. , 0. ],\n [ 4. , 0.5336131 , 1. , 0. , 0. ]], dtype=float32, weak_type=True),\n Array([[ 1. , 0. , 4. , 0.13149254],\n [ 1. , 1. , 4. , 0.02001922],\n [ 1. , 2. , 4. , -0.79229796],\n [ 1. , 4. , 3. , -0.57102853],\n [ 1. , nan, nan, nan]], dtype=float32))"
|
||||
},
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# single flatten\n",
|
||||
"nodes, conns = genome.restore(state, transformed)\n",
|
||||
"nodes, conns"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T09:01:46.660023600Z",
|
||||
"start_time": "2024-05-31T09:01:45.724970700Z"
|
||||
}
|
||||
},
|
||||
"id": "f2c65de38fdcff8f"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array([[ 1. , 3. , 0. , 1. , 4. ,\n 0.13149254],\n [ 1. , 3. , 1. , 1. , 4. ,\n 0.02001922],\n [ 1. , 3. , 2. , 1. , 4. ,\n -0.79229796],\n [ 1. , 3. , 4. , 1. , 3. ,\n -0.57102853],\n [ 1. , 3. , nan, 1. , nan,\n nan]], dtype=float32)"
|
||||
},
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"conns = jnp.insert(conns, obj=3, values=1, axis=1)\n",
|
||||
"conns"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-31T09:03:35.665080500Z",
|
||||
"start_time": "2024-05-31T09:03:35.013654700Z"
|
||||
}
|
||||
},
|
||||
"id": "10bcb665c32fb728"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "((3, 10, 5), (3, 10, 4))"
|
||||
},
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# batch_flatten\n",
|
||||
"key = jax.random.PRNGKey(0)\n",
|
||||
"keys = jax.random.split(key, 3)\n",
|
||||
"pop_nodes, pop_conns = jax.vmap(genome.initialize, in_axes=(None, 0))(state, keys)\n",
|
||||
"pop_nodes.shape, pop_conns.shape"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T11:43:09.287012800Z",
|
||||
"start_time": "2024-05-30T11:43:09.230179800Z"
|
||||
}
|
||||
},
|
||||
"id": "fe89b178b721d656"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(3, 2, 10, 10)"
|
||||
},
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"pop_unflatten = jax.vmap(unflatten_conns)(pop_nodes, pop_conns)\n",
|
||||
"pop_unflatten.shape"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T11:43:10.004429100Z",
|
||||
"start_time": "2024-05-30T11:43:09.404949800Z"
|
||||
}
|
||||
},
|
||||
"id": "14bbb257e5ddeab"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(3, 10, 4)"
|
||||
},
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"flatten = jax.vmap(flatten_conns, in_axes=(0, 0, None))(pop_nodes, pop_unflatten, 10)\n",
|
||||
"flatten.shape"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T11:43:39.983690700Z",
|
||||
"start_time": "2024-05-30T11:43:39.208549Z"
|
||||
}
|
||||
},
|
||||
"id": "8e5cdf6140c81dc0"
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 2
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython2",
|
||||
"version": "2.7.6"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
27
test/test_genome.py
Normal file
27
test/test_genome.py
Normal file
@@ -0,0 +1,27 @@
|
||||
import jax
|
||||
from algorithm.neat import *
|
||||
|
||||
genome = DefaultGenome(
|
||||
num_inputs=3,
|
||||
num_outputs=1,
|
||||
max_nodes=5,
|
||||
max_conns=10,
|
||||
)
|
||||
|
||||
|
||||
def test_output_work():
|
||||
randkey = jax.random.PRNGKey(0)
|
||||
state = genome.setup()
|
||||
nodes, conns = genome.initialize(state, randkey)
|
||||
transformed = genome.transform(state, nodes, conns)
|
||||
inputs = jax.random.normal(randkey, (3,))
|
||||
output = genome.forward(state, transformed, inputs)
|
||||
print(output)
|
||||
|
||||
batch_inputs = jax.random.normal(randkey, (10, 3))
|
||||
batch_output = jax.vmap(genome.forward, in_axes=(None, None, 0))(
|
||||
state, transformed, batch_inputs
|
||||
)
|
||||
print(batch_output)
|
||||
|
||||
assert True
|
||||
164
test/test_kan.ipynb
Normal file
164
test/test_kan.ipynb
Normal file
File diff suppressed because one or more lines are too long
35
test/test_nan_fitness.py
Normal file
35
test/test_nan_fitness.py
Normal file
@@ -0,0 +1,35 @@
|
||||
import jax, jax.numpy as jnp
|
||||
from tensorneat.common import Act
|
||||
from algorithm.neat import *
|
||||
import numpy as np
|
||||
|
||||
|
||||
def main():
|
||||
node_path = "../examples/brax/nan_node.npy"
|
||||
conn_path = "../examples/brax/nan_conn.npy"
|
||||
nodes = np.load(node_path)
|
||||
conns = np.load(conn_path)
|
||||
nodes, conns = jax.device_put([nodes, conns])
|
||||
|
||||
genome = DefaultGenome(
|
||||
num_inputs=8,
|
||||
num_outputs=2,
|
||||
max_nodes=20,
|
||||
max_conns=20,
|
||||
node_gene=DefaultNodeGene(
|
||||
activation_options=(Act.tanh,),
|
||||
activation_default=Act.tanh,
|
||||
),
|
||||
)
|
||||
|
||||
transformed = genome.transform(nodes, conns)
|
||||
print(*transformed, sep="\n")
|
||||
|
||||
key = jax.random.key(0)
|
||||
dummy_input = jnp.zeros((8,))
|
||||
output = genome.forward(dummy_input, transformed)
|
||||
print(output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
458
test/test_record_episode.ipynb
Normal file
458
test/test_record_episode.ipynb
Normal file
@@ -0,0 +1,458 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "initial_id",
|
||||
"metadata": {
|
||||
"collapsed": true,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:53:04.429593300Z",
|
||||
"start_time": "2024-05-30T08:53:02.326728600Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import jax, jax.numpy as jnp\n",
|
||||
"from tensorneat.utils import State\n",
|
||||
"from problem.rl_env import BraxEnv\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def random_policy(state: State, obs, randkey):\n",
|
||||
" return jax.random.uniform(randkey, (8,)) * 2 - 1"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array(24.975231, dtype=float32)"
|
||||
},
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# single evaluation without recording episode\n",
|
||||
"randkey = jax.random.key(0)\n",
|
||||
"env_key, policy_key = jax.random.split(randkey)\n",
|
||||
"problem = BraxEnv(env_name=\"ant\", max_step=100)\n",
|
||||
"state = problem.setup()\n",
|
||||
"evaluate_using_random_policy_without_record = lambda state, env_key, policy_key: problem.evaluate(state, env_key, random_policy,\n",
|
||||
" policy_key)\n",
|
||||
"score = jax.jit(evaluate_using_random_policy_without_record)(state, env_key, policy_key)\n",
|
||||
"score"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:53:18.928839600Z",
|
||||
"start_time": "2024-05-30T08:53:04.435561800Z"
|
||||
}
|
||||
},
|
||||
"id": "e62882e782d7e54e"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array([ -3.274895 , -6.016205 , -6.9032974, 9.187286 ,\n -120.19688 , 12.389805 , -4.6393256, -50.27197 ,\n 9.650737 , -73.77956 ], dtype=float32)"
|
||||
},
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# batch evaluation without recording episode\n",
|
||||
"batch = 10\n",
|
||||
"env_keys = jax.random.split(env_key, batch)\n",
|
||||
"policy_keys = jax.random.split(policy_key, batch)\n",
|
||||
"\n",
|
||||
"score = jax.jit(\n",
|
||||
" jax.vmap(\n",
|
||||
" evaluate_using_random_policy_without_record, \n",
|
||||
" in_axes=(None, 0, 0)\n",
|
||||
" ))(\n",
|
||||
" state, env_keys, policy_keys\n",
|
||||
" )\n",
|
||||
"score"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:53:29.458960600Z",
|
||||
"start_time": "2024-05-30T08:53:18.928839600Z"
|
||||
}
|
||||
},
|
||||
"id": "d01997be61038ea2"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(Array(18.354952, dtype=float32), (100, 27), (100, 8), (100,))"
|
||||
},
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# single evaluation with recording episode\n",
|
||||
"randkey = jax.random.key(0)\n",
|
||||
"env_key, policy_key = jax.random.split(randkey)\n",
|
||||
"problem = BraxEnv(env_name=\"ant\", max_step=100, record_episode=True)\n",
|
||||
"evaluate_using_random_policy_with_record = lambda state, env_key, policy_key: problem.evaluate(state, env_key, random_policy,\n",
|
||||
" policy_key)\n",
|
||||
"score, episode = jax.jit(evaluate_using_random_policy_with_record)(state, env_key, policy_key)\n",
|
||||
"score, episode[\"obs\"].shape, episode[\"action\"].shape, episode[\"reward\"].shape"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:53:40.372461Z",
|
||||
"start_time": "2024-05-30T08:53:29.455962200Z"
|
||||
}
|
||||
},
|
||||
"id": "ac6f72e21dd12ee8"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(Array(18.354952, dtype=float32), (10, 100, 27), (10, 100, 8), (10, 100))"
|
||||
},
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# batch evaluation without recording episode\n",
|
||||
"batch = 10\n",
|
||||
"env_keys = jax.random.split(env_key, batch)\n",
|
||||
"policy_keys = jax.random.split(policy_key, batch)\n",
|
||||
"\n",
|
||||
"scores, episodes = jax.jit(\n",
|
||||
" jax.vmap(\n",
|
||||
" evaluate_using_random_policy_with_record, \n",
|
||||
" in_axes=(None, 0, 0)\n",
|
||||
" ))(\n",
|
||||
" state, env_keys, policy_keys\n",
|
||||
" )\n",
|
||||
"score, episodes[\"obs\"].shape, episodes[\"action\"].shape, episodes[\"reward\"].shape"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:53:51.261470500Z",
|
||||
"start_time": "2024-05-30T08:53:40.368462Z"
|
||||
}
|
||||
},
|
||||
"id": "1c55341b054ee2e8"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array(18.354952, dtype=float32)"
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"evaluate_using_random_policy_with_record = jax.jit(evaluate_using_random_policy_with_record)\n",
|
||||
"evaluate_using_random_policy_without_record = jax.jit(evaluate_using_random_policy_without_record)\n",
|
||||
"evaluate_using_random_policy_with_record(state, env_key, policy_key)\n",
|
||||
"evaluate_using_random_policy_without_record(state, env_key, policy_key)"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:53:55.402886Z",
|
||||
"start_time": "2024-05-30T08:53:51.255470600Z"
|
||||
}
|
||||
},
|
||||
"id": "274ca4fd0d0b8663"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for _ in range(20):\n",
|
||||
" evaluate_using_random_policy_with_record(state, env_key, policy_key)\n",
|
||||
"# 47s384ms"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:54:42.782425800Z",
|
||||
"start_time": "2024-05-30T08:53:55.397887700Z"
|
||||
}
|
||||
},
|
||||
"id": "fdb34361d19cb78d"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for _ in range(20):\n",
|
||||
" evaluate_using_random_policy_without_record(state, env_key, policy_key)\n",
|
||||
"# 48s559ms"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:55:31.344699500Z",
|
||||
"start_time": "2024-05-30T08:54:42.785428500Z"
|
||||
}
|
||||
},
|
||||
"id": "9afdf6923051c9f1"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array(9., dtype=float32, weak_type=True)"
|
||||
},
|
||||
"execution_count": 12,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# single evaluation without recording episode\n",
|
||||
"from problem.rl_env import GymNaxEnv\n",
|
||||
"\n",
|
||||
"def random_policy(state: State, obs, randkey):\n",
|
||||
" return jax.random.uniform(randkey, ()) \n",
|
||||
"\n",
|
||||
"randkey = jax.random.key(0)\n",
|
||||
"env_key, policy_key = jax.random.split(randkey)\n",
|
||||
"problem = GymNaxEnv(env_name=\"CartPole-v1\", max_step=500)\n",
|
||||
"state = problem.setup()\n",
|
||||
"evaluate_using_random_policy_without_record = lambda state, env_key, policy_key: problem.evaluate(state, env_key, random_policy,\n",
|
||||
" policy_key)\n",
|
||||
"score = jax.jit(evaluate_using_random_policy_without_record)(state, env_key, policy_key)\n",
|
||||
"score"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:58:46.652406400Z",
|
||||
"start_time": "2024-05-30T08:58:45.606288800Z"
|
||||
}
|
||||
},
|
||||
"id": "1de25fb23f519284"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array([13., 19., 11., 12., 14., 21., 13., 11., 11., 28.], dtype=float32, weak_type=True)"
|
||||
},
|
||||
"execution_count": 13,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# batch evaluation without recording episode\n",
|
||||
"batch = 10\n",
|
||||
"env_keys = jax.random.split(env_key, batch)\n",
|
||||
"policy_keys = jax.random.split(policy_key, batch)\n",
|
||||
"\n",
|
||||
"score = jax.jit(\n",
|
||||
" jax.vmap(\n",
|
||||
" evaluate_using_random_policy_without_record, \n",
|
||||
" in_axes=(None, 0, 0)\n",
|
||||
" ))(\n",
|
||||
" state, env_keys, policy_keys\n",
|
||||
" )\n",
|
||||
"score"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:58:58.323528300Z",
|
||||
"start_time": "2024-05-30T08:58:57.272024400Z"
|
||||
}
|
||||
},
|
||||
"id": "99e745dce6f2872d"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 14,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(Array(9., dtype=float32, weak_type=True), (500, 4), (500,), (500,))"
|
||||
},
|
||||
"execution_count": 14,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# single evaluation with recording episode\n",
|
||||
"randkey = jax.random.key(0)\n",
|
||||
"env_key, policy_key = jax.random.split(randkey)\n",
|
||||
"problem = GymNaxEnv(env_name=\"CartPole-v1\", max_step=500, record_episode=True)\n",
|
||||
"evaluate_using_random_policy_with_record = lambda state, env_key, policy_key: problem.evaluate(state, env_key, random_policy,\n",
|
||||
" policy_key)\n",
|
||||
"score, episode = jax.jit(evaluate_using_random_policy_with_record)(state, env_key, policy_key)\n",
|
||||
"score, episode[\"obs\"].shape, episode[\"action\"].shape, episode[\"reward\"].shape"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:59:18.830495600Z",
|
||||
"start_time": "2024-05-30T08:59:17.568087200Z"
|
||||
}
|
||||
},
|
||||
"id": "257e340ebf24c10d"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 15,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(Array(9., dtype=float32, weak_type=True), (10, 500, 4), (10, 500), (10, 500))"
|
||||
},
|
||||
"execution_count": 15,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# batch evaluation without recording episode\n",
|
||||
"batch = 10\n",
|
||||
"env_keys = jax.random.split(env_key, batch)\n",
|
||||
"policy_keys = jax.random.split(policy_key, batch)\n",
|
||||
"\n",
|
||||
"scores, episodes = jax.jit(\n",
|
||||
" jax.vmap(\n",
|
||||
" evaluate_using_random_policy_with_record, \n",
|
||||
" in_axes=(None, 0, 0)\n",
|
||||
" ))(\n",
|
||||
" state, env_keys, policy_keys\n",
|
||||
" )\n",
|
||||
"score, episodes[\"obs\"].shape, episodes[\"action\"].shape, episodes[\"reward\"].shape"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:59:34.182539200Z",
|
||||
"start_time": "2024-05-30T08:59:32.956339600Z"
|
||||
}
|
||||
},
|
||||
"id": "9ba8dc68085cd0fc"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 16,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array(9., dtype=float32, weak_type=True)"
|
||||
},
|
||||
"execution_count": 16,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"evaluate_using_random_policy_with_record = jax.jit(evaluate_using_random_policy_with_record)\n",
|
||||
"evaluate_using_random_policy_without_record = jax.jit(evaluate_using_random_policy_without_record)\n",
|
||||
"evaluate_using_random_policy_with_record(state, env_key, policy_key)\n",
|
||||
"evaluate_using_random_policy_without_record(state, env_key, policy_key)"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T08:59:46.472504900Z",
|
||||
"start_time": "2024-05-30T08:59:46.419192900Z"
|
||||
}
|
||||
},
|
||||
"id": "ea01b6663a7ca076"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 19,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for _ in range(20):\n",
|
||||
" evaluate_using_random_policy_with_record(state, env_key, policy_key)\n",
|
||||
"# 48ms"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T09:00:18.905094200Z",
|
||||
"start_time": "2024-05-30T09:00:18.809970900Z"
|
||||
}
|
||||
},
|
||||
"id": "989c39c8e20779d0"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 20,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for _ in range(20):\n",
|
||||
" evaluate_using_random_policy_without_record(state, env_key, policy_key)\n",
|
||||
"# 43ms"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T09:00:19.240415900Z",
|
||||
"start_time": "2024-05-30T09:00:19.190416700Z"
|
||||
}
|
||||
},
|
||||
"id": "bab4782fe674f2d5"
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 2
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython2",
|
||||
"version": "2.7.6"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
0
test/test_record_episode.py
Normal file
0
test/test_record_episode.py
Normal file
317
test/test_update_by_batch.ipynb
Normal file
317
test/test_update_by_batch.ipynb
Normal file
@@ -0,0 +1,317 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"id": "initial_id",
|
||||
"metadata": {
|
||||
"collapsed": true,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:07:59.805322900Z",
|
||||
"start_time": "2024-05-30T15:07:57.075364700Z"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import jax, jax.numpy as jnp\n",
|
||||
"from algorithm.neat.genome import *\n",
|
||||
"from algorithm.neat.gene import *\n",
|
||||
"\n",
|
||||
"jnp.set_printoptions(precision=2, linewidth=150)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# genome = DefaultGenome(num_inputs=3, num_outputs=2, max_nodes=10, max_conns=10)\n",
|
||||
"# state = genome.setup()\n",
|
||||
"# randkey = jax.random.key(0)\n",
|
||||
"# genome_key, input_key = jax.random.split(randkey)\n",
|
||||
"# nodes, conns = genome.initialize(state, genome_key)\n",
|
||||
"# inputs = jax.random.normal(input_key, (10, 3)) * 2 + 1 # std: 2, mean: 1\n",
|
||||
"# print(nodes, conns, sep='\\n')\n",
|
||||
"# print(inputs)"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:07:59.817325200Z",
|
||||
"start_time": "2024-05-30T15:07:59.809324300Z"
|
||||
}
|
||||
},
|
||||
"id": "c81fa2df52f01d93"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# transformed = genome.transform(state, nodes, conns)\n",
|
||||
"# batch_output = jax.vmap(genome.forward, in_axes=(None, 0, None))(state, inputs, transformed)\n",
|
||||
"# batch_output, transformed"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:07:59.817950Z",
|
||||
"start_time": "2024-05-30T15:07:59.812323Z"
|
||||
}
|
||||
},
|
||||
"id": "d4b9aa0449c8d706"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# batch_output2, new_transformed = genome.update_by_batch(state, inputs, transformed)\n",
|
||||
"# batch_output2, new_transformed"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:07:59.831323800Z",
|
||||
"start_time": "2024-05-30T15:07:59.821324100Z"
|
||||
}
|
||||
},
|
||||
"id": "d32986470dad3229"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# assert jnp.allclose(new_transformed[0], transformed[0], equal_nan=True)\n",
|
||||
"# assert jnp.allclose(new_transformed[1], transformed[1], equal_nan=True)\n",
|
||||
"# assert jnp.allclose(new_transformed[2], transformed[2], equal_nan=True)"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:07:59.832325200Z",
|
||||
"start_time": "2024-05-30T15:07:59.826324400Z"
|
||||
}
|
||||
},
|
||||
"id": "3c4007dfd6770faf"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[[ 0. 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [ 1. 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [ 2. 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [ 3. 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [ 4. 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [ 5. 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [nan 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [nan 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [nan 0. 0. 0. 0. 1. 1. 0.]\n",
|
||||
" [nan 0. 0. 0. 0. 1. 1. 0.]]\n",
|
||||
"[[ 0. 5. 1. 1.]\n",
|
||||
" [ 1. 5. 1. 1.]\n",
|
||||
" [ 2. 5. 1. 1.]\n",
|
||||
" [ 5. 3. 1. 1.]\n",
|
||||
" [ 5. 4. 1. 1.]\n",
|
||||
" [nan nan nan 1.]\n",
|
||||
" [nan nan nan 1.]\n",
|
||||
" [nan nan nan 1.]\n",
|
||||
" [nan nan nan 1.]\n",
|
||||
" [nan nan nan 1.]]\n",
|
||||
"[[-1.9 -3.53 0.94]\n",
|
||||
" [ 2.92 0.06 3.44]\n",
|
||||
" [-0.9 -0.06 2.94]\n",
|
||||
" ...\n",
|
||||
" [ 2.07 -1.43 1.55]\n",
|
||||
" [ 1.93 2.85 0.19]\n",
|
||||
" [ 0.91 -0.65 1.86]]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(Array([ 0, 1, 2, 5, 3, 4, 2147483647, 2147483647, 2147483647, 2147483647], dtype=int32, weak_type=True),\n Array([[ 0., 0., 0., 0., 0., 1., 1., 0.],\n [ 1., 0., 0., 0., 0., 1., 1., 0.],\n [ 2., 0., 0., 0., 0., 1., 1., 0.],\n [ 3., 0., 0., 0., 0., 1., 1., 0.],\n [ 4., 0., 0., 0., 0., 1., 1., 0.],\n [ 5., 0., 0., 0., 0., 1., 1., 0.],\n [nan, 0., 0., 0., 0., 1., 1., 0.],\n [nan, 0., 0., 0., 0., 1., 1., 0.],\n [nan, 0., 0., 0., 0., 1., 1., 0.],\n [nan, 0., 0., 0., 0., 1., 1., 0.]], dtype=float32, weak_type=True),\n Array([[[nan, nan, nan, nan, nan, 1., nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, 1., nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, 1., nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, 1., 1., nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan]]], dtype=float32, weak_type=True))"
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from algorithm.neat.gene.node.normalized import NormalizedNode\n",
|
||||
"from algorithm.neat.gene.conn import DefaultConnGene\n",
|
||||
"from tensorneat.utils import Act\n",
|
||||
"\n",
|
||||
"genome = DefaultGenome(num_inputs=3, num_outputs=2, max_nodes=10, max_conns=10,\n",
|
||||
" node_gene=NormalizedNode(activation_default=Act.identity, activation_options=(Act.identity,)),\n",
|
||||
" conn_gene=DefaultConnGene(weight_init_mean=1))\n",
|
||||
"state = genome.setup()\n",
|
||||
"randkey = jax.random.key(0)\n",
|
||||
"genome_key, input_key = jax.random.split(randkey)\n",
|
||||
"nodes, conns = genome.initialize(state, genome_key)\n",
|
||||
"nodes = nodes.at[:, 1:].set(genome.node_gene.new_custom_attrs(state))\n",
|
||||
"conns = conns.at[:, 3:].set(genome.conn_gene.new_custom_attrs(state))\n",
|
||||
"\n",
|
||||
"inputs = jax.random.normal(input_key, (10000, 3)) * 2 + 1 # std: 2, mean: 1\n",
|
||||
"print(nodes, conns, sep='\\n')\n",
|
||||
"print(inputs)\n",
|
||||
"transformed = genome.transform(state, nodes, conns)\n",
|
||||
"transformed"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:08:04.532243100Z",
|
||||
"start_time": "2024-05-30T15:07:59.832325200Z"
|
||||
}
|
||||
},
|
||||
"id": "da73909c3414366e"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array([[-4.49, -4.49],\n [ 6.42, 6.42],\n [ 1.98, 1.98],\n ...,\n [ 2.19, 2.19],\n [ 4.97, 4.97],\n [ 2.12, 2.12]], dtype=float32, weak_type=True)"
|
||||
},
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"batch_output2 = jax.vmap(genome.forward, in_axes=(None, 0, None))(state, inputs, transformed)\n",
|
||||
"batch_output2"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:08:04.901593900Z",
|
||||
"start_time": "2024-05-30T15:08:04.527245300Z"
|
||||
}
|
||||
},
|
||||
"id": "8ef2402bc4c7908d"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"batch_z: [-4.49 6.42 1.98 ... 2.19 4.97 2.12]\n",
|
||||
"batch_z_mean: 2.9496588706970215\n",
|
||||
"batch_z: [-2.15 1. -0.28 ... -0.22 0.58 -0.24]\n",
|
||||
"batch_z_mean: -2.1362303925798187e-08\n",
|
||||
"batch_z: [-2.15 1. -0.28 ... -0.22 0.58 -0.24]\n",
|
||||
"batch_z_mean: -2.1362303925798187e-08\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"batch_output, new_transformed = genome.update_by_batch(state, inputs, transformed)"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:08:05.269935400Z",
|
||||
"start_time": "2024-05-30T15:08:04.899594200Z"
|
||||
}
|
||||
},
|
||||
"id": "b3c085c7ca28f127"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "(Array([[-2.15, -2.15],\n [ 1. , 1. ],\n [-0.28, -0.28],\n ...,\n [-0.22, -0.22],\n [ 0.58, 0.58],\n [-0.24, -0.24]], dtype=float32, weak_type=True),\n (Array([ 0, 1, 2, 5, 3, 4, 2147483647, 2147483647, 2147483647, 2147483647], dtype=int32, weak_type=True),\n Array([[ 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 1.00e+00, 1.00e+00, 0.00e+00],\n [ 1.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 1.00e+00, 1.00e+00, 0.00e+00],\n [ 2.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 1.00e+00, 1.00e+00, 0.00e+00],\n [ 3.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, -2.14e-08, 1.00e+00, 1.00e+00, 0.00e+00],\n [ 4.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, -2.14e-08, 1.00e+00, 1.00e+00, 0.00e+00],\n [ 5.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 2.95e+00, 3.46e+00, 1.00e+00, 0.00e+00],\n [ nan, 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 1.00e+00, 1.00e+00, 0.00e+00],\n [ nan, 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 1.00e+00, 1.00e+00, 0.00e+00],\n [ nan, 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 1.00e+00, 1.00e+00, 0.00e+00],\n [ nan, 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00, 1.00e+00, 1.00e+00, 0.00e+00]], dtype=float32, weak_type=True),\n Array([[[nan, nan, nan, nan, nan, 1., nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, 1., nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, 1., nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, 1., 1., nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],\n [nan, nan, nan, nan, nan, nan, nan, nan, nan, nan]]], dtype=float32, weak_type=True)))"
|
||||
},
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"batch_output, new_transformed"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:08:05.270935800Z",
|
||||
"start_time": "2024-05-30T15:08:05.261936200Z"
|
||||
}
|
||||
},
|
||||
"id": "60ce6747ebd95e10"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": "Array([[-2.15, -2.15],\n [ 1. , 1. ],\n [-0.28, -0.28],\n ...,\n [-0.22, -0.22],\n [ 0.58, 0.58],\n [-0.24, -0.24]], dtype=float32, weak_type=True)"
|
||||
},
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"batch_output2 = jax.vmap(genome.forward, in_axes=(None, 0, None))(state, inputs, new_transformed)\n",
|
||||
"batch_output2"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:08:05.415934Z",
|
||||
"start_time": "2024-05-30T15:08:05.269935400Z"
|
||||
}
|
||||
},
|
||||
"id": "7b092224d8f33b7"
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"outputs": [],
|
||||
"source": [],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"ExecuteTime": {
|
||||
"end_time": "2024-05-30T15:08:05.416935400Z",
|
||||
"start_time": "2024-05-30T15:08:05.405934300Z"
|
||||
}
|
||||
},
|
||||
"id": "eec974242eb3867e"
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 2
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython2",
|
||||
"version": "2.7.6"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
Reference in New Issue
Block a user