complete show() in brax env

examples/brax_env.py

@@ -1,7 +1,14 @@
+import imageio
 import jax
 
 import brax
 from brax import envs
+from brax.io import image
+import matplotlib.pyplot as plt
+
+import time
+from tqdm import tqdm
+import numpy as np
 
 
 def inference_func(key, *args):
@@ -17,20 +24,50 @@ jit_env_reset = jax.jit(env.reset)
 jit_env_step = jax.jit(env.step)
 jit_inference_fn = jax.jit(inference_func)
 
-
-rollout = []
 rng = jax.random.PRNGKey(seed=1)
 ori_state = jit_env_reset(rng=rng)
 state = ori_state
 
-for _ in range(100):
-    rollout.append(state.pipeline_state)
+render_history = []
+
+for i in range(100):
     act_rng, rng = jax.random.split(rng)
+
+    tic = time.time()
     act = jit_inference_fn(act_rng, state.obs)
     state = jit_env_step(state, act)
+    print("step time: ", time.time() - tic)
+
+    render_history.append(state.pipeline_state)
+
+    # img = image.render_array(sys=env.sys, state=pipeline_state, width=512, height=512)
+    # print("render time: ", time.time() - tic)
+
+    # plt.imsave("../images/ant_{}.png".format(i), img)
+
     reward = state.reward
-    # print(reward)
+    done = state.done
+    print(i, reward)
 
-a = 1
+render_history = jax.device_get(render_history)
+# print(render_history)
+
+imgs = [image.render_array(sys=env.sys, state=s, width=512, height=512) for s in tqdm(render_history)]
 
 
+# for i, s in enumerate(tqdm(render_history)):
+#     img = image.render_array(sys=env.sys, state=s, width=512, height=512)
+#     print(img.shape)
+#     # print(type(img))
+#     plt.imsave("../images/ant_{}.png".format(i), img)
+
+
+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, "../images/ant.gif", 0.1)