| 75 | |
| 76 | |
| 77 | def main(): |
| 78 | env = PushT() |
| 79 | rng = jax.random.PRNGKey(1) |
| 80 | env_step = jax.jit(env.step) |
| 81 | env_reset = jax.jit(env.reset) |
| 82 | state = env_reset(rng) |
| 83 | rollout = [state.pipeline_state] |
| 84 | for _ in range(50): |
| 85 | rng, rng_act = jax.random.split(rng) |
| 86 | act = jax.random.uniform(rng_act, (env.action_size,), minval=-1.0, maxval=1.0) |
| 87 | state = env_step(state, act) |
| 88 | rollout.append(state.pipeline_state) |
| 89 | webpage = html.render(env.sys.replace(dt=env.dt), rollout) |
| 90 | path = f"{mbd.__path__[0]}/../results/pushT" |
| 91 | if not os.path.exists(path): |
| 92 | os.makedirs(path) |
| 93 | with open(f"{path}/vis.html", "w") as f: |
| 94 | f.write(webpage) |
| 95 | |
| 96 | |
| 97 | if __name__ == "__main__": |