Remove block in colab.
PiperOrigin-RevId: 574595549 Change-Id: I3814a8ecb909a581d96b614edd99002c05de014b
This commit is contained in:
committed by
Copybara-Service
parent
48ec2d465b
commit
b989a85364
@@ -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": {
|
||||
|
||||
Reference in New Issue
Block a user