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

Added terrain generation to mjspec tutorial

Browse Source 
PiperOrigin-RevId: 764229337
Change-Id: Ic541eefc4b561d3c39c5ee99c7d38466d7e75903
This commit is contained in:
Tom Power
2025-05-28 06:12:22 -07:00
committed by Copybara-Service
parent 56b88b59f0
commit 3a4b6e6c5b
3 changed files with 848 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mjspec.ipynb
+839 -5
View File
@@ -81,7 +81,7 @@
"print('Setting environment variable to use GPU rendering:')\n",
"%env MUJOCO_GL=egl\n",
"\n",
"# Check if installation was succesful.\n",
"# Check if installation was successful.\n",
"try:\n",
" print('Checking that the installation succeeded:')\n",
" import mujoco as mj\n",
@@ -539,7 +539,7 @@
"\n",
" Args:\n",
" shape: The shape of the generated array (tuple of two ints).\n",
" This must be a multple of res.\n",
" This must be a multiple 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",
@@ -851,6 +851,839 @@
"media.show_video(frames, fps=framerate )"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bZ-mpAKhSBHw"
},
"source": [
"## Terrain Generation\n",
"Here we will create a terrain out of different tiles. We start by creating each sinlge tile. Then we move on to putting tiles side by side to create a complete terrain."
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"cellView": "form",
"id": "d9oHHP6ZSZeK"
},
"outputs": [],
"source": [
"#@title Utilities\n",
"def render_tile(tile_func, direction=None, cam_distance=6, cam_elevation=-30):\n",
" arena_xml = \"\"\"\n",
" <mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global elevation=\"-10\" offwidth=\"2048\" offheight=\"1536\"/>\n",
" <quality shadowsize=\"8192\"/>\n",
" </visual>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\".5 .5 .5\" rgb2=\"0 0 0\" 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=\"300\" height=\"300\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texuniform=\"true\" texrepeat=\"5 5\" reflectance=\"0.3\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" <light pos=\"0 0 5\" diffuse=\"1 1 1\" specular=\"1 1 1\"/>\n",
" </worldbody>\n",
" </mujoco>\n",
" \"\"\"\n",
"\n",
" spec = mj.MjSpec.from_string(arena_xml)\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" name = 'base_tile'\n",
"\n",
" spec.worldbody.add_body(pos=[-3, 0, 0], name=name)\n",
" if direction:\n",
" tile_func(spec, direction=direction)\n",
" else:\n",
" tile_func(spec)\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 = cam_distance\n",
" cam.elevation = cam_elevation\n",
"\n",
" height = 300\n",
"\n",
" with mj.Renderer(model, 480, 640) as renderer:\n",
" mj.mj_forward(model, data)\n",
" renderer.update_scene(data,cam)\n",
" media.show_image(renderer.render(), height=height)\n",
"\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 multiple 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",
"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"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "dtFWSSwWSgeD"
},
"outputs": [],
"source": [
"# @title Stairs\n",
"def stairs(spec=None, grid_loc=[0, 0] , num_stairs=4, direction=1, name='stair'):\n",
" SQUARE_LENGTH = 2\n",
" V_SIZE = 0.076\n",
" H_SIZE = 0.12\n",
" H_STEP = H_SIZE * 2\n",
" V_STEP = V_SIZE * 2\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" # Offset\n",
" x_beginning, y_end = [-SQUARE_LENGTH + H_SIZE] * 2\n",
" x_end, y_beginning = [SQUARE_LENGTH - H_SIZE] * 2\n",
" # Dimension\n",
" size_one = [H_SIZE, SQUARE_LENGTH, V_SIZE]\n",
" size_two = [SQUARE_LENGTH, H_SIZE, V_SIZE]\n",
" # Geoms positions\n",
" x_pos_l = [x_beginning, 0, direction * V_SIZE]\n",
" x_pos_r = [x_end, 0, direction * V_SIZE]\n",
" y_pos_up = [0, y_beginning, direction * V_SIZE]\n",
" y_pos_down = [0, y_end, direction * V_SIZE]\n",
"\n",
" for i in range(num_stairs):\n",
" size_one[1] = SQUARE_LENGTH - H_STEP * i\n",
" size_two[0] = SQUARE_LENGTH - H_STEP * i\n",
"\n",
" x_pos_l[2], x_pos_r[2], y_pos_up[2], y_pos_down[2] = [\n",
" direction * ( V_SIZE + V_STEP * i)] * 4\n",
"\n",
" # Left side\n",
" x_pos_l[0] = x_beginning + H_STEP * i\n",
" body.add_geom(pos=x_pos_l, size=size_one, rgba=BROWN)\n",
" # Right side\n",
" x_pos_r[0] = x_end - H_STEP * i\n",
" body.add_geom(pos=x_pos_r, size=size_one, rgba=BROWN)\n",
" # Top\n",
" y_pos_up[1] = y_beginning - H_STEP * i\n",
" body.add_geom(pos=y_pos_up, size=size_two, rgba=BROWN)\n",
" # Bottom\n",
" y_pos_down[1] = y_end + H_STEP * i\n",
" body.add_geom(pos=y_pos_down, size=size_two, rgba=BROWN)\n",
"\n",
" # Closing\n",
" size = [SQUARE_LENGTH - H_STEP * num_stairs,\n",
" SQUARE_LENGTH - H_STEP * num_stairs,\n",
" V_SIZE]\n",
" pos = [0, 0,\n",
" direction * (V_SIZE + V_STEP * num_stairs)]\n",
" body.add_geom(pos=pos, size=size, rgba=BROWN)\n",
"\n",
"render_tile(stairs, direction=random.choice([-1, 1]))"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "YwX50AtzSmnT"
},
"outputs": [],
"source": [
"# @title Debris (Geoms)\n",
"def debris_with_simple_geoms(spec=None, grid_loc=[0, 0], name='plane'):\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
" RED = [0.6, 0.12, 0.15, 1.0]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN )\n",
"\n",
" # Simple Geoms\n",
" x_beginning, y_end = [-SQUARE_LENGTH + THICKNESS] * 2\n",
" x_end, y_beginning = [SQUARE_LENGTH - THICKNESS] * 2\n",
"\n",
" x_grid = np.linspace(x_beginning, x_end, 10)\n",
" y_grid = np.linspace(y_beginning, y_end, 10)\n",
"\n",
" for i in range(10):\n",
" x = np.random.choice(x_grid)\n",
" y = np.random.choice(y_grid)\n",
"\n",
" pos=[grid_loc[0] + x, grid_loc[1] + y, 0.2]\n",
"\n",
" g_type = None\n",
" size = None\n",
" if random.randint(0, 1):\n",
" g_type = mj.mjtGeom.mjGEOM_BOX\n",
" size = [0.1, 0.1, 0.02]\n",
" else:\n",
" g_type = mj.mjtGeom.mjGEOM_CYLINDER\n",
" size = [0.1, 0.02, 0]\n",
"\n",
" body = spec.worldbody.add_body(pos=pos, name=f'g{i}_{name}', mass=1)\n",
" body.add_geom(type=g_type, size=size, rgba=RED)\n",
" body.add_freejoint()\n",
"\n",
"render_tile(debris_with_simple_geoms)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "Wm_0DHK8S2kQ"
},
"outputs": [],
"source": [
"# @title Debris (Mesh)\n",
"def debris(spec=None, grid_loc=[0, 0] , name='debris'):\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
" STEP = THICKNESS * 8\n",
" SCALE = 0.1\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
" RED = [0.6, 0.12, 0.15, 1.0]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
" main.mesh.scale = np.array([SCALE]*3, dtype=np.float64)\n",
"\n",
" x_beginning = -SQUARE_LENGTH + THICKNESS\n",
" y_beginning = SQUARE_LENGTH - THICKNESS\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n",
"\n",
" # Place debris on the tile\n",
" for i in range(10):\n",
" for j in range(10):\n",
" # draw on xy plane\n",
" drawing = np.random.normal(size=(4, 2))\n",
" drawing /= np.linalg.norm(drawing, axis=1, keepdims=True)\n",
" z = np.zeros((drawing.shape[0], 1))\n",
" # Add z value to drawing\n",
" base = np.concatenate((drawing, z), axis=1)\n",
" # Extrude drawing\n",
" z_extrusion = np.full((drawing.shape[0], 1), THICKNESS * 4)\n",
" top = np.concatenate((drawing, z_extrusion), axis=1)\n",
" # Combine to get a mesh\n",
" mesh = np.vstack((base, top))\n",
"\n",
" # Create body and add the mesh to the geom of the body\n",
" spec.add_mesh(name=f'd{i}_{j}_{name}', uservert=mesh.flatten())\n",
" pos=[grid_loc[0] + x_beginning + i * STEP,\n",
" grid_loc[1] + y_beginning - j * STEP,\n",
" 0.2]\n",
"\n",
" body = spec.worldbody.add_body(pos=pos, name=f'd{i}_{j}_{name}', mass=1)\n",
" body.add_geom(type=mj.mjtGeom.mjGEOM_MESH, meshname=f'd{i}_{j}_{name}',\n",
" rgba=RED)\n",
" body.add_freejoint()\n",
"\n",
"render_tile(debris)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "BqTG_K8uS3v_"
},
"outputs": [],
"source": [
"# @title Boxy Terrain\n",
"def boxy_terrain(spec=None, grid_loc=[0, 0], name='boxy_terrain'):\n",
" SQUARE_LENGTH = 2\n",
" CUBE_LENGTH = 0.05\n",
" GRID_SIZE = int(SQUARE_LENGTH / CUBE_LENGTH)\n",
" STEP = CUBE_LENGTH * 2\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
"\n",
" if spec == None:\n",
" spec=mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
"\n",
" x_beginning = -SQUARE_LENGTH + CUBE_LENGTH\n",
" y_beginning = SQUARE_LENGTH - CUBE_LENGTH\n",
" for i in range(GRID_SIZE):\n",
" for j in range(GRID_SIZE):\n",
" body.add_geom(\n",
" pos=[x_beginning + i * STEP ,\n",
" y_beginning - j * STEP ,\n",
" random.randint(-1, 1) * CUBE_LENGTH\n",
" ],\n",
" size=[CUBE_LENGTH] * 3,\n",
" rgba=BROWN\n",
" )\n",
"\n",
"render_tile(boxy_terrain)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "gfQ1CafnS7hs"
},
"outputs": [],
"source": [
"# @title Box (Extrusion | Cut)\n",
"def box_extrusions(spec=None, grid_loc=[0, 0], complex=False,\n",
" name='box_extrusions'):\n",
" # Warning! complex sometimes leads to creation of holes\n",
" SQUARE_LENGTH = 2\n",
" CUBE_LENGTH = 0.05\n",
" GRID_SIZE = int(SQUARE_LENGTH / CUBE_LENGTH)\n",
" STEP = CUBE_LENGTH * 2\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
"\n",
" x_beginning = -SQUARE_LENGTH + CUBE_LENGTH\n",
" y_beginning = SQUARE_LENGTH - CUBE_LENGTH\n",
"\n",
" # Create initial grid and store geoms ref\n",
" grid = [[ 0 for _ in range(GRID_SIZE)] for _ in range(GRID_SIZE)]\n",
" for i in range(GRID_SIZE):\n",
" for j in range(GRID_SIZE):\n",
" ref = body.add_geom(\n",
" pos=[x_beginning + i * STEP, y_beginning - j * STEP, 0],\n",
" size=[CUBE_LENGTH] * 3,\n",
" rgba = BROWN\n",
" )\n",
" grid[i][j] = ref\n",
"\n",
" # Extrude or Cut operation using the boxes\n",
" for _ in range(random.randint(4, 50)):\n",
" box = None\n",
" while box == None:\n",
" # Create a box\n",
" start = (random.randint(0, GRID_SIZE - 2), random.randint(0, GRID_SIZE - 2))\n",
" dim = (random.randint(0, GRID_SIZE - 2), random.randint(0, GRID_SIZE-2))\n",
" # Make suer box is valid\n",
" if start[0] + dim [0] < len(grid) and start[1] + dim [1] < len(grid):\n",
" box = {\"start\":start, \"dim\":dim}\n",
"\n",
" # Use the box to Cut or Extrude\n",
" operation = random.choice([1, -1])\n",
" start = box[\"start\"]\n",
" dim = box[\"dim\"]\n",
" for i in range(start[0], dim[0]):\n",
" for j in range(start[1], dim[1]):\n",
" tile = grid[i][j]\n",
" if complex:\n",
" tile.pos[2] += operation * CUBE_LENGTH\n",
" else:\n",
" tile.pos[2] = operation * CUBE_LENGTH\n",
"\n",
"render_tile(box_extrusions)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "e1PEGQ37S_ee"
},
"outputs": [],
"source": [
"# @title Heightfield\n",
"def h_field(spec=None, grid_loc=[0, 0], name='h_field'):\n",
" SQUARE_LENGTH = 2\n",
" HEIGHT = 0.1\n",
" BROWN_RGBA = [0.460, 0.362, 0.216, 1.0]\n",
"\n",
" if spec is None:\n",
" spec = mj.MjSpec()\n",
"\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=name,\n",
" size=[SQUARE_LENGTH, SQUARE_LENGTH,\n",
" HEIGHT, HEIGHT/10],\n",
" nrow=noise.shape[0],\n",
" ncol=noise.shape[1],\n",
" userdata=noise.flatten())\n",
"\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(type=mj.mjtGeom.mjGEOM_HFIELD, hfieldname=name,\n",
" rgba=BROWN_RGBA)\n",
"\n",
"render_tile(h_field)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "UwDVd6aUTRKF"
},
"outputs": [],
"source": [
"# @title Floating platform\n",
"def floating_platform(spec=None, gird_loc=[0, 0, 0], name='platform'):\n",
" PLATFORM_LENGTH = 0.5\n",
" WIDTH = 0.12\n",
" INWARD_OFFSET = 0.008\n",
" THICKNESS = 0.005\n",
" SIZE = [PLATFORM_LENGTH, WIDTH, THICKNESS]\n",
" TENDON_LENGTH = 0.5\n",
" Z_OFFSET = 0.1\n",
"\n",
" GOLD = [0.850, 0.838, 0.119, 1]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Platform with sites\n",
" gird_loc[2] += Z_OFFSET\n",
" platform = spec.worldbody.add_body(pos=gird_loc, name=name)\n",
" platform.add_geom(size=SIZE, rgba=GOLD)\n",
" platform.add_freejoint()\n",
"\n",
" for x_dir in [-1, 1]:\n",
" for y_dir in [-1, 1]:\n",
" # Add site to world\n",
" vector = np.array([x_dir * PLATFORM_LENGTH,\n",
" y_dir * (WIDTH - INWARD_OFFSET)])\n",
" x_w = gird_loc[0] + vector[0]\n",
" y_w = gird_loc[1] + vector[1]\n",
" z_w = gird_loc[2] + TENDON_LENGTH\n",
" # Rotate sites by theta\n",
" spec.worldbody.add_site(name=f'{name}_hook_{x_dir}_{y_dir}',\n",
" pos=[ x_w, y_w, z_w],\n",
" size=[0.01, 0, 0])\n",
" # Add site to platform\n",
" x_p = x_dir * PLATFORM_LENGTH\n",
" y_p = y_dir * (WIDTH - INWARD_OFFSET)\n",
" platform.add_site(name=f'{name}_anchor_{x_dir}_{y_dir}',\n",
" pos=[ x_p, y_p, THICKNESS * 2],\n",
" size=[0.01, 0, 0])\n",
"\n",
" # Connect tendon to sites\n",
" thread = spec.add_tendon(name=f'{name}_thread_{x_dir}_{y_dir}',\n",
" limited=True,\n",
" range=[0, TENDON_LENGTH], width=0.01 )\n",
" thread.wrap_site(f'{name}_hook_{x_dir}_{y_dir}')\n",
" thread.wrap_site(f'{name}_anchor_{x_dir}_{y_dir}')\n",
"\n",
"render_tile(floating_platform, cam_distance=2, cam_elevation=-20)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "Wb4GmZewTSVI"
},
"outputs": [],
"source": [
"# @title Simple stairs\n",
"def simple_suspended_stair(spec=None, grid_loc=[0, 0], num_stair=20,\n",
" name=\"simple_suspended_stair\"):\n",
" BROWN= [0.460, 0.362, 0.216, 1.0]\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
" OFFSET_Y = -4/5 * SQUARE_LENGTH\n",
"\n",
" V_STEP = 0.076\n",
" H_STEP = 0.12\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Create tile\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n",
"\n",
" # Create Stairs\n",
" for i in range(num_stair):\n",
" floating_platform(spec,[grid_loc[0],\n",
" OFFSET_Y + grid_loc[1] + i * 2 * H_STEP,\n",
" i * V_STEP],\n",
" name =f'{name}_p_{i}')\n",
"\n",
"render_tile(simple_suspended_stair,cam_distance=7, cam_elevation=-30)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "AaoLjrKeTXu-"
},
"outputs": [],
"source": [
"# @title Sinusoidal stairs\n",
"def sin_suspended_stair(spec, grid_loc=[0, 0], num_stair=40,\n",
" name=\"sin_suspended_stair\"):\n",
" BROWN = [0.460, 0.362, 0.216, 1.0]\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
" OFFSET_Y = -4/5 * SQUARE_LENGTH\n",
"\n",
" V_STEP = 0.076\n",
" H_STEP = 0.12\n",
" AMPLITUDE = 0.2\n",
" FREQUENCY = 0.5\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
"\n",
" # Plane\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size=[SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba=BROWN)\n",
"\n",
" for i in range(num_stair):\n",
" x_step = AMPLITUDE * np.sin(2 * np.pi * FREQUENCY * (i * H_STEP))\n",
" floating_platform(spec, [grid_loc[0] + x_step,\n",
" OFFSET_Y + grid_loc[1] + i * 2 * H_STEP,\n",
" i * V_STEP],\n",
" name=f'{name}_p_{i}')\n",
"\n",
"render_tile(sin_suspended_stair,cam_distance=7, cam_elevation=-30)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "2_Ux8h-PTqbW"
},
"outputs": [],
"source": [
"# @title Floating platform for circular stair\n",
"def floating_platform_for_circular_stair(spec=None, gird_loc=[0, 0, 0] ,theta=0,\n",
" name='platform'):\n",
" PLATFORM_LENGTH = 0.5\n",
" TENDON_LENGTH = 0.5\n",
" WIDTH = 0.12/4 # Platform (body) is made of 4 separate geoms\n",
" THICKNESS = 0.005\n",
" SIZE = [PLATFORM_LENGTH, WIDTH, THICKNESS]\n",
" Z_OFFSET = 0.1\n",
"\n",
" GOLD = [0.850, 0.838, 0.119, 1]\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
" spec.compiler.degree = False\n",
"\n",
" # Platform with sites\n",
" gird_loc[2] += Z_OFFSET\n",
" platform = spec.worldbody.add_body(pos=gird_loc, name=name, euler=[0, 0, theta])\n",
" platform.add_geom(pos=[0, 0, 0] , size=SIZE, euler=[0, 0, 0],rgba=GOLD)\n",
" platform.add_geom(pos=[0, 0.02, 0], size=SIZE, euler=[0, 0, 0.05],rgba=GOLD)\n",
" platform.add_geom(pos=[0, 0.05, 0], size=SIZE, euler=[0, 0, 0.1],rgba=GOLD)\n",
" platform.add_geom(pos=[0, 0.08, 0], size=SIZE, euler=[0, 0, 0.15],rgba=GOLD)\n",
" platform.add_freejoint()\n",
"\n",
" for i, x_dir in enumerate([-1, 1]):\n",
" for j, y_dir in enumerate([-1, 1]):\n",
" # Rotate sites by theta\n",
" rotation_matrix = np.array([[np.cos(-theta), -np.sin(-theta)],\n",
" [np.sin(-theta), np.cos(-theta)]])\n",
" vector = np.array([x_dir * PLATFORM_LENGTH, y_dir * WIDTH ])\n",
" if i + j == 2:\n",
" vector = np.array([x_dir * PLATFORM_LENGTH, y_dir * 6 * WIDTH ])\n",
" vector = np.dot(vector , rotation_matrix)\n",
" x_w = gird_loc[0] + vector[0]\n",
" y_w = gird_loc[1] + vector[1]\n",
" z_w = gird_loc[2] + TENDON_LENGTH\n",
"\n",
" # Add site to world\n",
" spec.worldbody.add_site(name=f'{name}_hook_{x_dir}_{y_dir}',\n",
" pos=[ x_w, y_w, z_w],\n",
" size=[0.01, 0, 0])\n",
" # Add site to platform\n",
" x_p = x_dir * PLATFORM_LENGTH\n",
" y_p = y_dir * WIDTH\n",
" if i + j == 2:\n",
" y_p = y_dir * 6 * WIDTH\n",
" platform.add_site(name=f'{name}_anchor_{x_dir}_{y_dir}',\n",
" pos=[x_p, y_p, THICKNESS * 2],\n",
" size=[0.01, 0, 0])\n",
"\n",
" # Connect tendon to sites\n",
" thread = spec.add_tendon(name=f'{name}_thread_{x_dir}_{y_dir}', limited=True,\n",
" range=[0, TENDON_LENGTH], width=0.01 )\n",
" thread.wrap_site(f'{name}_hook_{x_dir}_{y_dir}')\n",
" thread.wrap_site(f'{name}_anchor_{x_dir}_{y_dir}')\n",
"\n",
"render_tile(floating_platform_for_circular_stair,cam_distance=2, cam_elevation=-40)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "UXSDdVF8TwmQ"
},
"outputs": [],
"source": [
"# @title Circular stairs\n",
"def circular_stairs(spec, grid_loc=[0, 0], num_stair=60, name=\"circular_stairs\"):\n",
" BROWN_RGBA = [0.460, 0.362, 0.216, 1.0]\n",
" SQUARE_LENGTH = 2\n",
" THICKNESS = 0.05\n",
"\n",
" RADIUS = 1.5\n",
" V_STEP = 0.076\n",
"\n",
" if spec == None:\n",
" spec = mj.MjSpec()\n",
"\n",
" # Defaults\n",
" main = spec.default\n",
" main.geom.type = mj.mjtGeom.mjGEOM_BOX\n",
" spec.compiler.degree = False\n",
"\n",
" # Plane\n",
" body = spec.worldbody.add_body(pos=grid_loc + [0], name=name)\n",
" body.add_geom(size = [SQUARE_LENGTH, SQUARE_LENGTH, THICKNESS], rgba = BROWN_RGBA )\n",
"\n",
" theta_step = 2 * np.pi / num_stair\n",
" for i in range(num_stair):\n",
" theta = i * theta_step\n",
" x = grid_loc[0] + RADIUS * np.cos(theta)\n",
" y = grid_loc[1] + RADIUS * np.sin(theta)\n",
" z = i * V_STEP\n",
"\n",
" floating_platform_for_circular_stair(spec, [x, y, z], theta=theta, name=f'{name}_p_{i}')\n",
"\n",
"render_tile(circular_stairs,cam_distance=12, cam_elevation=-30)"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "lsDky1KjT30W"
},
"outputs": [],
"source": [
"# @title Tile Generator\n",
"def add_tile(spec=None, grid_loc=[0, 0]):\n",
" if spec is None:\n",
" spec = mj.MjSpec()\n",
"\n",
" tile_type = random.randint(0, 9)\n",
"\n",
" if tile_type == 0:\n",
" debris_with_simple_geoms(spec, grid_loc, name=f\"plane_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 1:\n",
" stairs(spec, grid_loc, name=f\"stairs_up_{grid_loc[0]}_{grid_loc[1]}\",direction=1)\n",
" elif tile_type == 2:\n",
" stairs(spec, grid_loc, name=f\"stairs_down_{grid_loc[0]}_{grid_loc[1]}\",direction=-1)\n",
" elif tile_type == 3:\n",
" debris(spec, grid_loc, name=f\"debris_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 4:\n",
" box_extrusions(spec, grid_loc, name=f\"box_extrusions_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 5:\n",
" boxy_terrain(spec, grid_loc, name=f\"boxy_terrain_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 6:\n",
" h_field(spec, grid_loc, name=f\"h_field_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 7:\n",
" simple_suspended_stair(spec, grid_loc, name=f\"sss_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 8:\n",
" sin_suspended_stair(spec, grid_loc, name=f\"sinss_{grid_loc[0]}_{grid_loc[1]}\")\n",
" elif tile_type == 9:\n",
" circular_stairs(spec, grid_loc, name=f\"circular_s_{grid_loc[0]}_{grid_loc[1]}\")\n",
" return spec"
]
},
{
"cell_type": "code",
"execution_count": 0,
"metadata": {
"id": "P2F0ObC2T8w8"
},
"outputs": [],
"source": [
"# @title Generate Terrain\n",
"arena_xml = \"\"\"\n",
"<mujoco>\n",
" <visual>\n",
" <headlight diffuse=\".5 .5 .5\" specular=\"1 1 1\"/>\n",
" <global elevation=\"-10\" offwidth=\"2048\" offheight=\"1536\"/>\n",
" <quality shadowsize=\"8192\"/>\n",
" </visual>\n",
"\n",
" <asset>\n",
" <texture type=\"skybox\" builtin=\"gradient\" rgb1=\".5 .5 .5\" rgb2=\"0 0 0\" 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=\"300\" height=\"300\"/>\n",
" <material name=\"groundplane\" texture=\"groundplane\" texuniform=\"true\" texrepeat=\"5 5\" reflectance=\"0.3\"/>\n",
" </asset>\n",
"\n",
" <worldbody>\n",
" </worldbody>\n",
"</mujoco>\n",
"\"\"\"\n",
"\n",
"spec = mj.MjSpec.from_string(arena_xml)\n",
"\n",
"spec.option.enableflags |= mj.mjtEnableBit.mjENBL_OVERRIDE\n",
"spec.option.enableflags |= mj.mjtEnableBit.mjENBL_MULTICCD\n",
"spec.option.timestep = 0.0001\n",
"spec.compiler.degree = False\n",
"\n",
"main = spec.default\n",
"main.geom.solref = [0.001, 1]\n",
"\n",
"# Add lights\n",
"for x in [-1, 1]:\n",
" for y in [-1, 1]:\n",
" spec.worldbody.add_light(pos=[x, y, 40], dir=[-x, -y, -15])\n",
"\n",
"SQUARE_LENGTH = 2\n",
"for i in range(-2, 2):\n",
" for j in range(-2, 2):\n",
" add_tile(spec=spec, grid_loc=[i * 2 * SQUARE_LENGTH, j * 2 * SQUARE_LENGTH])\n",
"\n",
"model = spec.compile()\n",
"data = mj.MjData(model)\n",
"\n",
"cam = mj.MjvCamera()\n",
"mj.mjv_defaultCamera(cam)\n",
"cam.lookat = [-2, 0, -2]\n",
"cam.distance = 18\n",
"cam.elevation = -30\n",
"\n",
"with mj.Renderer(model, 720, 1280) as renderer:\n",
" mj.mj_forward(model, data)\n",
" renderer.update_scene(data,cam)\n",
" media.show_image(renderer.render())"
]
},
{
"cell_type": "markdown",
"metadata": {
@@ -864,6 +1697,7 @@
"cell_type": "code",
"execution_count": 0,
"metadata": {
"cellView": "form",
"id": "223KzKAzLdEJ"
},
"outputs": [],
@@ -1263,7 +2097,7 @@
"id": "RYbaTPNmLdEK"
},
"source": [
"We can scale the size of a model by traversing the kinematic tree and applying the the scale to the relevant geoms. Above we can see humanoids of three different sizes."
"We can scale the size of a model by traversing the kinematic tree and applying the scale to the relevant geoms. Above we can see humanoids of three different sizes."
]
},
{
@@ -1605,7 +2439,7 @@
"Note that:\n",
"\n",
"- MJCF attributes correspond directly to arguments of the `add_()` methods.\n",
"- When referencing elements, e.g when specifying the joint to which an actuator is attached, the name string of the MJCF elements is used."
"- When referencing elements, e.g. when specifying the joint to which an actuator is attached, the name string of the MJCF elements is used."
]
},
{
@@ -1792,7 +2626,7 @@
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"yXY7HGfVsVlo"
"sJFuNetilv4m"
],
"gpuClass": "premium",
"private_outputs": true,