Reorder model editing notebook, placing dm_control example last
PiperOrigin-RevId: 760642275 Change-Id: I7d8a625ab1d6d8364de93571390e8c14289e070a
This commit is contained in:
Tom Power
committed by
Copybara-Service
Copybara-Service
parent
6c201d3e1c
commit
670c97c4f7
+286
-288
@@ -853,269 +853,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "3N4YEIVt75_T"
|
||||
},
|
||||
"source": [
|
||||
"# `dm_control` example"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "TcQuv56BwaJf"
|
||||
},
|
||||
"source": [
|
||||
"A key feature is the ability to easily attach multiple models into a larger one. Disambiguation of duplicated names from different\n",
|
||||
"models, or multiple instances of the same model is handled via user-defined namespacing.\n",
|
||||
"\n",
|
||||
"One example use case is when we want robots with a variable number of joints, as this is a fundamental change to the kinematic structure. The snippets below follow the lines of the [example in dm_control](https://arxiv.org/abs/2006.12983), an older package with similar capabilities."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "7C-hfbtj8nRV"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"leg_model = \"\"\"\n",
|
||||
"<mujoco>\n",
|
||||
" <compiler angle=\"radian\"/>\n",
|
||||
"\n",
|
||||
" <default>\n",
|
||||
" <joint damping=\"2\" type=\"hinge\"/>\n",
|
||||
" <geom type=\"capsule\"/>\n",
|
||||
" </default>\n",
|
||||
"\n",
|
||||
" <worldbody>\n",
|
||||
" <body name=\"thigh\">\n",
|
||||
" <joint name=\"hip\" axis=\"0 0 1\"/>\n",
|
||||
" <body name=\"shin\">\n",
|
||||
" <joint name=\"knee\" axis=\"0 1 0\"/>\n",
|
||||
" </body>\n",
|
||||
" </body>\n",
|
||||
" </worldbody>\n",
|
||||
"\n",
|
||||
" <actuator>\n",
|
||||
" <position joint=\"hip\" kp=\"10\" name=\"hip\"/>\n",
|
||||
" <position joint=\"knee\" kp=\"10\" name=\"knee\"/>\n",
|
||||
" </actuator>\n",
|
||||
"</mujoco>\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"class Leg(object):\n",
|
||||
" \"\"\"A 2-DoF leg with position actuators.\"\"\"\n",
|
||||
" def __init__(self, length, rgba):\n",
|
||||
" self.spec = mj.MjSpec.from_string(leg_model)\n",
|
||||
"\n",
|
||||
" # Thigh:\n",
|
||||
" thigh = self.spec.body('thigh')\n",
|
||||
" thigh.add_geom(fromto=[0, 0, 0, length, 0, 0], size=[length/4, 0, 0], rgba=rgba)\n",
|
||||
"\n",
|
||||
" # Hip:\n",
|
||||
" shin = self.spec.body('shin')\n",
|
||||
" shin.add_geom(fromto=[0, 0, 0, 0, 0, -length], size=[length/5, 0, 0], rgba=rgba)\n",
|
||||
" shin.pos[0] = length"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "MQGsxnIB_RLO"
|
||||
},
|
||||
"source": [
|
||||
"The `Leg` class describes an abstract articulated leg, with two joints and corresponding proportional-derivative actuators.\n",
|
||||
"\n",
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "kMiuMyZW_XoB"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"BODY_RADIUS = 0.1\n",
|
||||
"random_state = np.random.RandomState(42)\n",
|
||||
"creature_model = \"\"\"\n",
|
||||
"<mujoco>\n",
|
||||
" <compiler angle=\"radian\"/>\n",
|
||||
"\n",
|
||||
" <worldbody>\n",
|
||||
" <geom name=\"torso\" type=\"ellipsoid\" size=\"{} {} {}\"/>\n",
|
||||
" </worldbody>\n",
|
||||
"</mujoco>\n",
|
||||
"\"\"\".format(BODY_RADIUS, BODY_RADIUS, BODY_RADIUS / 2)\n",
|
||||
"\n",
|
||||
"def make_creature(num_legs):\n",
|
||||
" \"\"\"Constructs a creature with `num_legs` legs.\"\"\"\n",
|
||||
" rgba = random_state.uniform([0, 0, 0, 1], [1, 1, 1, 1])\n",
|
||||
" spec = mj.MjSpec.from_string(creature_model)\n",
|
||||
" spec.copy_during_attach = True\n",
|
||||
"\n",
|
||||
" # Attach legs to equidistant sites on the circumference.\n",
|
||||
" spec.worldbody.first_geom().rgba = rgba\n",
|
||||
" leg = Leg(length=BODY_RADIUS, rgba=rgba)\n",
|
||||
" for i in range(num_legs):\n",
|
||||
" theta = 2 * i * np.pi / num_legs\n",
|
||||
" hip_pos = BODY_RADIUS * np.array([np.cos(theta), np.sin(theta), 0])\n",
|
||||
" hip_site = spec.worldbody.add_site(pos=hip_pos, euler=[0, 0, theta])\n",
|
||||
" hip_site.attach_body(leg.spec.body('thigh'), '', '-' + str(i))\n",
|
||||
"\n",
|
||||
" return spec"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QMQ3jc6-_toj"
|
||||
},
|
||||
"source": [
|
||||
"The `make_creature` function uses the `attach()` method to procedurally attach legs to the torso. Note that at this stage both the torso and hip attachment sites are children of the `worldbody`, since their parent body has yet to be instantiated. We'll now make an arena with a chequered floor and two lights, and place our creatures in a grid."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "vt2JwXd__1cT"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Six Creatures on a floor {vertical-output: true}\n",
|
||||
"\n",
|
||||
"arena = mj.MjSpec()\n",
|
||||
"\n",
|
||||
"if hasattr(arena, 'compiler'):\n",
|
||||
" arena.compiler.degree = False # MuJoCo dev (next release).\n",
|
||||
"else:\n",
|
||||
" arena.degree = False # MuJoCo release\n",
|
||||
"\n",
|
||||
"# Make arena with textured floor.\n",
|
||||
"chequered = arena.add_texture(\n",
|
||||
" name=\"chequered\", type=mj.mjtTexture.mjTEXTURE_2D,\n",
|
||||
" builtin=mj.mjtBuiltin.mjBUILTIN_CHECKER,\n",
|
||||
" width=300, height=300, rgb1=[.2, .3, .4], rgb2=[.3, .4, .5])\n",
|
||||
"grid = arena.add_material(\n",
|
||||
" name='grid', texrepeat=[5, 5], reflectance=.2\n",
|
||||
" ).textures[mj.mjtTextureRole.mjTEXROLE_RGB] = 'chequered'\n",
|
||||
"arena.worldbody.add_geom(\n",
|
||||
" type=mj.mjtGeom.mjGEOM_PLANE, size=[2, 2, .1], material='grid')\n",
|
||||
"for x in [-2, 2]:\n",
|
||||
" arena.worldbody.add_light(pos=[x, -1, 3], dir=[-x, 1, -2])\n",
|
||||
"\n",
|
||||
"# Instantiate 6 creatures with 3 to 8 legs.\n",
|
||||
"creatures = [make_creature(num_legs=num_legs) for num_legs in range(3, 9)]\n",
|
||||
"\n",
|
||||
"# Place them on a grid in the arena.\n",
|
||||
"height = .15\n",
|
||||
"grid = 5 * BODY_RADIUS\n",
|
||||
"xpos, ypos, zpos = np.meshgrid([-grid, 0, grid], [0, grid], [height])\n",
|
||||
"for i, spec in enumerate(creatures):\n",
|
||||
" # Place spawn sites on a grid.\n",
|
||||
" spawn_pos = (xpos.flat[i], ypos.flat[i], zpos.flat[i])\n",
|
||||
" spawn_site = arena.worldbody.add_site(pos=spawn_pos, group=3)\n",
|
||||
" # Attach to the arena at the spawn sites, with a free joint.\n",
|
||||
" spawn_body = spawn_site.attach_body(spec.worldbody, '', '-' + str(i))\n",
|
||||
" spawn_body.add_freejoint()\n",
|
||||
"\n",
|
||||
"# Instantiate the physics and render.\n",
|
||||
"model = arena.compile()\n",
|
||||
"render(model)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "mPUGkrCzAFMg"
|
||||
},
|
||||
"source": [
|
||||
"Multi-legged creatures, ready to roam! Let's inject some controls and watch them move. We'll generate a sinusoidal open-loop control signal of fixed frequency and random phase, recording both video frames and the horizontal positions of the torso geoms, in order to plot the movement trajectories."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "7gz9FfNzGxPO"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Video of the movement {vertical-output: true}\n",
|
||||
"\n",
|
||||
"data = mj.MjData(model)\n",
|
||||
"duration = 10 # (Seconds)\n",
|
||||
"framerate = 30 # (Hz)\n",
|
||||
"video = []\n",
|
||||
"pos_x = []\n",
|
||||
"pos_y = []\n",
|
||||
"geoms = arena.worldbody.find_all(mj.mjtObj.mjOBJ_GEOM)\n",
|
||||
"torsos_data = [data.bind(geom) for geom in geoms if 'torso' in geom.name]\n",
|
||||
"torsos_model = [model.bind(geom) for geom in geoms if 'torso' in geom.name]\n",
|
||||
"actuators = [data.bind(actuator) for actuator in arena.actuators]\n",
|
||||
"\n",
|
||||
"# Control signal frequency, phase, amplitude.\n",
|
||||
"freq = 5\n",
|
||||
"phase = 2 * np.pi * random_state.rand(len(arena.actuators))\n",
|
||||
"amp = 0.9\n",
|
||||
"\n",
|
||||
"# Simulate, saving video frames and torso locations.\n",
|
||||
"mj.mj_resetData(model, data)\n",
|
||||
"with mj.Renderer(model) as renderer:\n",
|
||||
" while data.time < duration:\n",
|
||||
" # Inject controls and step the physics.\n",
|
||||
" for i, actuator in enumerate(actuators):\n",
|
||||
" actuator.ctrl = amp * np.sin(freq * data.time + phase[i])\n",
|
||||
" mj.mj_step(model, data)\n",
|
||||
"\n",
|
||||
" # Save torso horizontal positions using name indexing.\n",
|
||||
" pos_x.append([torso.xpos[0] for torso in torsos_data])\n",
|
||||
" pos_y.append([torso.xpos[1] for torso in torsos_data])\n",
|
||||
"\n",
|
||||
" # Save video frames.\n",
|
||||
" if len(video) < data.time * framerate:\n",
|
||||
" renderer.update_scene(data)\n",
|
||||
" pixels = renderer.render()\n",
|
||||
" video.append(pixels.copy())\n",
|
||||
"\n",
|
||||
"media.show_video(video, fps=framerate)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "qt2L52e_Tcgt"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Movement trajectories {vertical-output: true}\n",
|
||||
"\n",
|
||||
"creature_colors = [torso.rgba[:3] for torso in torsos_model]\n",
|
||||
"fig, ax = plt.subplots(figsize=(4, 4))\n",
|
||||
"ax.set_prop_cycle(color=creature_colors)\n",
|
||||
"_ = ax.plot(pos_x, pos_y, linewidth=4)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "kSEUoxifxYJ4"
|
||||
},
|
||||
"source": [
|
||||
"The plot above shows the corresponding movement trajectories of creature positions. Note how `mjSpec` attribute `id` were used to access both `xpos` and `rgba` values. This attribute is valid only after a model is compiled."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "QZ8alJZz8cB1"
|
||||
"id": "IGd0uD64LdEJ"
|
||||
},
|
||||
"source": [
|
||||
"# Model editing"
|
||||
@@ -1125,8 +863,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "m4sppBqvf7yd"
|
||||
"id": "223KzKAzLdEJ"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1175,7 +912,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "JN3Z4v0PyXKa"
|
||||
"id": "eGgXNjQ8LdEK"
|
||||
},
|
||||
"source": [
|
||||
"`mjSpec` elements can be traversed in two ways:\n",
|
||||
@@ -1188,7 +925,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "8IcB7nezblyT"
|
||||
"id": "Len0o_idLdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1215,7 +952,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "hcGI4orhyzvc"
|
||||
"id": "GeiFFBYxLdEK"
|
||||
},
|
||||
"source": [
|
||||
"An `mjSpec` can be compiled multiple times. If the state has to be preserved between different compilations, then the function `recompile()` must be used, which returns a new `mjData` that contains the mapped state, possibly having a different dimension from the origin."
|
||||
@@ -1225,7 +962,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "uh_N1Fkqk-Mi"
|
||||
"id": "eiRXgh9OLdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1271,7 +1008,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "XmSlXirVzLqt"
|
||||
"id": "kuTWD415LdEK"
|
||||
},
|
||||
"source": [
|
||||
"Let us load the humanoid model and inspect it."
|
||||
@@ -1281,7 +1018,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "UywMzsp5Hnk2"
|
||||
"id": "5d1wmQM2LdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1296,7 +1033,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "owcmKeuSzQRy"
|
||||
"id": "38PXB1rWLdEK"
|
||||
},
|
||||
"source": [
|
||||
"We wish to remove the arms and replace them with the legs. This can be done by first storing the arm positions into frames attached to the torso. Then we can detach the arms and self-attach the legs into the frames."
|
||||
@@ -1306,7 +1043,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "qZCyv-B0IGiG"
|
||||
"id": "0eaNq0Q7LdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1341,7 +1078,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "LnEwEjW3zdua"
|
||||
"id": "HfhxL2EqLdEK"
|
||||
},
|
||||
"source": [
|
||||
"Similarly, different models can be attach together. Here, the right arm is detached and a robot arm from a different model is attached in its place."
|
||||
@@ -1351,7 +1088,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "w-NdFhSIIrLL"
|
||||
"id": "uS4LGbI7LdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1382,7 +1119,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "e_idaggAznXu"
|
||||
"id": "CWXYy_1uLdEK"
|
||||
},
|
||||
"source": [
|
||||
"When doing this, the actuators and all other objects referenced by the attached sub-tree are imported in the new model. All assets are currently imported, referenced or not."
|
||||
@@ -1392,8 +1129,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"id": "50lOgJ7mQ2bV"
|
||||
"id": "UwWDD-NHLdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1406,7 +1142,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "APDoWK4mz0aJ"
|
||||
"id": "hDvt3vcxLdEK"
|
||||
},
|
||||
"source": [
|
||||
"Domain randomization can be performed by attaching multiple times the same spec, edited each time with a new instance of randomized parameters."
|
||||
@@ -1416,7 +1152,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "oHjdgkISNLKy"
|
||||
"id": "oPPFbWawLdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1459,7 +1195,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "iXgYCVzEWFTU"
|
||||
"id": "PML2pxYgLdEK"
|
||||
},
|
||||
"source": [
|
||||
"## Model scaling"
|
||||
@@ -1469,7 +1205,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "-hSJKyH4A2VY"
|
||||
"id": "pcUNLmQBLdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1523,7 +1259,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tBH5vmqJXleD"
|
||||
"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."
|
||||
@@ -1533,7 +1269,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "cV4tkG6siFQp"
|
||||
"id": "u-ejx8lKLdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1646,7 +1382,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZSOra3S2YpIB"
|
||||
"id": "uKhDI_IfLdEK"
|
||||
},
|
||||
"source": [
|
||||
"We can also apply scaling to the actuators. In the humanoid case, scaling the geoms without scaling the `gear` parameter for the actuators results in a humanoid that can jump higher proportional to its size."
|
||||
@@ -1656,7 +1392,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "9IuwQQ0F2ddA"
|
||||
"id": "ovQIAxn7LdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1717,7 +1453,7 @@
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "D2DrPBvBZjI0"
|
||||
"id": "8o1daXIOLdEK"
|
||||
},
|
||||
"source": [
|
||||
"We can also apply scaling to the model non-uniformly. In this instance we scale the humanoid to have long limbs, by only applying the scale to the length of the capsule geoms for the arms, legs and feet."
|
||||
@@ -1727,7 +1463,7 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "1G8VO45v2ddA"
|
||||
"id": "UBSo2nfQLdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -1787,6 +1523,268 @@
|
||||
"model = spec.compile()\n",
|
||||
"render(model, height=400, camera=cam)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "qhXwxLe3LdEK"
|
||||
},
|
||||
"source": [
|
||||
"# dm_control example"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "enuJ_YIqLdEK"
|
||||
},
|
||||
"source": [
|
||||
"A key feature is the ability to easily attach multiple models into a larger one. Disambiguation of duplicated names from different\n",
|
||||
"models, or multiple instances of the same model is handled via user-defined namespacing.\n",
|
||||
"\n",
|
||||
"One example use case is when we want robots with a variable number of joints, as this is a fundamental change to the kinematic structure. The snippets below follow the lines of the [example in dm_control](https://arxiv.org/abs/2006.12983), an older package with similar capabilities."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "4p3P_dP8LdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"leg_model = \"\"\"\n",
|
||||
"<mujoco>\n",
|
||||
" <compiler angle=\"radian\"/>\n",
|
||||
"\n",
|
||||
" <default>\n",
|
||||
" <joint damping=\"2\" type=\"hinge\"/>\n",
|
||||
" <geom type=\"capsule\"/>\n",
|
||||
" </default>\n",
|
||||
"\n",
|
||||
" <worldbody>\n",
|
||||
" <body name=\"thigh\">\n",
|
||||
" <joint name=\"hip\" axis=\"0 0 1\"/>\n",
|
||||
" <body name=\"shin\">\n",
|
||||
" <joint name=\"knee\" axis=\"0 1 0\"/>\n",
|
||||
" </body>\n",
|
||||
" </body>\n",
|
||||
" </worldbody>\n",
|
||||
"\n",
|
||||
" <actuator>\n",
|
||||
" <position joint=\"hip\" kp=\"10\" name=\"hip\"/>\n",
|
||||
" <position joint=\"knee\" kp=\"10\" name=\"knee\"/>\n",
|
||||
" </actuator>\n",
|
||||
"</mujoco>\n",
|
||||
"\"\"\"\n",
|
||||
"\n",
|
||||
"class Leg(object):\n",
|
||||
" \"\"\"A 2-DoF leg with position actuators.\"\"\"\n",
|
||||
" def __init__(self, length, rgba):\n",
|
||||
" self.spec = mj.MjSpec.from_string(leg_model)\n",
|
||||
"\n",
|
||||
" # Thigh:\n",
|
||||
" thigh = self.spec.body('thigh')\n",
|
||||
" thigh.add_geom(fromto=[0, 0, 0, length, 0, 0], size=[length/4, 0, 0], rgba=rgba)\n",
|
||||
"\n",
|
||||
" # Hip:\n",
|
||||
" shin = self.spec.body('shin')\n",
|
||||
" shin.add_geom(fromto=[0, 0, 0, 0, 0, -length], size=[length/5, 0, 0], rgba=rgba)\n",
|
||||
" shin.pos[0] = length"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Mqr8rXLILdEK"
|
||||
},
|
||||
"source": [
|
||||
"The `Leg` class describes an abstract articulated leg, with two joints and corresponding proportional-derivative actuators.\n",
|
||||
"\n",
|
||||
"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."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "1z2NBpAPLdEK"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"BODY_RADIUS = 0.1\n",
|
||||
"random_state = np.random.RandomState(42)\n",
|
||||
"creature_model = \"\"\"\n",
|
||||
"<mujoco>\n",
|
||||
" <compiler angle=\"radian\"/>\n",
|
||||
"\n",
|
||||
" <worldbody>\n",
|
||||
" <geom name=\"torso\" type=\"ellipsoid\" size=\"{} {} {}\"/>\n",
|
||||
" </worldbody>\n",
|
||||
"</mujoco>\n",
|
||||
"\"\"\".format(BODY_RADIUS, BODY_RADIUS, BODY_RADIUS / 2)\n",
|
||||
"\n",
|
||||
"def make_creature(num_legs):\n",
|
||||
" \"\"\"Constructs a creature with `num_legs` legs.\"\"\"\n",
|
||||
" rgba = random_state.uniform([0, 0, 0, 1], [1, 1, 1, 1])\n",
|
||||
" spec = mj.MjSpec.from_string(creature_model)\n",
|
||||
" spec.copy_during_attach = True\n",
|
||||
"\n",
|
||||
" # Attach legs to equidistant sites on the circumference.\n",
|
||||
" spec.worldbody.first_geom().rgba = rgba\n",
|
||||
" leg = Leg(length=BODY_RADIUS, rgba=rgba)\n",
|
||||
" for i in range(num_legs):\n",
|
||||
" theta = 2 * i * np.pi / num_legs\n",
|
||||
" hip_pos = BODY_RADIUS * np.array([np.cos(theta), np.sin(theta), 0])\n",
|
||||
" hip_site = spec.worldbody.add_site(pos=hip_pos, euler=[0, 0, theta])\n",
|
||||
" hip_site.attach_body(leg.spec.body('thigh'), '', '-' + str(i))\n",
|
||||
"\n",
|
||||
" return spec"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "865FGuntLdEL"
|
||||
},
|
||||
"source": [
|
||||
"The `make_creature` function uses the `attach()` method to procedurally attach legs to the torso. Note that at this stage both the torso and hip attachment sites are children of the `worldbody`, since their parent body has yet to be instantiated. We'll now make an arena with a chequered floor and two lights, and place our creatures in a grid."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "2fPaSkgfLdEL"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Six Creatures on a floor {vertical-output: true}\n",
|
||||
"\n",
|
||||
"arena = mj.MjSpec()\n",
|
||||
"\n",
|
||||
"if hasattr(arena, 'compiler'):\n",
|
||||
" arena.compiler.degree = False # MuJoCo dev (next release).\n",
|
||||
"else:\n",
|
||||
" arena.degree = False # MuJoCo release\n",
|
||||
"\n",
|
||||
"# Make arena with textured floor.\n",
|
||||
"chequered = arena.add_texture(\n",
|
||||
" name=\"chequered\", type=mj.mjtTexture.mjTEXTURE_2D,\n",
|
||||
" builtin=mj.mjtBuiltin.mjBUILTIN_CHECKER,\n",
|
||||
" width=300, height=300, rgb1=[.2, .3, .4], rgb2=[.3, .4, .5])\n",
|
||||
"grid = arena.add_material(\n",
|
||||
" name='grid', texrepeat=[5, 5], reflectance=.2\n",
|
||||
" ).textures[mj.mjtTextureRole.mjTEXROLE_RGB] = 'chequered'\n",
|
||||
"arena.worldbody.add_geom(\n",
|
||||
" type=mj.mjtGeom.mjGEOM_PLANE, size=[2, 2, .1], material='grid')\n",
|
||||
"for x in [-2, 2]:\n",
|
||||
" arena.worldbody.add_light(pos=[x, -1, 3], dir=[-x, 1, -2])\n",
|
||||
"\n",
|
||||
"# Instantiate 6 creatures with 3 to 8 legs.\n",
|
||||
"creatures = [make_creature(num_legs=num_legs) for num_legs in range(3, 9)]\n",
|
||||
"\n",
|
||||
"# Place them on a grid in the arena.\n",
|
||||
"height = .15\n",
|
||||
"grid = 5 * BODY_RADIUS\n",
|
||||
"xpos, ypos, zpos = np.meshgrid([-grid, 0, grid], [0, grid], [height])\n",
|
||||
"for i, spec in enumerate(creatures):\n",
|
||||
" # Place spawn sites on a grid.\n",
|
||||
" spawn_pos = (xpos.flat[i], ypos.flat[i], zpos.flat[i])\n",
|
||||
" spawn_site = arena.worldbody.add_site(pos=spawn_pos, group=3)\n",
|
||||
" # Attach to the arena at the spawn sites, with a free joint.\n",
|
||||
" spawn_body = spawn_site.attach_body(spec.worldbody, '', '-' + str(i))\n",
|
||||
" spawn_body.add_freejoint()\n",
|
||||
"\n",
|
||||
"# Instantiate the physics and render.\n",
|
||||
"model = arena.compile()\n",
|
||||
"render(model)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tq5mKlc_LdEL"
|
||||
},
|
||||
"source": [
|
||||
"Multi-legged creatures, ready to roam! Let's inject some controls and watch them move. We'll generate a sinusoidal open-loop control signal of fixed frequency and random phase, recording both video frames and the horizontal positions of the torso geoms, in order to plot the movement trajectories."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "i37FpwCeLdEL"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Video of the movement {vertical-output: true}\n",
|
||||
"\n",
|
||||
"data = mj.MjData(model)\n",
|
||||
"duration = 10 # (Seconds)\n",
|
||||
"framerate = 30 # (Hz)\n",
|
||||
"video = []\n",
|
||||
"pos_x = []\n",
|
||||
"pos_y = []\n",
|
||||
"geoms = arena.worldbody.find_all(mj.mjtObj.mjOBJ_GEOM)\n",
|
||||
"torsos_data = [data.bind(geom) for geom in geoms if 'torso' in geom.name]\n",
|
||||
"torsos_model = [model.bind(geom) for geom in geoms if 'torso' in geom.name]\n",
|
||||
"actuators = [data.bind(actuator) for actuator in arena.actuators]\n",
|
||||
"\n",
|
||||
"# Control signal frequency, phase, amplitude.\n",
|
||||
"freq = 5\n",
|
||||
"phase = 2 * np.pi * random_state.rand(len(arena.actuators))\n",
|
||||
"amp = 0.9\n",
|
||||
"\n",
|
||||
"# Simulate, saving video frames and torso locations.\n",
|
||||
"mj.mj_resetData(model, data)\n",
|
||||
"with mj.Renderer(model) as renderer:\n",
|
||||
" while data.time < duration:\n",
|
||||
" # Inject controls and step the physics.\n",
|
||||
" for i, actuator in enumerate(actuators):\n",
|
||||
" actuator.ctrl = amp * np.sin(freq * data.time + phase[i])\n",
|
||||
" mj.mj_step(model, data)\n",
|
||||
"\n",
|
||||
" # Save torso horizontal positions using name indexing.\n",
|
||||
" pos_x.append([torso.xpos[0] for torso in torsos_data])\n",
|
||||
" pos_y.append([torso.xpos[1] for torso in torsos_data])\n",
|
||||
"\n",
|
||||
" # Save video frames.\n",
|
||||
" if len(video) < data.time * framerate:\n",
|
||||
" renderer.update_scene(data)\n",
|
||||
" pixels = renderer.render()\n",
|
||||
" video.append(pixels.copy())\n",
|
||||
"\n",
|
||||
"media.show_video(video, fps=framerate)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 0,
|
||||
"metadata": {
|
||||
"id": "uFrvaih4LdEL"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title Movement trajectories {vertical-output: true}\n",
|
||||
"\n",
|
||||
"creature_colors = [torso.rgba[:3] for torso in torsos_model]\n",
|
||||
"fig, ax = plt.subplots(figsize=(4, 4))\n",
|
||||
"ax.set_prop_cycle(color=creature_colors)\n",
|
||||
"_ = ax.plot(pos_x, pos_y, linewidth=4)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "FMW4l-fSLdEL"
|
||||
},
|
||||
"source": [
|
||||
"The plot above shows the corresponding movement trajectories of creature positions. Note how `mjSpec` attribute `id` were used to access both `xpos` and `rgba` values. This attribute is valid only after a model is compiled."
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
|
||||
Reference in New Issue
Block a user