chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

Update tutorial notebooks to use the built-in Renderer class.

Browse Source 
PiperOrigin-RevId: 493303419
Change-Id: I5ddca087d522e7d43f8880641e8bc8d82d4f07a3
This commit is contained in:
Saran Tunyasuvunakool
2022-12-06 07:45:25 -08:00
committed by Copybara-Service
parent 0d608282dd
commit 8f3a771209
2 changed files with 43 additions and 427 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/LQR.ipynb
+20 -161
View File
@@ -38,21 +38,33 @@
{
"cell_type": "markdown",
"metadata": {
"id": "pXd9W7hMHYWx"
"id": "QPdJNe3k62mx"
},
"source": [
"## Setup"
"### Install MuJoCo\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Xqo7pyX-n72M"
},
"outputs": [],
"source": [
"!pip install mujoco"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "IbZxYDxzoz5R"
},
"outputs": [],
"source": [
"#@title Install MuJoCo\n",
"#@title Check if installation was successful\n",
"\n",
"from google.colab import files\n",
"\n",
@@ -62,10 +74,7 @@
" raise RuntimeError(\n",
" 'Cannot communicate with GPU. '\n",
" 'Make sure you are using a GPU Colab runtime. '\n",
" 'Go to the Runtime menu and select Choose runtime type.')\n",
"\n",
"print('Installing mujoco:')\n",
"!pip install -q mujoco\n",
" 'Go to the Runtime menu and select Choose runtime type.')",
"\n",
"# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
"print('Setting environment variable to use GPU rendering:')\n",
@@ -89,7 +98,8 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "gKc1FNhKiVJX"
"cellView": "form",
"id": "T5f4w3Kq2X14"
},
"outputs": [],
"source": [
@@ -109,157 +119,6 @@
"np.set_printoptions(precision=3, suppress=True, linewidth=100)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "zcXO67XM5-Kc"
},
"source": [
"## Renderer class\n",
"\n",
"We define a simple `Renderer` to render MuJoCo scenes. It is similar to the `Camera` class provided by `dm_control`.\n",
"\n",
"This class is still under development, please help us improve it."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "l5EFG9Mi6Joy"
},
"outputs": [],
"source": [
"class Renderer:\n",
" \"\"\"Renders MuJoCo scenes.\"\"\"\n",
"\n",
" def __init__(\n",
" self,\n",
" model: mujoco.MjModel,\n",
" height: int = 240,\n",
" width: int = 320,\n",
" max_geom: int = 5000,\n",
" ) -\u003e None:\n",
" \"\"\"Initializes a new `Renderer`.\n",
"\n",
" Args:\n",
" model: an MjModel instance.\n",
" height: image height in pixels.\n",
" width: image width in pixels.\n",
" max_geom: integer specifying the maximum number of geoms that can be\n",
" rendered in the same scene.\n",
"\n",
" Raises:\n",
" ValueError: If `camera_id` is outside the valid range, or if `width` or\n",
" `height` exceed the dimensions of MuJoCo's offscreen framebuffer.\n",
" \"\"\"\n",
" buffer_width = model.vis.global_.offwidth\n",
" buffer_height = model.vis.global_.offheight\n",
" if width \u003e buffer_width:\n",
" raise ValueError('Image width {} \u003e framebuffer width {}. Either reduce '\n",
" 'the image width or specify a larger offscreen '\n",
" 'framebuffer in the model XML using the clause\\n'\n",
" '\u003cvisual\u003e\\n'\n",
" ' \u003cglobal offwidth=\"my_width\"/\u003e\\n'\n",
" '\u003c/visual\u003e'.format(width, buffer_width))\n",
" if height \u003e buffer_height:\n",
" raise ValueError('Image height {} \u003e framebuffer height {}. Either reduce '\n",
" 'the image height or specify a larger offscreen '\n",
" 'framebuffer in the model XML using the clause\\n'\n",
" '\u003cvisual\u003e\\n'\n",
" ' \u003cglobal offheight=\"my_height\"/\u003e\\n'\n",
" '\u003c/visual\u003e'.format(height, buffer_height))\n",
"\n",
" self._width = width\n",
" self._height = height\n",
" self._model = model\n",
"\n",
" self._scene = mujoco.MjvScene(model=model, maxgeom=max_geom)\n",
" self._scene_option = mujoco.MjvOption()\n",
"\n",
" self._rect = mujoco.MjrRect(0, 0, self._width, self._height)\n",
"\n",
" # Internal buffers.\n",
" self._rgb_buffer = np.empty((self._height, self._width, 3), dtype=np.uint8)\n",
" self._depth_buffer = np.empty((self._height, self._width), dtype=np.float32)\n",
"\n",
" # Create render contexts.\n",
" self._gl_context = mujoco.GLContext(self._width, self._height)\n",
" self._gl_context.make_current()\n",
" self._mjr_context = mujoco.MjrContext(\n",
" model, mujoco.mjtFontScale.mjFONTSCALE_150\n",
" )\n",
" mujoco.mjr_setBuffer(\n",
" mujoco.mjtFramebuffer.mjFB_OFFSCREEN, self._mjr_context\n",
" )\n",
"\n",
" def render(self) -\u003e np.ndarray:\n",
" \"\"\"Renders the scene as a numpy array of pixel values.\n",
"\n",
" Returns:\n",
" A numpy array of pixels with dimensions (H, W, 3). The array will be\n",
" mutated by future calls to `render`.\n",
" \"\"\"\n",
" self._gl_context.make_current()\n",
"\n",
" # Render scene and read contents of RGB buffer.\n",
" mujoco.mjr_render(self._rect, self._scene, self._mjr_context)\n",
" mujoco.mjr_readPixels(self._rgb_buffer, None, self._rect, self._mjr_context)\n",
"\n",
" pixels = self._rgb_buffer\n",
" return np.flipud(pixels)\n",
"\n",
" def update_scene(\n",
" self,\n",
" data: mujoco.MjData,\n",
" camera: Union[int, str, mujoco.MjvCamera] = -1,\n",
" scene_option: Optional[mujoco.MjvOption] = None,\n",
" ):\n",
" \"\"\"Updates geometry used for rendering.\n",
"\n",
" Args:\n",
" data: An instance of `mujoco.MjData`.\n",
" camera: An instance of `mujoco.MjvCamera`, a string or an integer\n",
" scene_option: A custom `mujoco.MjvOption` instance to use to render\n",
" the scene instead of the default.\n",
" \"\"\"\n",
" if not isinstance(camera, mujoco.MjvCamera):\n",
" camera_id = camera\n",
" if isinstance(camera_id, str):\n",
" camera_id = self._model.camera(camera_id).id\n",
" if camera_id \u003c -1:\n",
" raise ValueError('camera_id cannot be smaller than -1.')\n",
" if camera_id \u003e= self._model.ncam:\n",
" raise ValueError(\n",
" f'model has {self._model.ncam} fixed cameras. '\n",
" f'camera_id={camera_id} is invalid.'\n",
" )\n",
" camera = mujoco.MjvCamera()\n",
" camera.fixedcamid = camera_id\n",
"\n",
" # -1 corresponds to free camera.\n",
" if camera_id == -1:\n",
" camera.type = mujoco.mjtCamera.mjCAMERA_FREE\n",
" mujoco.mjv_defaultFreeCamera(self._model, camera)\n",
" # Else index into the corresponding fixed camera.\n",
" else:\n",
" camera.type = mujoco.mjtCamera.mjCAMERA_FIXED\n",
"\n",
" scene_option = scene_option or self._scene_option\n",
" mujoco.mjv_updateScene(\n",
" self._model,\n",
" data,\n",
" scene_option,\n",
" None,\n",
" camera, mujoco.mjtCatBit.mjCAT_ALL,\n",
" self._scene,\n",
" )\n",
"\n",
" @property\n",
" def scene(self) -\u003e mujoco.MjvScene:\n",
" return self._scene"
]
},
{
"cell_type": "markdown",
"metadata": {
@@ -302,7 +161,7 @@
"source": [
"model = mujoco.MjModel.from_xml_string(xml)\n",
"data = mujoco.MjData(model)\n",
"renderer = Renderer(model)"
"renderer = mujoco.Renderer(model)"
]
},
{
@@ -1038,7 +897,7 @@
"data.qpos = qpos0\n",
"\n",
"# New renderer instance with higher resolution.\n",
"renderer = Renderer(model, width=1280, height=720)\n",
"renderer = mujoco.Renderer(model, width=1280, height=720)\n",
"\n",
"frames = []\n",
"step = 0\n",