⚠️ Running on Google Colab? Execution there can be up to ~50× slower than on a local machine or dedicated GPU, depending on the resources Colab makes available at the time. Use Colab only for testing and following along with the tutorials—not for production runs or benchmarking.
# --- Google Colab setup ---
# If you opened this notebook in Google Colab, run this cell first to
# install FlyGym and enable headless (EGL) rendering. It is a no-op when
# the notebook is run anywhere else.
import os
import sys
if "google.colab" in sys.modules:
os.environ["MUJOCO_GL"] = "egl"
os.environ["PYOPENGL_PLATFORM"] = "egl"
%pip install -q tqdm "flygym @ git+https://github.com/NeLy-EPFL/flygym.git@v2.1.0"
Advanced model composition¶
The previous tutorial demonstrated that you can compose "standard" scenes using the flygym.compose module. What if you need to customize it more? By design, FlyGym does not expose all available options in MuJoCo (that'd just be reinventing the wheel). Instead, MuJoCo's native model-editing API, mujoco.MjSpec, can be used to adjust these settings. Essentially, the entire model being simulated in MuJoCo is fully specified in an MJCF file (which follows the XML format). Libraries like mujoco.MjSpec and flygym.compose are simply "parsers and writers" that generate and modify these XML files in a human-friendly way.
In this tutorial, we will demonstrate how you can use mujoco.MjSpec to customize your model beyond what's natively supported by FlyGym. For all details, refer to the official model editing documentation of MuJoCo.
Modifying the MJCF model directly through mujoco.MjSpec¶
Let's compose a Fly and add it to a FlatGroundWorld, just as we did in the previous tutorial:
from flygym.compose import NeuroMechFly, KinematicPosePreset
from flygym.anatomy import AxisOrder, JointPreset, Skeleton, ActuatedDOFPreset
fly = NeuroMechFly(name="my_fly_model")
axis_order = AxisOrder.ROLL_PITCH_YAW
joint_preset = JointPreset.ALL_BIOLOGICAL
skeleton = Skeleton(joint_preset=joint_preset, axis_order=axis_order)
neutral_pose = KinematicPosePreset.NEUTRAL
joints = fly.add_joints(skeleton, neutral_pose=neutral_pose)
actuation_preset = ActuatedDOFPreset.LEGS_ACTIVE_ONLY
actuator_gain = 50
actuated_jointdofs = skeleton.get_actuated_dofs_from_preset(actuation_preset)
actuators = fly.add_actuators(
actuated_jointdofs,
actuator_type="position",
neutral_input=neutral_pose,
kp=actuator_gain,
)
fly.colorize()
camera = fly.add_tracking_camera()
A Fly on its own is just a model of the animal's body; to actually simulate it, we attach it to a World that provides the environment. Here we add our fly to a FlatGroundWorld — a flat ground plane that the fly stands on, connected to the world by a free joint. (See the previous tutorial for a detailed walkthrough of spawn poses and ground-contact presets.)
from flygym.compose import FlatGroundWorld
from flygym.anatomy import ContactBodiesPreset
from flygym.utils.math import Rotation3D
world = FlatGroundWorld()
spawn_pos = [0, 0, 0.7] # center of thorax 0.7 mm above the ground
spawn_rot = Rotation3D(format="quat", values=[1, 0, 0, 0]) # upright, facing forward
contact_preset = ContactBodiesPreset.LEGS_THORAX_ABDOMEN_HEAD
world.add_fly(fly, spawn_pos, spawn_rot, bodysegs_with_ground_contact=contact_preset)
But what do these translate to when we build the MJCF model specification? Let's export the world model and inspect the MJCF XML file:
from pathlib import Path
export_dir = Path("./demo_output/world_model/")
world.save_xml_with_assets(export_dir)
When we export the model to a directory, the MJCF file (in XML format) as well as all asset files (e.g. body geometric meshes) are saved under that directory:
!ls ./demo_output/world_model/
c_abdomen12.stl lf_tarsus2.stl lh_tibia.stl c_abdomen3.stl lf_tarsus3.stl lh_trochanterfemur.stl c_abdomen4.stl lf_tarsus4.stl lm_coxa.stl c_abdomen5.stl lf_tarsus5.stl lm_tarsus1.stl c_abdomen6.stl lf_tibia.stl lm_tarsus2.stl c_haustellum.stl lf_trochanterfemur.stl lm_tarsus3.stl c_head.stl l_funiculus.stl lm_tarsus4.stl c_rostrum.stl l_haltere.stl lm_tarsus5.stl c_thorax.stl lh_coxa.stl lm_tibia.stl flat_ground_world.xml lh_tarsus1.stl lm_trochanterfemur.stl l_arista.stl lh_tarsus2.stl l_pedicel.stl l_eye.stl lh_tarsus3.stl l_wing.stl lf_coxa.stl lh_tarsus4.stl lf_tarsus1.stl lh_tarsus5.stl
Please check the content of the actual XML file by downloading and opening the path below.
Before we start,
- If you are not familiar with the XML file format, see this short explanation from IBM.
- All possible XML elements and their full specifications can be found in the XML reference of the official MuJoCo documentation.
- Anything that is not explicitly defined in the XML file take the default value mentioned in the XML reference above.
xml_path = export_dir / f"{world.mjcf_root.modelname}.xml" # path to the actual XML
print(f"View XML file at this path: {xml_path}")
View XML file at this path: demo_output/world_model/flat_ground_world.xml
The first thing that you will notice is that the whole model is under a <mujoco model="flat_ground_world"> element, where "flat_ground_world" is the default name of the FlatGroundWorld we created. Our fly keeps its own name, "my_fly_model", as a namespace inside the world (more on that below).
The next few sections define global settings for the simulation. Note that MuJoCo only writes out attributes that differ from their defaults, so the exact set of attributes you see depends on the model.
The
<compiler ...>element defines settings such as lower bounds on the mass and inertia of simulated bodies and the native angle unit:<compiler angle="radian" boundmass="1e-06" boundinertia="1e-12"/>
The
<option ...>element defines core simulation settings such as timestep, integrator, solver, and contact handling. In particular,gravity="0 0 -9810"represents Earth gravity in the model's length unit. Because this model uses millimeters, gravity is written as -9810 mm/s⁻² (equivalent to -9.81 m/s⁻²) (see discussion below).<option timestep="0.0001" gravity="0 0 -9810" noslip_iterations="5"> <flag energy="enable"/> </option>
MuJoCo itself is unitless, so keeping units consistent across the model is the user's responsibility. When performing real-number arithmetics with discrete binary bits on a computer, keeping numbers closer in their orders of magnitude generally helps reduce round-off errors caused by floating point error. Therefore, since the fly is small, we use millimeter and gram as base units for length and mass instead of their SI counterparts, meter and kilogram. This means that the forces that we read out from the simulation are in g·mm·s⁻² (i.e., micronewton, μN), instead of kg·m·s⁻² (i.e., Newton). This helps us bring numerically computed numbers about 6 orders of magnitude closer (compared to length, which is around 1 mm).
The
<visual>element controls rendering-only parameters. Here,<headlight ...>sets scene lighting, and<map ...>sets scaling factors used to visualize forces and torques in the interactive viewer.<visual> <global offwidth="2048" offheight="2048"/> <headlight ambient="0.5 0.5 0.5" diffuse="0.6 0.6 0.6" specular="0.3 0.3 0.3"/> <map stiffness="1000000" stiffnessrot="5000000" force="1e-05" torque="0.0002"/> </visual>
Next, we have the <asset> section. It contains reusable resources such as mesh files, textures, and materials that are referenced by bodies and geometries elsewhere in the model. Assets that belong to the world itself (the ground's checker texture and grid material, and the skybox) keep plain names, while assets that belong to the fly are prefixed with the fly's name (e.g. my_fly_model/headthorax).
<asset>
<!-- Assets that belong to the world (the background and the ground plane) -->
<texture name="skybox" type="skybox" builtin="gradient" .../>
<texture name="checker" type="2d" builtin="checker" .../>
<material name="grid" texture="checker" .../>
<!-- Assets that belong to the fly, namespaced under the fly's name. -->
<!-- Note that the meshes are scaled up 1000x in x, y, and z. This is because we -->
<!-- use mm as the base unit for length in the simulation, but mesh geometries are -->
<!-- specified in SI units (m). The `file` attribute points to the mesh on disk; -->
<!-- when you export with `save_xml_with_assets`, the meshes are copied next to the -->
<!-- XML and the paths are rewritten to bare filenames so the export is self-contained. -->
<mesh
name="my_fly_model/c_thorax"
file="c_thorax.stl"
scale="1000 1000 1000"/>
<mesh
name="my_fly_model/c_head"
file="c_head.stl"
scale="1000 1000 1000"/>
...
<material
name="my_fly_model/headthorax"
texture="my_fly_model/headthorax"
specular="0.2"
shininess="0.2"/>
<texture
name="my_fly_model/headthorax"
builtin="flat"
rgb1="0.59 0.39 0.12"
rgb2="0.59 0.39 0.12"
mark="random"
markrgb="0.7 0.49 0.2"
random="0.3"
width="50"
height="300"/>
...
</asset>
The next section is very important, as it defines the kinematic tree of our model. The top-level element is <worldbody>, which encompasses all bodies in the model and the joints that connect them. In MuJoCo, body objects are abstract containers that hold "things" such as geometries, joints, and sites. The actual shape of a body is defined by its child geom elements, which can be primitive (spheres, cylinders, boxes, or cylinders with hemispherical ends, called "capsules") or mesh-based. In our case, body-segment meshes from the <asset> section are used for geometry. This section is very long and repetitive, so we include only a simplified, reordered snippet for demonstration.
The <worldbody> starts with elements that belong to the world itself: the ground_plane geom that the fly stands on, and a <site> marking the fly's spawn position. The fly is then attached as a child <body>. Two things to note about the fly:
- Namespacing. Every element belonging to the fly is prefixed with the fly's name, e.g.
my_fly_model/c_thorax. This is what lets a single world hold several flies without name clashes. For readability, we abbreviate the prefix as…/in the snippet below. - The root body and the free joint. The root of the fly's kinematic tree, the thorax, is a proper
<body>, and it is connected to the world by a free joint (<joint type="free">, named after the fly), which gives the fly 6 degrees of freedom (3 translational + 3 rotational). All the other body segments are nested inside the thorax.
<worldbody>
<!-- Elements belonging to the world: the ground plane and the fly's spawn site -->
<geom name="ground_plane" type="plane" contype="0" conaffinity="0" material="grid"/>
<site name="my_fly_model" pos="0 0 0.7"/>
<!-- The root of the fly's kinematic tree is the thorax. It is a <body> connected -->
<!-- to the world by a free joint, and it wraps the entire kinematic tree below. -->
<body name="…/c_thorax" pos="0.496 0 2">
<joint name="my_fly_model" type="free"/> <!-- 6-DoF connection to the world -->
<geom name="…/c_thorax" type="mesh" contype="0" conaffinity="0"
material="…/headthorax" mass="0.000307" mesh="…/c_thorax"/>
<camera name="…/trackcam" ... />
<!-- Extending from the thorax, we add body segments. The order in which -->
<!-- elements are defined does not matter, but the level of the element is -->
<!-- important (e.g., you can define the geom first and then the child bodies, -->
<!-- or vice versa). For example, here we define the thorax-head-proboscis -->
<!-- kinematic chain. -->
<body name="…/c_head" ... >
<geom name="…/c_head" ... />
<!-- Under the "head" body, we can add joints that connect it to its parent -->
<!-- in the kinematic tree (thorax). Note that in NeuroMechFly, all joints -->
<!-- within the fly body are "hinge" joints, which means they define rotations -->
<!-- along only one rotational axis. In other words, they correspond to the -->
<!-- JointDOF class in FlyGym, not AnatomicalJoint (which can contain >1 DoFs. -->
<!-- Here, the thorax-head joint is a 3-DoF joint, so we define 3 joints. -->
<!-- Joints can have biophysical properties. For example, their passive -->
<!-- biomechanics are approximated as a spring-damper system (passive as in -->
<!-- the spring-damper system can dynamically respond to forces even in the -->
<!-- absence of active actuators like motors or muscles). The "springref" -->
<!-- parameter defines the neutral position (omitted below when it is the -->
<!-- default of 0; MuJoCo only writes out non-default attributes). -->
<joint name="…/c_thorax-c_head-roll"
type="hinge" axis="0 0 1"
stiffness="10" damping="0.5" ... />
<joint name="…/c_thorax-c_head-pitch" type="hinge" axis="0 1 0"
stiffness="10" damping="0.5" ... />
<joint name="…/c_thorax-c_head-yaw" type="hinge" axis="1 0 0"
stiffness="10" damping="0.5" ... />
<!-- Now we have further body segments extending from the head -->
<body name="…/c_rostrum" ... > <!-- rostrum: proximal part of the proboscis -->
<geom name="…/c_rostrum" ... />
<!-- From here on, we omit joint type, axis, spring-damper parameters, etc. -->
<joint name="…/c_head-c_rostrum-roll" ... />
<joint name="…/c_head-c_rostrum-pitch" ... />
<joint name="…/c_head-c_rostrum-yaw" ... />
<!-- ... and so forth -->
<body name="…/c_haustellum" ... > <!-- distal part of the proboscis -->
<geom name="…/c_haustellum" ... />
<joint name="…/c_rostrum-c_haustellum-roll" ... />
<joint name="…/c_rostrum-c_haustellum-pitch" ... />
<joint name="…/c_rostrum-c_haustellum-yaw" ... />
</body>
</body>
<!-- The kinematic tree can branch out. For example, here we define the left -->
<!-- antenna. Note that the base link of the antenna (pedicel) is defined -->
<!-- immediately under the head (i.e., at the same level as the rostrum). -->
<!-- In other words, we have two branches in the kinematic tree: -->
<!-- thorax-head-rostrum-haustellum and thorax-head-l_pedicel-... -->
<body name="…/l_pedicel" ... >
<geom name="…/l_pedicel"... />
<joint name="…/c_head-l_pedicel-roll" ... />
<joint name="…/c_head-l_pedicel-pitch" ... />
<joint name="…/c_head-l_pedicel-yaw" ... />
<body name="…/l_funiculus" ... >
<geom name="…/l_funiculus" ... />
<joint name="…/l_pedicel-l_funiculus-roll" ... />
<joint name="…/l_pedicel-l_funiculus-pitch" ... />
<joint name="…/l_pedicel-l_funiculus-yaw" ... />
<body name="…/l_arista" ... >
<geom name="…/l_arista" ... />
<joint name="…/l_funiculus-l_arista-roll" ... />
<joint name="…/l_funiculus-l_arista-pitch" ... />
<joint name="…/l_funiculus-l_arista-yaw" ... />
</body>
</body>
</body>
</body>
<!-- Similarly, at the same level as the head, other kinematic chains extend -->
<!-- from the thorax. Here we have the left-front leg's coxa, which rotates -->
<!-- about all 3 axes of rotation at its joint with the thorax -->
<body name="…/lf_coxa" ... >
<geom name="…/lf_coxa" ... />
<joint name="…/c_thorax-lf_coxa-roll" ... />
<joint name="…/c_thorax-lf_coxa-pitch" ... />
<joint name="…/c_thorax-lf_coxa-yaw" ... />
<!-- The second link of the kinematic chain is the fused trochanter-femur -->
<!-- segment, which has 2 axes of rotation at its joint with the coxa -->
<body name="…/lf_trochanterfemur" ... >
<geom name="…/lf_trochanterfemur" ... />
<joint name="…/lf_coxa-lf_trochanterfemur-roll" ... />
<joint name="…/lf_coxa-lf_trochanterfemur-pitch" ... />
<!-- Now the tibia, which only has one axis of rotation relative to -->
<!-- trochanter-femur (pitch) -->
<body name="…/lf_tibia" ... >
<geom name="…/lf_tibia" ... />
<joint name="…/lf_trochanterfemur-lf_tibia-pitch" ... />
<!-- The tarsus is compliant (soft) and has 5 segments. In NeuroMechFly, -->
<!-- each tarsal segment is approximated as a rigid body, and their -->
<!-- connections are approximated as passive, non-rigid joints (i.e., with -->
<!-- stiffness and damping) -->
<body name="…/lf_tarsus1" ... >
<geom name="…/lf_tarsus1" ... />
<joint name="…/lf_tibia-lf_tarsus1-pitch" ...
stiffness="10" damping="0.5" springref="-0.17" ... />
<body name="…/lf_tarsus2" ... >
<geom name="…/lf_tarsus2" ... />
<joint name="…/lf_tarsus1-lf_tarsus2-pitch" ...
stiffness="10" damping="0.5" springref="-0.09" ... />
<body name="…/lf_tarsus3" ... >
<geom name="…/lf_tarsus3" ... />
<joint name="…/lf_tarsus2-lf_tarsus3-pitch" ...
stiffness="10" damping="0.5" springref="-0.09" ... />
<body name="…/lf_tarsus4" ... >
<geom name="…/lf_tarsus4" ... />
<joint name="…/lf_tarsus3-lf_tarsus4-pitch" ...
stiffness="10" damping="0.5" springref="-0.09" ... />
<body name="…/lf_tarsus5" ... >
<geom name="…/lf_tarsus5" ... />
<joint name="…/lf_tarsus4-lf_tarsus5-pitch" ...
stiffness="10" damping="0.5" springref="-0.09" ... />
</body>
</body>
</body>
</body>
</body>
</body>
</body>
</body>
<!-- Other legs are defined in parallel to the left front leg -->
<body name="…/lm_coxa" ... >
...
</body>
...
</body> <!-- end of the fly's thorax body -->
<!-- End of the model -->
</worldbody>
The next section defines actuators. Because the actuators act on the fly's joints, their names are namespaced under the fly's name as well:
<actuator>
<!-- MuJoCo represents every actuator internally as a "general" actuator. The -->
<!-- familiar shortcuts (position, velocity, motor, muscle, adhesion, ...) are just -->
<!-- convenient presets that expand into a general actuator with a particular -->
<!-- gain/bias parameterization. For example, the position actuators we requested -->
<!-- with `actuator_type="position", kp=50` expand into general actuators with -->
<!-- `biastype="affine"`, `gainprm="50 ..."`, and `biasprm="0 -50 ..."`. -->
<general
name="my_fly_model/c_thorax-lf_coxa-roll-position"
joint="my_fly_model/c_thorax-lf_coxa-roll"
forcelimited="true"
forcerange="-30 30"
biastype="affine"
gainprm="50"
biasprm="0 -50"/>
<general
name="my_fly_model/c_thorax-lf_coxa-pitch-position"
joint="my_fly_model/c_thorax-lf_coxa-pitch"
forcelimited="true"
forcerange="-30 30"
biastype="affine"
gainprm="50"
biasprm="0 -50"/>
...
</actuator>
There are additional sections that can be defined. For these, see the MuJoCo XML reference for details.
Now, you might have realized that flygym.compose is simply a tool to generate these MJCF specification strings. There is literally no physics simulation involved at all—just string manipulation and formatting that result in this text-only XML file (in fact, so is mujoco.MjSpec). If you want to customize certain things in ways that are not supported by flygym.compose? You can quite literally modify the appropriate parameters in the XML file after composing your model using flygym.compose.
However, we do not recommend that you literally edit the XML file directly. We have learned through our experience that custom edits to XML files are hard to reproduce and hard to document. In a year's time, you might have mysterious semi-manually edited XML files passed among collaborators that no one really understands how they were made. For this reason, we strongly recommend using mujoco.MjSpec to edit the model procedurally (i.e., with human-readable code).
As a concrete example, let's say you want to change the stiffness of all tarsal joints (tarsus1-2, 2-3, 3-4, 4-5) from 10 to 5. With flygym.compose, you can use a non-default stiffness value (fly.add_joints(stiffness=...)), but this would apply to all joints. Therefore, the best way to selectively modify tarsal joints is to use mujoco.MjSpec. We will briefly demonstrate how you can do this, but for all details about the mujoco.MjSpec API, please refer to the official model editing documentation.
The MJCF root element (an mujoco.MjSpec) is exposed through FlyGym as world.mjcf_root (and as fly.mjcf_root on Fly objects). Because our fly is attached to the world, the two point into the same merged specification; below we go through world.mjcf_root:
world.mjcf_root
<mujoco._specs.MjSpec at 0x769cb10976b0>
We can access its element trees, for example the worldbody:
world.mjcf_root.worldbody
<mujoco._specs.MjsBody at 0x769caef16430>
We can use .find_all(tag_name) to find all elements of that type under any element. In this case, let's find all joints:
all_joints = world.mjcf_root.worldbody.find_all("joint")
print(f"Found {len(all_joints)} joints in the worldbody (printing first 5):")
for joint in all_joints[:5]:
print(joint.name)
Found 127 joints in the worldbody (printing first 5): my_fly_model my_fly_model/c_thorax-c_head-roll my_fly_model/c_thorax-c_head-pitch my_fly_model/c_thorax-c_head-yaw my_fly_model/c_head-c_rostrum-roll
By filtering by name, we can find all tarsal joints. Recall from the previous tutorial that flygym.anatomy offers handy Python class representations like JointDOF. You can parse a MuJoCo <joint> name into a JointDOF object to make filtering easy.
One thing to keep in mind: now that the fly is attached to a world, all of its elements are namespaced under the fly's name. For example, the joint c_thorax-lf_coxa-roll becomes my_fly_model/c_thorax-lf_coxa-roll (you can see this in the output above), and the world adds a free joint (named after the fly, my_fly_model) that connects the fly to the ground. So before parsing a joint name with JointDOF, we strip the "<fly_name>/" prefix — and skip the free joint, which is not a JointDOF. JointDOF.from_name operates on the un-prefixed name:
from flygym.anatomy import JointDOF
dof = JointDOF.from_name("c_thorax-lf_coxa-roll")
print("DoF:", dof.name)
print("parent:", dof.parent)
print("child:", dof.child)
print("child.is_leg:", dof.child.is_leg())
print("child.pos:", dof.parent.pos)
print("child.link:", dof.child.link)
DoF: c_thorax-lf_coxa-roll parent: BodySegment(name='c_thorax') child: BodySegment(name='lf_coxa') child.is_leg: True child.pos: c child.link: coxa
With this, we can modify tarsal segments as follows:
prefix = f"{fly.name}/" # the fly's elements are namespaced under its name
for joint_element in world.mjcf_root.worldbody.find_all("joint"):
if not joint_element.name.startswith(prefix):
continue # e.g. the free joint that attaches the fly to the world
dof = JointDOF.from_name(joint_element.name.removeprefix(prefix))
if dof.child.link.startswith("tarsus") and dof.child.link != "tarsus1":
# As of MuJoCo 3.7, a joint's stiffness/damping are stored as polynomial
# coefficients; index 0 is the linear term (the scalar we want to change).
old_stiffness = joint_element.stiffness[0]
joint_element.stiffness[0] = 5
print(
f"Joint {joint_element.name}: "
f"changed stiffness from {old_stiffness} to {joint_element.stiffness[0]}"
)
Joint my_fly_model/lf_tarsus1-lf_tarsus2-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lf_tarsus2-lf_tarsus3-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lf_tarsus3-lf_tarsus4-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lf_tarsus4-lf_tarsus5-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lm_tarsus1-lm_tarsus2-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lm_tarsus2-lm_tarsus3-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lm_tarsus3-lm_tarsus4-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lm_tarsus4-lm_tarsus5-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lh_tarsus1-lh_tarsus2-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lh_tarsus2-lh_tarsus3-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lh_tarsus3-lh_tarsus4-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/lh_tarsus4-lh_tarsus5-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rf_tarsus1-rf_tarsus2-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rf_tarsus2-rf_tarsus3-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rf_tarsus3-rf_tarsus4-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rf_tarsus4-rf_tarsus5-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rm_tarsus1-rm_tarsus2-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rm_tarsus2-rm_tarsus3-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rm_tarsus3-rm_tarsus4-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rm_tarsus4-rm_tarsus5-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rh_tarsus1-rh_tarsus2-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rh_tarsus2-rh_tarsus3-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rh_tarsus3-rh_tarsus4-pitch: changed stiffness from 10.0 to 5.0 Joint my_fly_model/rh_tarsus4-rh_tarsus5-pitch: changed stiffness from 10.0 to 5.0
As you can see, these few lines of code are much more readable, version-trackable (e.g., with Git), and maintainable than someone changing these numbers manually in a 500-line XML file and passing it on to someone else.
We can also add new elements with dm_control.mjcf. Purely for the sake of demonstration, we will add a red sphere at the tibia-tarsus joint of every leg, and remove it specifically for the right front leg:
import mujoco as mj
from flygym.anatomy import BodySegment
for body_element in world.mjcf_root.worldbody.find_all("body"):
if not body_element.name.startswith(prefix):
continue # skip any non-fly bodies
body_segment = BodySegment(body_element.name.removeprefix(prefix))
if body_segment.link == "tarsus1":
leg_name = body_segment.pos
# Add a sphere under the tarsus1 body at position (0, 0, 0), which, by
# definition, is the joint with its parent (tibia). Note that MjSpec uses
# typed `add_<element>` methods (e.g. `add_geom`, `add_body`, `add_joint`)
# and expects enum values for fields like the geom type.
sphere_geom = body_element.add_geom(
name=f"visualization_sphere_{leg_name}",
type=mj.mjtGeom.mjGEOM_SPHERE,
size=[0.1],
rgba=[1, 0, 0, 1],
pos=[0, 0, 0],
density=0, # important! Make the sphere massless and visualization-only
)
print(f"Added sphere {sphere_geom.name} to body {body_element.name}")
Added sphere visualization_sphere_lf to body my_fly_model/lf_tarsus1 Added sphere visualization_sphere_lm to body my_fly_model/lm_tarsus1 Added sphere visualization_sphere_lh to body my_fly_model/lh_tarsus1 Added sphere visualization_sphere_rf to body my_fly_model/rf_tarsus1 Added sphere visualization_sphere_rm to body my_fly_model/rm_tarsus1 Added sphere visualization_sphere_rh to body my_fly_model/rh_tarsus1
Now, we can find the sphere for the right front leg by name. Element lookups are done on the spec with spec.geom(name) (and likewise spec.body(name), spec.joint(name), etc.), which returns the element or None if it does not exist:
world.mjcf_root.geom("visualization_sphere_rf")
<mujoco._specs.MjsGeom at 0x769cb10fa9b0>
... and remove it with spec.delete(element):
world.mjcf_root.delete(world.mjcf_root.geom("visualization_sphere_rf"))
Now, if we visualize the fly, we will notice the new sphere objects (except for the right front leg).
from flygym.rendering import preview_model
mj_model, mj_data = world.compile()
preview_model(mj_model, mj_data, camera, show_in_notebook=True)
my_fly_model/trackcam |
Modifying physics parameters¶
Beyond the body kinematic tree, you can also modify global simulation parameters by editing attributes in the appropriate sections. These global settings (e.g. the <option> and <visual> sections) belong to the world once the fly is attached, so we edit them on world.mjcf_root. For example, to change the simulation timestep from 0.1 ms to 1 ms (which makes the simulation 10x faster but less stable):
print("Old timestep:", world.mjcf_root.option.timestep)
world.mjcf_root.option.timestep = 0.001
print("New timestep:", world.mjcf_root.option.timestep)
Old timestep: 0.0001 New timestep: 0.001
Or, to make background lighting less bright:
print("Old headlight ambience:", world.mjcf_root.visual.headlight.ambient)
world.mjcf_root.visual.headlight.ambient = [0.1, 0.1, 0.1] # in RGB
print("New headlight ambience:", world.mjcf_root.visual.headlight.ambient)
Old headlight ambience: [0.5 0.5 0.5] New headlight ambience: [0.1 0.1 0.1]
mj_model, mj_data = world.compile()
preview_model(mj_model, mj_data, camera, show_in_notebook=True)
my_fly_model/trackcam |
As we would expect, the physics simulation becomes much less stable now that the timestep is increased by a factor of 10.
Where does flygym.compose get its data from?¶
Defaults used by flygym.compose are defined in these files:
from flygym import assets_dir
model_dir = (assets_dir / "model").relative_to(Path.cwd(), walk_up=True)
for path in sorted(model_dir.glob("*")):
print(path)
../src/flygym/assets/model/flybody ../src/flygym/assets/model/musculoskeletal ../src/flygym/assets/model/neuromechfly
Here,
- The
legacy/folder can be ignored. - The
meshes/folder contains geometric meshes for body parts. mujoco_globals.yamldefines parameters in the<compiler>,<option>,<size>and<visual>sections. Anything left unset defaults to its MuJoCo default.- The
pose/folder contains built-in kinematic poses (e.g., thepose/neutral/subfolder contains joint angles for the default pose in different joint axis orders). rigging.yamldefines the kinematic tree (which body parts are connected by which joints).visuals.yamldefines the appearance of different body parts.
For reasons similar to our arguments against manually modifying XML files, we recommend that you don't modify these files yourself and use mujoco.MjSpec to change them instead. An exception is that you might add additional kinematic poses under the pose/ folder.