From 19624aef9d4d19e961759eca782d72566700d1a2 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 19 Feb 2025 15:33:25 -0800 Subject: [PATCH] Add an example of a textured height field to mjSpec notebook. https://youtu.be/i6eXrzJNKeY PiperOrigin-RevId: 728841769 Change-Id: I1aca838af69d8474aab13b06d8c1e5cf87fc6440 --- python/mjspec.ipynb | 271 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 266 insertions(+), 5 deletions(-) diff --git a/python/mjspec.ipynb b/python/mjspec.ipynb index e818800a..2eb89566 100644 --- a/python/mjspec.ipynb +++ b/python/mjspec.ipynb @@ -96,6 +96,7 @@ "\n", "# Other imports and helper functions\n", "import numpy as np\n", + "from scipy.signal import convolve2d\n", "\n", "# Graphics and plotting.\n", "print('Installing mediapy:')\n", @@ -103,6 +104,7 @@ "!pip install -q mediapy\n", "import mediapy as media\n", "import matplotlib.pyplot as plt\n", + "import matplotlib.colors as mcolors\n", "\n", "# Printing.\n", "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n", @@ -208,7 +210,7 @@ "id": "NolxAaRn9N9r" }, "source": [ - "# Constructing models from scratch" + "# Procedural models" ] }, { @@ -261,7 +263,7 @@ "id": "Y4rV2NDh92Ga" }, "source": [ - "## Procedural tree\n", + "## Tree\n", "\n", "Let's use procedural model creation to make a simple model of a tree.\n", "\n", @@ -303,11 +305,12 @@ "cell_type": "code", "execution_count": 0, "metadata": { + "cellView": "form", "id": "IQ9G54Yu-Cse" }, "outputs": [], "source": [ - "#@title utility functions\n", + "#@title Utilities\n", "def branch_frames(num_samples, phi_lower=np.pi / 8, phi_upper=np.pi / 3):\n", " \"\"\"Returns branch direction vectors and normalized attachment heights.\"\"\"\n", " directions = []\n", @@ -504,6 +507,262 @@ "media.show_video(frames, fps=framerate / 2)" ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "0wR6XAnKdB6s" + }, + "source": [ + "## Height field\n", + "\n", + "Height fields represent uneven terrain. There are many ways to generate procedural terrain. Here, we will use [Perlin Noise](https://www.youtube.com/watch?v=9x6NvGkxXhU)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "yXY7HGfVsVlo" + }, + "source": [ + "### Utilities" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "JWlgTJiHevOg" + }, + "outputs": [], + "source": [ + "#@title Perlin noise generator\n", + "\n", + "# adapted from https://github.com/pvigier/perlin-numpy\n", + "\n", + "def interpolant(t):\n", + " return t*t*t*(t*(t*6 - 15) + 10)\n", + "\n", + "def perlin(shape, res, tileable=(False, False), interpolant=interpolant):\n", + " \"\"\"Generate a 2D numpy array of perlin noise.\n", + "\n", + " Args:\n", + " shape: The shape of the generated array (tuple of two ints).\n", + " This must be a multple of res.\n", + " res: The number of periods of noise to generate along each\n", + " axis (tuple of two ints). Note shape must be a multiple of\n", + " res.\n", + " tileable: If the noise should be tileable along each axis\n", + " (tuple of two bools). Defaults to (False, False).\n", + " interpolant: The interpolation function, defaults to\n", + " t*t*t*(t*(t*6 - 15) + 10).\n", + "\n", + " Returns:\n", + " A numpy array of shape shape with the generated noise.\n", + "\n", + " Raises:\n", + " ValueError: If shape is not a multiple of res.\n", + " \"\"\"\n", + " delta = (res[0] / shape[0], res[1] / shape[1])\n", + " d = (shape[0] // res[0], shape[1] // res[1])\n", + " grid = np.mgrid[0:res[0]:delta[0], 0:res[1]:delta[1]]\\\n", + " .transpose(1, 2, 0) % 1\n", + " # Gradients\n", + " angles = 2*np.pi*np.random.rand(res[0]+1, res[1]+1)\n", + " gradients = np.dstack((np.cos(angles), np.sin(angles)))\n", + " if tileable[0]:\n", + " gradients[-1,:] = gradients[0,:]\n", + " if tileable[1]:\n", + " gradients[:,-1] = gradients[:,0]\n", + " gradients = gradients.repeat(d[0], 0).repeat(d[1], 1)\n", + " g00 = gradients[ :-d[0], :-d[1]]\n", + " g10 = gradients[d[0]: , :-d[1]]\n", + " g01 = gradients[ :-d[0],d[1]: ]\n", + " g11 = gradients[d[0]: ,d[1]: ]\n", + " # Ramps\n", + " n00 = np.sum(np.dstack((grid[:,:,0] , grid[:,:,1] )) * g00, 2)\n", + " n10 = np.sum(np.dstack((grid[:,:,0]-1, grid[:,:,1] )) * g10, 2)\n", + " n01 = np.sum(np.dstack((grid[:,:,0] , grid[:,:,1]-1)) * g01, 2)\n", + " n11 = np.sum(np.dstack((grid[:,:,0]-1, grid[:,:,1]-1)) * g11, 2)\n", + " # Interpolation\n", + " t = interpolant(grid)\n", + " n0 = n00*(1-t[:,:,0]) + t[:,:,0]*n10\n", + " n1 = n01*(1-t[:,:,0]) + t[:,:,0]*n11\n", + " return np.sqrt(2)*((1-t[:,:,1])*n0 + t[:,:,1]*n1)\n", + "\n", + "noise = perlin((256, 256), (8, 8))\n", + "plt.imshow(noise, cmap = 'gray', interpolation = 'lanczos')\n", + "plt.title('Perlin noise example')\n", + "plt.colorbar();" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "WKoagEORmCz-" + }, + "outputs": [], + "source": [ + "#@title Soft edge slope\n", + "def edge_slope(size, border_width=5, blur_iterations=20):\n", + " \"\"\"Creates a grayscale image with a white center and fading black edges using convolution.\"\"\"\n", + " img = np.ones((size, size), dtype=np.float32)\n", + " img[:border_width, :] = 0\n", + " img[-border_width:, :] = 0\n", + " img[:, :border_width] = 0\n", + " img[:, -border_width:] = 0\n", + "\n", + " kernel = np.array([[1, 1, 1],\n", + " [1, 1, 1],\n", + " [1, 1, 1]]) / 9.0\n", + "\n", + " for _ in range(blur_iterations):\n", + " img = convolve2d(img, kernel, mode='same', boundary='symm')\n", + "\n", + " return img\n", + "\n", + "image = edge_slope(256)\n", + "plt.imshow(image, cmap='gray')\n", + "plt.title('Smooth sloped edges')\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "MCSks-w3spoJ" + }, + "source": [ + "### Textured height-field generator" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "hzroDjVLfxHs" + }, + "outputs": [], + "source": [ + "def add_hfield(spec=None, hsize=10, vsize=4):\n", + " \"\"\" Function that adds a heighfield with countours\"\"\"\n", + "\n", + " # Initialize spec\n", + " if spec is None:\n", + " spec = mj.MjSpec()\n", + "\n", + " # Generate Perlin noise\n", + " size = 128\n", + " noise = perlin((size, size), (8, 8))\n", + "\n", + " # Remap noise to 0 to 1\n", + " noise = (noise + 1)/2\n", + " noise -= np.min(noise)\n", + " noise /= np.max(noise)\n", + "\n", + " # Makes the edges slope down to avoid sharp boundary\n", + " noise *= edge_slope(size)\n", + "\n", + " # Create height field\n", + " hfield = spec.add_hfield(name ='hfield',\n", + " size = [hsize, hsize, vsize, vsize/10],\n", + " nrow = noise.shape[0],\n", + " ncol = noise.shape[1],\n", + " userdata = noise.flatten())\n", + "\n", + " # Add texture\n", + " texture = spec.add_texture(name = \"contours\",\n", + " type = mj.mjtTexture.mjTEXTURE_2D,\n", + " width = 128, height = 128)\n", + "\n", + " # Create texture map, assign to texture\n", + " h = noise\n", + " s = 0.7 * np.ones(h.shape)\n", + " v = 0.7 * np.ones(h.shape)\n", + " hsv = np.stack([h, s, v], axis=-1)\n", + " rgb = mcolors.hsv_to_rgb(hsv)\n", + " rgb = np.flipud((rgb * 255).astype(np.uint8))\n", + " texture.data = rgb.tobytes()\n", + "\n", + " # Assign texture to material\n", + " grid = spec.add_material( name = 'contours')\n", + " grid.textures[mj.mjtTextureRole.mjTEXROLE_RGB] = 'contours'\n", + " spec.worldbody.add_geom(type = mj.mjtGeom.mjGEOM_HFIELD,\n", + " material = 'contours', hfieldname = 'hfield')\n", + "\n", + " return spec" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "id": "Q2b5BfoZgV79" + }, + "outputs": [], + "source": [ + "#@title Video\n", + "\n", + "arena_xml = \"\"\"\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + "\n", + "\"\"\"\n", + "\n", + "spec = add_hfield(mj.MjSpec.from_string(arena_xml))\n", + "\n", + "# Add lights\n", + "for x in [-15, 15]:\n", + " for y in [-15, 15]:\n", + " spec.worldbody.add_light(pos = [x, y, 10], dir = [-x, -y, -15])\n", + "\n", + "# Add balls\n", + "for x in np.linspace(-8, 8, 8):\n", + " for y in np.linspace(-8, 8, 8):\n", + " pos = [x, y, 4 + np.random.uniform(0, 10)]\n", + " ball = spec.worldbody.add_body(pos=pos)\n", + " ball.add_geom(type = mj.mjtGeom.mjGEOM_SPHERE, size = [0.5, 0, 0],\n", + " rgba = [np.random.uniform()]*3 + [1])\n", + " ball.add_freejoint()\n", + "\n", + "model = spec.compile()\n", + "data = mj.MjData(model)\n", + "\n", + "cam = mj.MjvCamera()\n", + "mj.mjv_defaultCamera(cam)\n", + "cam.lookat = [0, 0, 0]\n", + "cam.distance = 30\n", + "cam.elevation = -30\n", + "\n", + "duration = 6 # (seconds)\n", + "framerate = 60 # (Hz)\n", + "frames = []\n", + "with mj.Renderer(model, width=1920 // 3, height=1080 // 3) as renderer:\n", + " while data.time < duration:\n", + " mj.mj_step(model, data)\n", + " if len(frames) < data.time * framerate:\n", + " cam.azimuth = 20 + 30 * (1 - np.cos(np.pi*data.time / duration))\n", + " renderer.update_scene(data, cam)\n", + " pixels = renderer.render()\n", + " frames.append(pixels)\n", + "\n", + "media.show_video(frames, fps=framerate )" + ] + }, { "cell_type": "markdown", "metadata": { @@ -1073,10 +1332,12 @@ "accelerator": "GPU", "colab": { "collapsed_sections": [ - "sJFuNetilv4m" + "sJFuNetilv4m", + "yXY7HGfVsVlo" ], "gpuClass": "premium", - "private_outputs": true + "private_outputs": true, + "toc_visible": true }, "gpuClass": "premium", "kernelspec": {