complete show() in brax env

examples/brax/ant.py

@@ -12,7 +12,7 @@ def example_conf():
         basic=BasicConfig(
             seed=42,
             fitness_target=10000,
-            pop_size=100
+            pop_size=1000
         ),
         neat=NeatConfig(
             inputs=27,
@@ -23,6 +23,7 @@ def example_conf():
             activation_options=(Act.tanh,),
         ),
         problem=BraxConfig(
+            env_name="ant"
         )
     )
 

examples/brax/half_cheetah.py

@@ -15,7 +15,8 @@ def example_conf():
         basic=BasicConfig(
             seed=42,
             fitness_target=10000,
-            pop_size=10000
+            generation_limit=10,
+            pop_size=100
         ),
         neat=NeatConfig(
             inputs=17,
@@ -33,9 +34,9 @@ def example_conf():
 
 if __name__ == '__main__':
     conf = example_conf()
-
     algorithm = NEAT(conf, NormalGene)
     pipeline = Pipeline(conf, algorithm, BraxEnv)
     state = pipeline.setup()
     pipeline.pre_compile(state)
     state, best = pipeline.auto_run(state)
+    pipeline.show(state, best, save_path="half_cheetah.gif", )

examples/brax/reacher.py

@@ -12,7 +12,7 @@ def example_conf():
         basic=BasicConfig(
             seed=42,
             fitness_target=10000,
-            pop_size=10000
+            pop_size=1000
         ),
         neat=NeatConfig(
             inputs=11,

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)

examples/brax_render.py

@@ -0,0 +1,54 @@
+import brax
+from brax import envs
+from brax.envs.wrappers import gym as gym_wrapper
+from brax.io import image
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import traceback
+
+# print(f"Using Brax {brax.__version__}, Jax {jax.__version__}")
+# print("From GymWrapper, env.reset()")
+# try:
+#     env = envs.create("inverted_pendulum",
+#                       batch_size=1,
+#                       episode_length=150,
+#                       backend='generalized')
+#     env = gym_wrapper.GymWrapper(env)
+#     env.reset()
+#     img = env.render(mode='rgb_array')
+#     plt.imshow(img)
+# except Exception:
+#     traceback.print_exc()
+#
+# print("From GymWrapper, env.reset() and action")
+# try:
+#     env = envs.create("inverted_pendulum",
+#                       batch_size=1,
+#                       episode_length=150,
+#                       backend='generalized')
+#     env = gym_wrapper.GymWrapper(env)
+#     env.reset()
+#     action = jnp.zeros(env.action_space.shape)
+#     env.step(action)
+#     img = env.render(mode='rgb_array')
+#     plt.imshow(img)
+# except Exception:
+#     traceback.print_exc()
+
+print("From brax env")
+try:
+    env = envs.create("inverted_pendulum",
+                      batch_size=1,
+                      episode_length=150,
+                      backend='generalized')
+    key = jax.random.PRNGKey(0)
+    initial_env_state = env.reset(key)
+    base_state = initial_env_state.pipeline_state
+    pipeline_state = env.pipeline_init(base_state.q.ravel(), base_state.qd.ravel())
+    img = image.render_array(sys=env.sys, state=pipeline_state, width=256, height=256)
+    print(f"pixel values: [{img.min()}, {img.max()}]")
+    plt.imshow(img)
+    plt.show()
+except Exception:
+    traceback.print_exc()