update "show" in brax

tensorneat/examples/brax/show_test.py

@@ -0,0 +1,19 @@
+import jax
+from problem.rl_env import BraxEnv
+
+
+def random_policy(randkey, forward_func, obs):
+    return jax.random.uniform(randkey, (6,), minval=-1, maxval=1)
+
+
+if __name__ == "__main__":
+    problem = BraxEnv(env_name="walker2d", max_step=1000, action_policy=random_policy)
+    state = problem.setup()
+    randkey = jax.random.key(0)
+    problem.show(
+        state,
+        randkey,
+        act_func=lambda state, params, obs: obs,
+        params=None,
+        save_path="walker2d_random_policy",
+    )