complete fully stateful!

use black to format all files!

tensorneat/problem/rl_env/brax_env.py

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

tensorneat/problem/rl_env/gymnax_env.py

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

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