initialize methods

tensorneat/problem/rl_env/rl_jit.py

@@ -9,32 +9,55 @@ class RLEnv(BaseProblem):
     jitable = True
 
     # TODO: move output transform to algorithm
-    def __init__(self):
+    def __init__(self, max_step=1000):
         super().__init__()
+        self.max_step = max_step
+
+    # def evaluate(self, randkey, state, 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, _ = carry
+    #         return ~done
+
+    #     def body_func(carry):
+    #         obs, env_state, rng, _, tr = carry  # total reward
+    #         action = act_func(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
+
+    #     _, _, _, _, total_reward = jax.lax.while_loop(
+    #         cond_func,
+    #         body_func,
+    #         (init_obs, init_env_state, rng_episode, False, 0.0)
+    #     )
+
+    #     return total_reward
 
     def evaluate(self, randkey, state, 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, _ = carry
-            return ~done
+            _, _, _, done, _, count = carry
+            return ~done & (count < self.max_step)  
 
         def body_func(carry):
-            obs, env_state, rng, _, tr = carry  # total reward
-            action = act_func(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
+            obs, env_state, rng, done, tr, count = carry  # tr -> total reward
+            action = act_func(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
 
-        _, _, _, _, total_reward = jax.lax.while_loop(
+        _, _, _, _, total_reward, _ = jax.lax.while_loop(
             cond_func,
             body_func,
-            (init_obs, init_env_state, rng_episode, False, 0.0)
+            (init_obs, init_env_state, rng_episode, False, 0.0, 0)  
         )
 
-        return 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)