complete fully stateful!

use black to format all files!

tensorneat/problem/rl_env/rl_jit.py

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