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

use black format all files;

Browse Source 
remove "return state" for functions which will be executed in vmap;
recover randkey as args in mutation methods
This commit is contained in:
wls2002
2024-05-26 15:46:04 +08:00
parent 79d53ea7af
commit cf69b916af
38 changed files with 932 additions and 582 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -17,29 +17,29 @@ class RLEnv(BaseProblem):
init_obs, init_env_state = self.reset(rng_reset)
def cond_func(carry):
_, _, _, done, _, count = carry
_, _, _, _, done, _, count = carry
return ~done & (count < self.max_step)
def body_func(carry):
obs, env_state, rng, done, tr, count = carry # tr -> total reward
action = act_func(obs, params)
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 next_obs, next_env_state, next_rng, done, tr + reward, count + 1
return state_, next_obs, next_env_state, next_rng, done, tr + reward, count + 1
_, _, _, _, total_reward, _ = jax.lax.while_loop(
state, _, _, _, _, total_reward, _ = jax.lax.while_loop(
cond_func,
body_func,
(init_obs, init_env_state, rng_episode, False, 0.0, 0)
(state, init_obs, init_env_state, rng_episode, False, 0.0, 0)
)
return total_reward
return state, total_reward
@partial(jax.jit, static_argnums=(0,))
# @partial(jax.jit, static_argnums=(0,))
def step(self, randkey, env_state, action):
return self.env_step(randkey, env_state, action)
@partial(jax.jit, static_argnums=(0,))
# @partial(jax.jit, static_argnums=(0,))
def reset(self, randkey):
return self.env_reset(randkey)