complete show() in brax env

problem/func_fit/func_fit.py

@@ -44,7 +44,7 @@ class FuncFit(Problem):
 
         return -loss
 
-    def show(self, randkey, state: State, act_func: Callable, params):
+    def show(self, randkey, state: State, act_func: Callable, params, *args, **kwargs):
         predict = act_func(state, self.inputs, params)
         inputs, target, predict = jax.device_get([self.inputs, self.targets, predict])
         loss = -self.evaluate(randkey, state, act_func, params)

problem/rl_env/brax_env.py

@@ -34,12 +34,50 @@ class BraxEnv(RLEnv):
 
     @property
     def input_shape(self):
-        return (self.env.observation_size, )
+        return (self.env.observation_size,)
 
     @property
     def output_shape(self):
-        return (self.env.action_size, )
+        return (self.env.action_size,)
+
+    def show(self, randkey, state: State, act_func: Callable, params, save_path=None, height=512, width=512,
+             duration=0.1, *args,
+             **kwargs):
+
+        import jax
+        import imageio
+        import numpy as np
+        from brax.io import image
+        from tqdm import tqdm
+
+        obs, env_state = self.reset(randkey)
+        reward, done = 0.0, False
+        state_histories = []
+
+        def step(key, env_state, obs):
+            key, _ = jax.random.split(key)
+            net_out = act_func(state, obs, params)
+            action = self.config.output_transform(net_out)
+            obs, env_state, r, done, _ = self.step(randkey, env_state, action)
+            return key, env_state, obs, r, done
+
+        while not done:
+            state_histories.append(env_state.pipeline_state)
+            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")]
+
+        def create_gif(image_list, gif_name, duration):
+            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)
+
+        create_gif(imgs, save_path, duration=0.1)
+        print("Gif saved to: ", save_path)
+        print("Total reward: ", reward)
+
 
-    def show(self, randkey, state: State, act_func: Callable, params):
-        # TODO
-        raise NotImplementedError("im busy! to de done!")

problem/rl_env/rl_jit.py

@@ -32,7 +32,6 @@ class RLEnv(Problem):
         def body_func(carry):
             obs, env_state, rng, _, tr = carry  # total reward
             net_out = act_func(state, obs, params)
-
             action = self.config.output_transform(net_out)
             next_obs, next_env_state, reward, done, _ = self.step(rng, env_state, action)
             next_rng, _ = jax.random.split(rng)
@@ -68,5 +67,5 @@ class RLEnv(Problem):
     def output_shape(self):
         raise NotImplementedError
 
-    def show(self, randkey, state: State, act_func: Callable, params):
+    def show(self, randkey, state: State, act_func: Callable, params, *args, **kwargs):
         raise NotImplementedError