use black format all files;

remove "return state" for functions which will be executed in vmap;
recover randkey as args in mutation methods

tensorneat/problem/base.py

@@ -6,9 +6,9 @@ from utils import State
 class BaseProblem:
     jitable = None
 
-    def setup(self, randkey, state: State = State()):
+    def setup(self, state: State = State()):
         """initialize the state of the problem"""
-        pass
+        return state
 
     def evaluate(self, randkey, state: State, act_func: Callable, params):
         """evaluate one individual"""

tensorneat/problem/func_fit/func_fit.py

@@ -16,12 +16,12 @@ class FuncFit(BaseProblem):
         assert error_method in {'mse', 'rmse', 'mae', 'mape'}
         self.error_method = error_method
 
-    def setup(self, randkey, state: State = State()):
+    def setup(self, state: State = State()):
         return state
 
     def evaluate(self, randkey, state, act_func, params):
 
-        predict = jax.vmap(act_func, in_axes=(0, None))(self.inputs, params)
+        state, predict = jax.vmap(act_func, in_axes=(None, 0, None), out_axes=(None, 0))(state, self.inputs, params)
 
         if self.error_method == 'mse':
             loss = jnp.mean((predict - self.targets) ** 2)
@@ -38,12 +38,14 @@ class FuncFit(BaseProblem):
         else:
             raise NotImplementedError
 
-        return -loss
+        return state, -loss
 
     def show(self, randkey, state, act_func, params, *args, **kwargs):
-        predict = jax.vmap(act_func, in_axes=(0, None))(self.inputs, params)
+        state, predict = jax.vmap(act_func, in_axes=(None, 0, None), out_axes=(None, 0))(state, self.inputs, params)
         inputs, target, predict = jax.device_get([self.inputs, self.targets, predict])
-        loss = -self.evaluate(randkey, state, act_func, params)
+        state, loss = self.evaluate(randkey, state, act_func, params)
+        loss = -loss
+
         msg = ""
         for i in range(inputs.shape[0]):
             msg += f"input: {inputs[i]}, target: {target[i]}, predict: {predict[i]}\n"

tensorneat/problem/rl_env/rl_jit.py

@@ -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)