Add an example of a textured height field to mjSpec notebook.
https://youtu.be/i6eXrzJNKeY PiperOrigin-RevId: 728841769 Change-Id: I1aca838af69d8474aab13b06d8c1e5cf87fc6440
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
7a2ad8fd97
commit
19624aef9d
+266
-5
@@ -96,6 +96,7 @@
|
|||||||
"\n",
|
"\n",
|
||||||
"# Other imports and helper functions\n",
|
"# Other imports and helper functions\n",
|
||||||
"import numpy as np\n",
|
"import numpy as np\n",
|
||||||
|
"from scipy.signal import convolve2d\n",
|
||||||
"\n",
|
"\n",
|
||||||
"# Graphics and plotting.\n",
|
"# Graphics and plotting.\n",
|
||||||
"print('Installing mediapy:')\n",
|
"print('Installing mediapy:')\n",
|
||||||
@@ -103,6 +104,7 @@
|
|||||||
"!pip install -q mediapy\n",
|
"!pip install -q mediapy\n",
|
||||||
"import mediapy as media\n",
|
"import mediapy as media\n",
|
||||||
"import matplotlib.pyplot as plt\n",
|
"import matplotlib.pyplot as plt\n",
|
||||||
|
"import matplotlib.colors as mcolors\n",
|
||||||
"\n",
|
"\n",
|
||||||
"# Printing.\n",
|
"# Printing.\n",
|
||||||
"np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
|
"np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
|
||||||
@@ -208,7 +210,7 @@
|
|||||||
"id": "NolxAaRn9N9r"
|
"id": "NolxAaRn9N9r"
|
||||||
},
|
},
|
||||||
"source": [
|
"source": [
|
||||||
"# Constructing models from scratch"
|
"# Procedural models"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -261,7 +263,7 @@
|
|||||||
"id": "Y4rV2NDh92Ga"
|
"id": "Y4rV2NDh92Ga"
|
||||||
},
|
},
|
||||||
"source": [
|
"source": [
|
||||||
"## Procedural tree\n",
|
"## Tree\n",
|
||||||
"\n",
|
"\n",
|
||||||
"Let's use procedural model creation to make a simple model of a tree.\n",
|
"Let's use procedural model creation to make a simple model of a tree.\n",
|
||||||
"\n",
|
"\n",
|
||||||
@@ -303,11 +305,12 @@
|
|||||||
"cell_type": "code",
|
"cell_type": "code",
|
||||||
"execution_count": 0,
|
"execution_count": 0,
|
||||||
"metadata": {
|
"metadata": {
|
||||||
|
"cellView": "form",
|
||||||
"id": "IQ9G54Yu-Cse"
|
"id": "IQ9G54Yu-Cse"
|
||||||
},
|
},
|
||||||
"outputs": [],
|
"outputs": [],
|
||||||
"source": [
|
"source": [
|
||||||
"#@title utility functions\n",
|
"#@title Utilities\n",
|
||||||
"def branch_frames(num_samples, phi_lower=np.pi / 8, phi_upper=np.pi / 3):\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",
|
" \"\"\"Returns branch direction vectors and normalized attachment heights.\"\"\"\n",
|
||||||
" directions = []\n",
|
" directions = []\n",
|
||||||
@@ -504,6 +507,262 @@
|
|||||||
"media.show_video(frames, fps=framerate / 2)"
|
"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",
|
"cell_type": "markdown",
|
||||||
"metadata": {
|
"metadata": {
|
||||||
@@ -1073,10 +1332,12 @@
|
|||||||
"accelerator": "GPU",
|
"accelerator": "GPU",
|
||||||
"colab": {
|
"colab": {
|
||||||
"collapsed_sections": [
|
"collapsed_sections": [
|
||||||
"sJFuNetilv4m"
|
"sJFuNetilv4m",
|
||||||
|
"yXY7HGfVsVlo"
|
||||||
],
|
],
|
||||||
"gpuClass": "premium",
|
"gpuClass": "premium",
|
||||||
"private_outputs": true
|
"private_outputs": true,
|
||||||
|
"toc_visible": true
|
||||||
},
|
},
|
||||||
"gpuClass": "premium",
|
"gpuClass": "premium",
|
||||||
"kernelspec": {
|
"kernelspec": {
|
||||||
|
|||||||
Reference in New Issue
Block a user