37 lines
694 B
Python
37 lines
694 B
Python
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
|
|
|
|
|