Remove block in colab.

PiperOrigin-RevId: 574595549
Change-Id: I3814a8ecb909a581d96b614edd99002c05de014b
This commit is contained in:
Baruch Tabanpour
2023-10-18 14:18:12 -07:00
committed by Copybara-Service
parent 48ec2d465b
commit b989a85364
-42
View File
@@ -1208,48 +1208,6 @@
"\n",
"media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
]
},
{
"metadata": {
"id": "lLXVB0Q3EaC-"
},
"cell_type": "markdown",
"source": [
"We can also perform the physics step using the original MuJoCo python bindings to show that a policy trained in MJX works in MuJoCo-C."
]
},
{
"metadata": {
"id": "dXHTu5AIEaC-"
},
"cell_type": "code",
"source": [
"mj_model = eval_env.model\n",
"mj_data = mujoco.MjData(mj_model)\n",
"\n",
"camera = 'track' if env_name == 'barkour' else 'side'\n",
"renderer = mujoco.Renderer(mj_model)\n",
"\n",
"ctrl = jp.zeros(mj_model.nu)\n",
"\n",
"images = []\n",
"for i in range(n_steps):\n",
" act_rng, rng = jax.random.split(rng)\n",
" obs = eval_env._get_obs(mjx.device_put(mj_data), ctrl)\n",
" ctrl, _ = jit_inference_fn(obs, act_rng)\n",
"\n",
" mj_data.ctrl = ctrl\n",
" for _ in range(eval_env._physics_steps_per_control_step):\n",
" mujoco.mj_step(mj_model, mj_data) # Physics step using MuJoCo mj_step.\n",
"\n",
" if i % render_every == 0:\n",
" renderer.update_scene(mj_data, camera=camera)\n",
" images.append(renderer.render())\n",
"\n",
"media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
],
"outputs": [],
"execution_count": null
}
],
"metadata": {