Skip to content

Commit

Permalink
[DROP] Add examples for parallel visualization
Browse files Browse the repository at this point in the history
  • Loading branch information
flferretti committed Jan 28, 2025
1 parent 9958e80 commit fd19ecd
Showing 1 changed file with 76 additions and 2 deletions.
78 changes: 76 additions & 2 deletions examples/jaxsim_as_physics_engine_advanced.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,17 @@
" )\n",
")(jnp.vstack(subkeys))\n",
"\n",
"# Reset the x and y position to a grid.\n",
"data_batch_t0 = data_batch_t0.reset_base_position(\n",
" jnp.array(\n",
" [\n",
" jnp.linspace(-1, 1, batch_size),\n",
" jnp.linspace(-1, 1, batch_size),\n",
" data_batch_t0.base_position()[:, 2],\n",
" ]\n",
" ).T\n",
")\n",
"\n",
"print(\"W_p_B(t0)=\\n\", data_batch_t0.base_position()[0:10])"
]
},
Expand Down Expand Up @@ -420,6 +431,69 @@
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import jaxsim.mujoco\n",
"\n",
"mjcf_string, assets = jaxsim.mujoco.ModelToMjcf.convert(model.built_from)\n",
"\n",
"# Create a helper for each parallel instance.\n",
"mj_model_helpers = [\n",
" jaxsim.mujoco.MujocoModelHelper.build_from_xml(\n",
" mjcf_description=mjcf_string, assets=assets\n",
" )\n",
" for _ in range(batch_size)\n",
"]\n",
"\n",
"# Create the video recorder.\n",
"recorder = jaxsim.mujoco.MujocoVideoRecorder(\n",
" model=mj_model_helpers[0].model,\n",
" data=[helper.data for helper in mj_model_helpers],\n",
" fps=int(1 / model.time_step),\n",
" width=320 * 2,\n",
" height=240 * 2,\n",
")\n",
"\n",
"for data_t in data_trajectory_list:\n",
"\n",
" for helper, base_position, base_quaternion, joint_position in zip(\n",
" mj_model_helpers,\n",
" data_t.base_position(),\n",
" data.base_orientation(True),\n",
" data.joint_positions(),\n",
" strict=True,\n",
" ):\n",
" helper.set_base_position(position=base_position)\n",
" helper.set_base_orientation(orientation=base_quaternion, dcm=True)\n",
"\n",
" if model.dofs() > 0:\n",
" helper.set_joint_positions(\n",
" positions=joint_position, joint_names=model.joint_names()\n",
" )\n",
"\n",
" # Record a new video frame.\n",
" recorder.record_frame()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import mediapy as media\n",
"\n",
"# from pathlib import Path\n",
"\n",
"# recorder.write_video(Path(\"sphere.mp4\"))\n",
"\n",
"media.show_video(recorder.frames, fps=recorder.fps)"
]
},
{
"cell_type": "markdown",
"metadata": {
Expand Down Expand Up @@ -454,7 +528,7 @@
"toc_visible": true
},
"kernelspec": {
"display_name": "rsl",
"display_name": "jaxsimsprint",
"language": "python",
"name": "python3"
},
Expand All @@ -468,7 +542,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.7"
"version": "3.13.1"
}
},
"nbformat": 4,
Expand Down

0 comments on commit fd19ecd

Please sign in to comment.