Skip to content

Commit

Permalink
Add an example of a textured height field to mjSpec notebook.
Browse files Browse the repository at this point in the history 
https://youtu.be/i6eXrzJNKeY

PiperOrigin-RevId: 728841769
Change-Id: I1aca838af69d8474aab13b06d8c1e5cf87fc6440
  • Loading branch information
@yuvaltassa @copybara-github
yuvaltassa authored and copybara-github committed Feb 19, 2025
1 parent 7a2ad8f commit 19624ae
Showing 1 changed file with 266 additions and 5 deletions.
271 changes: 266 additions & 5 deletions python/mjspec.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,13 +96,15 @@
"\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",
"!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
"!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",
Expand Down Expand Up @@ -208,7 +210,7 @@
"id": "NolxAaRn9N9r"
},
"source": [
"# Constructing models from scratch"
"# Procedural models"
]
},
{
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
"<mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global offwidth=\"2048\" offheight=\"1536\"/>\n",
" <quality shadowsize=\"8192\"/>\n",
" </visual>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" width=\"10\" height=\"10\"/>\n",
" <texture type=\"2d\" name=\"groundplane\" builtin=\"checker\" mark=\"edge\" rgb1=\"1 1 1\" rgb2=\"1 1 1\" markrgb=\"0 0 0\" width=\"400\" height=\"400\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texrepeat=\"45 45\" reflectance=\"0\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <geom name=\"floor\" size=\"150 150 0.1\" type=\"plane\" material=\"groundplane\"/>\n",
" </worldbody>\n",
"</mujoco>\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": {
Expand Down Expand Up @@ -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": {
Expand Down

0 comments on commit 19624ae

Please sign in to comment.