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

complete fully stateful!

Browse Source 
use black to format all files!
This commit is contained in:
wls2002
2024-05-26 18:08:43 +08:00
parent cf69b916af
commit 18c3d44c79
41 changed files with 620 additions and 495 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tensorneat/problem/base.py
View File

@@ -10,7 +10,7 @@ class BaseProblem:
"""initialize the state of the problem"""
return state
def evaluate(self, randkey, state: State, act_func: Callable, params):
def evaluate(self, state: State, randkey, act_func: Callable, params):
"""evaluate one individual"""
raise NotImplementedError
@@ -32,7 +32,7 @@ class BaseProblem:
"""
raise NotImplementedError
def show(self, randkey, state: State, act_func: Callable, params, *args, **kwargs):
def show(self, state: State, randkey, act_func: Callable, params, *args, **kwargs):
"""
show how a genome perform in this problem
"""

30
tensorneat/problem/func_fit/func_fit.py
View File

@@ -8,42 +8,44 @@ from .. import BaseProblem
class FuncFit(BaseProblem):
jitable = True
def __init__(self,
error_method: str = 'mse'
):
def __init__(self, error_method: str = "mse"):
super().__init__()
assert error_method in {'mse', 'rmse', 'mae', 'mape'}
assert error_method in {"mse", "rmse", "mae", "mape"}
self.error_method = error_method
def setup(self, state: State = State()):
return state
def evaluate(self, randkey, state, act_func, params):
def evaluate(self, state, randkey, act_func, params):
state, predict = jax.vmap(act_func, in_axes=(None, 0, None), out_axes=(None, 0))(state, self.inputs, params)
predict = jax.vmap(act_func, in_axes=(None, 0, None))(
state, self.inputs, params
)
if self.error_method == 'mse':
if self.error_method == "mse":
loss = jnp.mean((predict - self.targets) ** 2)
elif self.error_method == 'rmse':
elif self.error_method == "rmse":
loss = jnp.sqrt(jnp.mean((predict - self.targets) ** 2))
elif self.error_method == 'mae':
elif self.error_method == "mae":
loss = jnp.mean(jnp.abs(predict - self.targets))
elif self.error_method == 'mape':
elif self.error_method == "mape":
loss = jnp.mean(jnp.abs((predict - self.targets) / self.targets))
else:
raise NotImplementedError
return state, -loss
return -loss
def show(self, randkey, state, act_func, params, *args, **kwargs):
state, predict = jax.vmap(act_func, in_axes=(None, 0, None), out_axes=(None, 0))(state, self.inputs, params)
def show(self, state, randkey, act_func, params, *args, **kwargs):
predict = jax.vmap(act_func, in_axes=(None, 0, None))(
state, self.inputs, params
)
inputs, target, predict = jax.device_get([self.inputs, self.targets, predict])
state, loss = self.evaluate(randkey, state, act_func, params)
loss = self.evaluate(state, randkey, act_func, params)
loss = -loss
msg = ""

17
tensorneat/problem/func_fit/xor.py
View File

@@ -4,27 +4,16 @@ from .func_fit import FuncFit
class XOR(FuncFit):
def __init__(self, error_method: str = 'mse'):
def __init__(self, error_method: str = "mse"):
super().__init__(error_method)
@property
def inputs(self):
return np.array([
[0, 0],
[0, 1],
[1, 0],
[1, 1]
])
return np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
@property
def targets(self):
return np.array([
[0],
[1],
[1],
[0]
])
return np.array([[0], [1], [1], [0]])
@property
def input_shape(self):

36
tensorneat/problem/func_fit/xor3d.py
View File

@@ -4,35 +4,27 @@ from .func_fit import FuncFit
class XOR3d(FuncFit):
def __init__(self, error_method: str = 'mse'):
def __init__(self, error_method: str = "mse"):
super().__init__(error_method)
@property
def inputs(self):
return np.array([
[0, 0, 0],
[0, 0, 1],
[0, 1, 0],
[0, 1, 1],
[1, 0, 0],
[1, 0, 1],
[1, 1, 0],
[1, 1, 1],
])
return np.array(
[
[0, 0, 0],
[0, 0, 1],
[0, 1, 0],
[0, 1, 1],
[1, 0, 0],
[1, 0, 1],
[1, 1, 0],
[1, 1, 1],
]
)
@property
def targets(self):
return np.array([
[0],
[1],
[1],
[0],
[1],
[0],
[0],
[1]
])
return np.array([[0], [1], [1], [0], [1], [0], [0], [1]])
@property
def input_shape(self):

24
tensorneat/problem/rl_env/brax_env.py
View File

@@ -25,7 +25,19 @@ class BraxEnv(RLEnv):
def output_shape(self):
return (self.env.action_size,)
def show(self, randkey, state, act_func, params, save_path=None, height=512, width=512, duration=0.1, *args, **kwargs):
def show(
self,
state,
randkey,
act_func,
params,
save_path=None,
height=512,
width=512,
duration=0.1,
*args,
**kwargs
):
import jax
import imageio
@@ -48,11 +60,13 @@ class BraxEnv(RLEnv):
key, env_state, obs, r, done = jax.jit(step)(randkey, env_state, obs)
reward += r
imgs = [image.render_array(sys=self.env.sys, state=s, width=width, height=height) for s in
tqdm(state_histories, desc="Rendering")]
imgs = [
image.render_array(sys=self.env.sys, state=s, width=width, height=height)
for s in tqdm(state_histories, desc="Rendering")
]
def create_gif(image_list, gif_name, duration):
with imageio.get_writer(gif_name, mode='I', duration=duration) as writer:
with imageio.get_writer(gif_name, mode="I", duration=duration) as writer:
for image in image_list:
formatted_image = np.array(image, dtype=np.uint8)
writer.append_data(formatted_image)
@@ -60,5 +74,3 @@ class BraxEnv(RLEnv):
create_gif(imgs, save_path, duration=0.1)
print("Gif saved to: ", save_path)
print("Total reward: ", reward)

3
tensorneat/problem/rl_env/gymnax_env.py
View File

@@ -4,7 +4,6 @@ from .rl_jit import RLEnv
class GymNaxEnv(RLEnv):
def __init__(self, env_name):
super().__init__()
assert env_name in gymnax.registered_envs, f"Env {env_name} not registered"
@@ -24,5 +23,5 @@ class GymNaxEnv(RLEnv):
def output_shape(self):
return self.env.action_space(self.env_params).shape
def show(self, randkey, state, act_func, params, *args, **kwargs):
def show(self, state, randkey, act_func, params, *args, **kwargs):
raise NotImplementedError("GymNax render must rely on gym 0.19.0(old version).")

30
tensorneat/problem/rl_env/rl_jit.py
View File

@@ -12,29 +12,29 @@ class RLEnv(BaseProblem):
super().__init__()
self.max_step = max_step
def evaluate(self, randkey, state, act_func, params):
def evaluate(self, state, randkey, act_func, params):
rng_reset, rng_episode = jax.random.split(randkey)
init_obs, init_env_state = self.reset(rng_reset)
def cond_func(carry):
_, _, _, _, done, _, count = carry
return ~done & (count < self.max_step)
_, _, _, done, _, count = carry
return ~done & (count < self.max_step)
def body_func(carry):
state_, obs, env_state, rng, done, tr, count = carry # tr -> total reward
state_, action = act_func(state_, obs, params)
next_obs, next_env_state, reward, done, _ = self.step(rng, env_state, action)
next_rng, _ = jax.random.split(rng)
return state_, next_obs, next_env_state, next_rng, done, tr + reward, count + 1
obs, env_state, rng, done, tr, count = carry # tr -> total reward
action = act_func(state, obs, params)
next_obs, next_env_state, reward, done, _ = self.step(
rng, env_state, action
)
next_rng, _ = jax.random.split(rng)
return next_obs, next_env_state, next_rng, done, tr + reward, count + 1
state, _, _, _, _, total_reward, _ = jax.lax.while_loop(
cond_func,
body_func,
(state, init_obs, init_env_state, rng_episode, False, 0.0, 0)
_, _, _, _, total_reward, _ = jax.lax.while_loop(
cond_func, body_func, (init_obs, init_env_state, rng_episode, False, 0.0, 0)
)
return state, total_reward
return total_reward
# @partial(jax.jit, static_argnums=(0,))
def step(self, randkey, env_state, action):
return self.env_step(randkey, env_state, action)
@@ -57,5 +57,5 @@ class RLEnv(BaseProblem):
def output_shape(self):
raise NotImplementedError
def show(self, randkey, state, act_func, params, *args, **kwargs):
def show(self, state, randkey, act_func, params, *args, **kwargs):
raise NotImplementedError