add brax env

examples/brax_env.py

@@ -0,0 +1,36 @@
+import jax
+
+import brax
+from brax import envs
+
+
+def inference_func(key, *args):
+    return jax.random.normal(key, shape=(env.action_size,))
+
+
+env_name = "ant"
+backend = "generalized"
+
+env = envs.create(env_name=env_name, backend=backend)
+
+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)
+    act_rng, rng = jax.random.split(rng)
+    act = jit_inference_fn(act_rng, state.obs)
+    state = jit_env_step(state, act)
+    reward = state.reward
+    # print(reward)
+
+a = 1
+
+