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

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
+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",
+23 -266
View File
@@ -49,39 +49,7 @@
"id": "YvyGCsgSCxHQ"
},
"source": [
"# Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "cVAwWp0E6Zgq"
},
"outputs": [],
"source": [
"#@title Prerequisite checks for the Colab\n",
"from google.colab import files\n",
"\n",
"import distutils.util\n",
"import subprocess\n",
"if subprocess.run('nvidia-smi').returncode:\n",
" 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.')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QPdJNe3k62mx"
},
"source": [
"## Installing MuJoCo\n",
"\n",
"Here we install the MuJoCo package. We also check if the installation was successful. This part is optional."
"# Install MuJoCo"
]
},
{
@@ -106,6 +74,16 @@
"source": [
"#@title Check if installation was successful\n",
"\n",
"from google.colab import files\n",
"\n",
"import distutils.util\n",
"import subprocess\n",
"if subprocess.run('nvidia-smi').returncode:\n",
" 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",
"# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
"print('Setting environment variable to use GPU rendering:')\n",
"%env MUJOCO_GL=egl\n",
@@ -150,227 +128,6 @@
"np.set_printoptions(precision=3, suppress=True, linewidth=100)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-re3Szx-1Ias"
},
"source": [
"## The 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",
"In a future version of the MuJoCo Python bindings, this class will become a native library function. We include it in this colab as a draft, open to comments and bug reports."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "liZGoxlEOlMK"
},
"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",
" ) -\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",
"\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=10000)\n",
" self._scene_option = mujoco.MjvOption()\n",
"\n",
" self._rect = mujoco.MjrRect(0, 0, self._width, self._height)\n",
"\n",
" # Render camera.\n",
" self._render_camera = mujoco.MjvCamera()\n",
" self._render_camera.fixedcamid = -1\n",
" self._render_camera.type = mujoco.mjtCamera.mjCAMERA_FREE\n",
" mujoco.mjv_defaultFreeCamera(model, self._render_camera)\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",
" # Default render flags\n",
" self._render_flags = {'depth': False, 'segmentation': False}\n",
"\n",
" @property\n",
" def model(self):\n",
" return self._model\n",
"\n",
" @property\n",
" def scene(self) -\u003e mujoco.MjvScene:\n",
" return self._scene\n",
"\n",
" @property\n",
" def height(self):\n",
" return self._height\n",
"\n",
" @property\n",
" def width(self):\n",
" return self._width\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",
" original_flags = self._scene.flags.copy()\n",
"\n",
" if self._render_flags['segmentation']:\n",
" self._scene.flags[mujoco.mjtRndFlag.mjRND_SEGMENT] = True\n",
" self._scene.flags[mujoco.mjtRndFlag.mjRND_IDCOLOR] = True\n",
"\n",
" self._gl_context.make_current()\n",
"\n",
" # Render scene and read contents of RGB and depth buffers.\n",
" mujoco.mjr_render(self._rect, self._scene, self._mjr_context)\n",
" mujoco.mjr_readPixels(self._rgb_buffer, self._depth_buffer, self._rect,\n",
" self._mjr_context)\n",
"\n",
" if self._render_flags['depth']:\n",
" # Get the distances to the near and far clipping planes.\n",
" extent = self._model.stat.extent\n",
" near = self._model.vis.map.znear * extent\n",
" far = self._model.vis.map.zfar * extent\n",
"\n",
" # Convert from [0 1] to depth in meters, see links below:\n",
" # http://stackoverflow.com/a/6657284/1461210\n",
" # https://www.khronos.org/opengl/wiki/Depth_Buffer_Precision\n",
" pixels = near / (1 - self._depth_buffer * (1 - near / far))\n",
"\n",
" elif self._render_flags['segmentation']:\n",
" # Convert 3-channel uint8 to 1-channel uint32.\n",
" image3 = self._rgb_buffer.astype(np.uint32)\n",
" segimage = (image3[:, :, 0] +\n",
" image3[:, :, 1] * (2**8) +\n",
" image3[:, :, 2] * (2**16))\n",
" # Remap segid to 2-channel (object ID, object type) pair.\n",
" # Seg ID 0 is background -- will be remapped to (-1, -1).\n",
" ngeoms = self._scene.ngeom\n",
" segid2output = np.full((ngeoms + 1, 2), fill_value=-1,\n",
" dtype=np.int32) # Seg id cannot be \u003e ngeom + 1.\n",
" visible_geoms = [g for g in self._scene.geoms[:ngeoms] if g.segid != -1]\n",
" visible_segids = np.array([g.segid + 1 for g in visible_geoms], np.int32)\n",
" visible_objid = np.array([g.objid for g in visible_geoms], np.int32)\n",
" visible_objtype = np.array([g.objtype for g in visible_geoms], np.int32)\n",
" segid2output[visible_segids, 0] = visible_objid\n",
" segid2output[visible_segids, 1] = visible_objtype\n",
" pixels = segid2output[segimage]\n",
"\n",
" # Reset scene flags.\n",
" np.copyto(self._scene.flags, original_flags)\n",
" else:\n",
" pixels = self._rgb_buffer\n",
" return np.flipud(pixels)\n",
"\n",
" def enable_depth_rendering(self):\n",
" self._render_flags = {'depth': True, 'segmentation': False}\n",
"\n",
" def disable_depth_rendering(self):\n",
" self._render_flags = {'depth': False, 'segmentation': False}\n",
"\n",
" def enable_segmentation_rendering(self):\n",
" self._render_flags = {'depth': False, 'segmentation': True}\n",
"\n",
" def disable_segmentation_rendering(self):\n",
" self._render_flags = {'depth': False, 'segmentation': False}\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 = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, camera_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",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {
@@ -681,7 +438,7 @@
"data = mujoco.MjData(model)\n",
"\n",
"# Make renderer, render and show the pixels\n",
"renderer = Renderer(model)\n",
"renderer = mujoco.Renderer(model)\n",
"media.show_image(renderer.render())"
]
},
@@ -740,7 +497,7 @@
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(xml)\n",
"data = mujoco.MjData(model)\n",
"renderer = Renderer(model)\n",
"renderer = mujoco.Renderer(model)\n",
"\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data)\n",
@@ -842,7 +599,7 @@
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(xml)\n",
"data = mujoco.MjData(model)\n",
"renderer = Renderer(model)\n",
"renderer = mujoco.Renderer(model)\n",
"\n",
"# enable joint visualization option:\n",
"scene_option = mujoco.MjvOption()\n",
@@ -1017,7 +774,7 @@
"\u003c/mujoco\u003e\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(tippe_top)\n",
"renderer = Renderer(model)\n",
"renderer = mujoco.Renderer(model)\n",
"data = mujoco.MjData(model)\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data, camera=\"closeup\")\n",
@@ -1194,7 +951,7 @@
"\u003c/mujoco\u003e\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(chaotic_pendulum)\n",
"renderer = Renderer(model, 480, 640)\n",
"renderer = mujoco.Renderer(model, 480, 640)\n",
"data = mujoco.MjData(model)\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data, camera=\"fixed\")\n",
@@ -1224,7 +981,7 @@
"framerate = 30 # Hz\n",
"n_frames = int(n_seconds * framerate)\n",
"frames = []\n",
"renderer = Renderer(model, 240, 320)\n",
"renderer = mujoco.Renderer(model, 240, 320)\n",
"\n",
"\n",
"# set initial state\n",
@@ -1489,7 +1246,7 @@
"\u003c/mujoco\u003e\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(free_body_MJCF)\n",
"renderer = Renderer(model, 400, 600)\n",
"renderer = mujoco.Renderer(model, 400, 600)\n",
"data = mujoco.MjData(model)\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data, \"fixed\")\n",
@@ -1517,7 +1274,7 @@
"height = 240\n",
"width = 320\n",
"frames = []\n",
"renderer = Renderer(model, height, width)\n",
"renderer = mujoco.Renderer(model, height, width)\n",
"\n",
"# visualize contact frames and forces, make body transparent\n",
"options = mujoco.MjvOption()\n",
@@ -1696,7 +1453,7 @@
"# load\n",
"model = mujoco.MjModel.from_xml_string(MJCF)\n",
"data = mujoco.MjData(model)\n",
"renderer = Renderer(model, height, width)\n",
"renderer = mujoco.Renderer(model, height, width)\n",
"\n",
"# simulate and render\n",
"mujoco.mj_resetData(model, data)\n",
@@ -1774,7 +1531,7 @@
"\u003c/mujoco\u003e\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(MJCF)\n",
"renderer = Renderer(model, 480, 480)\n",
"renderer = mujoco.Renderer(model, 480, 480)\n",
"data = mujoco.MjData(model)\n",
"mujoco.mj_forward(model, data)\n",
"renderer.update_scene(data, \"fixed\")\n",
@@ -1806,7 +1563,7 @@
"times = []\n",
"sensordata = []\n",
"\n",
"renderer = Renderer(model, height, width)\n",
"renderer = mujoco.Renderer(model, height, width)\n",
"\n",
"# constant actuator signal\n",
"mujoco.mj_resetData(model, data)\n",
@@ -1897,7 +1654,7 @@
"\u003c/mujoco\u003e\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(xml)\n",
"renderer = Renderer(model)\n",
"renderer = mujoco.Renderer(model)\n",
"data = mujoco.MjData(model)\n",
"\n",
"mujoco.mj_forward(model, data)\n",