⚠️ 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.
⚠️ FlyBody integration is experimental: The API may change in future releases.
In [1]:
Copied!
# --- 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"
# --- 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"
Using the FlyBody model¶
In this tutorial, we use the FlyBody body model — published in Vaxenburg et al. (2025) (code) — instead of the default NeuroMechFly model, while retaining the FlyGym API for scene composition, simulation, and control.
We run the same hybrid turning controller from Tutorial 4d on flat terrain. A two-value descending signal modulates the left/right CPG amplitudes and directions; the hybrid correction rules then produce low-level joint angle and adhesion commands.
One FlyBody-specific detail arises from the DOF ordering. HybridController builds its joint-angle output using plain BodySegment/RotationAxis/JointDOF types from flygym.anatomy, while fly.get_actuated_jointdofs_order returns FlybodyJointDOF instances. Because frozen dataclasses compare equal only within the same concrete class, we convert the DOF order to plain types before passing it to the controller — the index order is preserved, so the resulting joint-angle array still maps 1-for-1 onto the FlyBody position actuators.
!!! warning "Experimental" Support for the FlyBody body model is experimental. The API may change in future releases, and not all features available for the default NeuroMechFly model are currently supported.
In [2]:
Copied!
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
from tqdm import trange
from flygym import Simulation
from flygym.anatomy import BodySegment, JointDOF, RotationAxis
from flygym.compose import ActuatorType, FlatGroundWorld, KinematicPosePreset
from flygym.compose.fly import FlyBody
from flygym.flybody import (
FlyBodyActuatedDOFPreset,
FlyBodyAxisOrder,
FlyBodyContactBodiesPreset,
FlyBodyJointPreset,
FlyBodySkeleton,
FlyBodyBodySegment,
)
from flygym.utils.math import Rotation3D
from flygym_demo.complex_terrain import (
HybridControllerObservation,
HybridTurningController,
FlyBodyPreprogrammedSteps,
apply_locomotion_action,
)
from flygym_demo.complex_terrain.common import LocomotionAction
output_dir = Path("demo_output/flybody_turning_controller")
output_dir.mkdir(parents=True, exist_ok=True)
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
from tqdm import trange
from flygym import Simulation
from flygym.anatomy import BodySegment, JointDOF, RotationAxis
from flygym.compose import ActuatorType, FlatGroundWorld, KinematicPosePreset
from flygym.compose.fly import FlyBody
from flygym.flybody import (
FlyBodyActuatedDOFPreset,
FlyBodyAxisOrder,
FlyBodyContactBodiesPreset,
FlyBodyJointPreset,
FlyBodySkeleton,
FlyBodyBodySegment,
)
from flygym.utils.math import Rotation3D
from flygym_demo.complex_terrain import (
HybridControllerObservation,
HybridTurningController,
FlyBodyPreprogrammedSteps,
apply_locomotion_action,
)
from flygym_demo.complex_terrain.common import LocomotionAction
output_dir = Path("demo_output/flybody_turning_controller")
output_dir.mkdir(parents=True, exist_ok=True)
Build the FlyBody fly¶
Legs-only joints, position-controlled active leg DoFs, and leg adhesion — same active DOF layout (42 position actuators, 6 adhesion outputs) as the standard locomotion fly, but using the FlyBody skeleton, axis convention, and neutral pose.
In [3]:
Copied!
axis_order = FlyBodyAxisOrder.YAW_ROLL_PITCH
neutral_pose = KinematicPosePreset.FLYBODY_NEUTRAL
fly = FlyBody(name="flybody")
skeleton = FlyBodySkeleton(
axis_order=axis_order,
joint_preset=FlyBodyJointPreset.LEGS_ONLY,
)
fly.add_joints(skeleton, neutral_pose)
actuated_dofs = fly.skeleton.get_actuated_dofs_from_preset(
FlyBodyActuatedDOFPreset.LEGS_ACTIVE_ONLY
)
fly.add_actuators(actuated_dofs, ActuatorType.POSITION, kp=100)
fly.add_leg_adhesion()
fly.colorize()
body_cam = fly.add_tracking_camera(
name="body_cam",
pos_offset=(0.0, -8.0, 0.0),
rotation=Rotation3D("euler", (1.57, 0.0, 0.0)),
fovy=35.0,
)
world = FlatGroundWorld()
world.add_fly(
fly,
[0, 0, 0.7],
Rotation3D("quat", [1, 0, 0, 0]),
bodysegs_with_ground_contact=FlyBodyContactBodiesPreset.TIBIA_TARSUS_ONLY,
add_ground_contact_sensors=False,
)
sim = Simulation(world)
renderer = sim.set_renderer(
[body_cam],
camera_res=(240, 320),
playback_speed=0.1,
output_fps=25,
)
axis_order = FlyBodyAxisOrder.YAW_ROLL_PITCH
neutral_pose = KinematicPosePreset.FLYBODY_NEUTRAL
fly = FlyBody(name="flybody")
skeleton = FlyBodySkeleton(
axis_order=axis_order,
joint_preset=FlyBodyJointPreset.LEGS_ONLY,
)
fly.add_joints(skeleton, neutral_pose)
actuated_dofs = fly.skeleton.get_actuated_dofs_from_preset(
FlyBodyActuatedDOFPreset.LEGS_ACTIVE_ONLY
)
fly.add_actuators(actuated_dofs, ActuatorType.POSITION, kp=100)
fly.add_leg_adhesion()
fly.colorize()
body_cam = fly.add_tracking_camera(
name="body_cam",
pos_offset=(0.0, -8.0, 0.0),
rotation=Rotation3D("euler", (1.57, 0.0, 0.0)),
fovy=35.0,
)
world = FlatGroundWorld()
world.add_fly(
fly,
[0, 0, 0.7],
Rotation3D("quat", [1, 0, 0, 0]),
bodysegs_with_ground_contact=FlyBodyContactBodiesPreset.TIBIA_TARSUS_ONLY,
add_ground_contact_sensors=False,
)
sim = Simulation(world)
renderer = sim.set_renderer(
[body_cam],
camera_res=(240, 320),
playback_speed=0.1,
output_fps=25,
)
Build the controller — converting the DOF order¶
fly.get_actuated_jointdofs_order(POSITION) returns FlyBody-typed FlybodyJointDOF objects. The hybrid controller builds an internal dict[JointDOF, angle] using plain JointDOF(BodySegment, BodySegment, RotationAxis) keys, so we rebuild the order with the plain types. Because we only rebuild the element types (not the order), the resulting joint-angle array still lines up with the FlyBody position actuators index-by-index.
In [4]:
Copied!
def to_plain_dof_order(flybody_dof_order):
"""Cast a list of FlyBodyJointDOF into plain JointDOF, preserving order."""
return [
JointDOF(
BodySegment(d.parent.name),
BodySegment(d.child.name),
RotationAxis(d.axis.value),
)
for d in flybody_dof_order
]
flybody_dof_order = fly.get_actuated_jointdofs_order(ActuatorType.POSITION)
controller_dof_order = to_plain_dof_order(flybody_dof_order)
preprogrammed_steps = FlyBodyPreprogrammedSteps()
controller = HybridTurningController(
timestep=sim.timestep,
preprogrammed_steps=preprogrammed_steps,
output_dof_order=controller_dof_order,
enable_adhesion=False, # disable adhesion as somehow it leads to the fly rearing ???
)
print("Position actuators:", len(flybody_dof_order))
print("Simulation timestep:", sim.timestep)
def to_plain_dof_order(flybody_dof_order):
"""Cast a list of FlyBodyJointDOF into plain JointDOF, preserving order."""
return [
JointDOF(
BodySegment(d.parent.name),
BodySegment(d.child.name),
RotationAxis(d.axis.value),
)
for d in flybody_dof_order
]
flybody_dof_order = fly.get_actuated_jointdofs_order(ActuatorType.POSITION)
controller_dof_order = to_plain_dof_order(flybody_dof_order)
preprogrammed_steps = FlyBodyPreprogrammedSteps()
controller = HybridTurningController(
timestep=sim.timestep,
preprogrammed_steps=preprogrammed_steps,
output_dof_order=controller_dof_order,
enable_adhesion=False, # disable adhesion as somehow it leads to the fly rearing ???
)
print("Position actuators:", len(flybody_dof_order))
print("Simulation timestep:", sim.timestep)
Position actuators: 42 Simulation timestep: 0.0001
Apply the initial pose¶
In [5]:
Copied!
sim.reset()
controller.reset(seed=0)
initial_action = LocomotionAction(
joint_angles=preprogrammed_steps.default_pose_by_dof_order(controller_dof_order),
adhesion_onoff=np.ones(6, dtype=bool),
)
apply_locomotion_action(sim, fly.name, initial_action)
sim.warmup()
sim.reset()
controller.reset(seed=0)
initial_action = LocomotionAction(
joint_angles=preprogrammed_steps.default_pose_by_dof_order(controller_dof_order),
adhesion_onoff=np.ones(6, dtype=bool),
)
apply_locomotion_action(sim, fly.name, initial_action)
sim.warmup()
Run the turning controller¶
Descending command turns one way for the first half of the run, then the other way.
In [6]:
Copied!
run_time = 2.0
nsteps_sim = int(run_time / sim.timestep)
time_grid = np.arange(nsteps_sim) * sim.timestep
thorax_idx = fly.get_bodysegs_order().index(FlyBodyBodySegment("c_thorax"))
thorax_positions = np.full((nsteps_sim, 3), np.nan, dtype=np.float32)
cpg_magnitudes = np.full((nsteps_sim, 6), np.nan, dtype=np.float32)
descending_signals = np.full((nsteps_sim, 2), np.nan, dtype=np.float32)
for step_idx in trange(nsteps_sim, desc="Running flybody turning controller"):
curr_time = step_idx * sim.timestep
if curr_time < run_time / 2:
descending_signal = np.array([0.4, 1.2])
else:
descending_signal = np.array([1.2, 0.4])
obs = HybridControllerObservation.from_sim(sim, fly.name)
action = controller.step(descending_signal, obs)
apply_locomotion_action(sim, fly.name, action)
sim.step_with_profile()
thorax_positions[step_idx] = sim.get_body_positions(fly.name)[thorax_idx]
cpg_magnitudes[step_idx] = controller.cpg_network.curr_magnitudes
descending_signals[step_idx] = descending_signal
sim.render_as_needed_with_profile()
print("Rendered frames:", len(renderer.frames[body_cam.name]))
print(
"Final thorax displacement:",
np.round(thorax_positions[-1] - thorax_positions[0], 3),
"mm",
)
sim.print_performance_report()
run_time = 2.0
nsteps_sim = int(run_time / sim.timestep)
time_grid = np.arange(nsteps_sim) * sim.timestep
thorax_idx = fly.get_bodysegs_order().index(FlyBodyBodySegment("c_thorax"))
thorax_positions = np.full((nsteps_sim, 3), np.nan, dtype=np.float32)
cpg_magnitudes = np.full((nsteps_sim, 6), np.nan, dtype=np.float32)
descending_signals = np.full((nsteps_sim, 2), np.nan, dtype=np.float32)
for step_idx in trange(nsteps_sim, desc="Running flybody turning controller"):
curr_time = step_idx * sim.timestep
if curr_time < run_time / 2:
descending_signal = np.array([0.4, 1.2])
else:
descending_signal = np.array([1.2, 0.4])
obs = HybridControllerObservation.from_sim(sim, fly.name)
action = controller.step(descending_signal, obs)
apply_locomotion_action(sim, fly.name, action)
sim.step_with_profile()
thorax_positions[step_idx] = sim.get_body_positions(fly.name)[thorax_idx]
cpg_magnitudes[step_idx] = controller.cpg_network.curr_magnitudes
descending_signals[step_idx] = descending_signal
sim.render_as_needed_with_profile()
print("Rendered frames:", len(renderer.frames[body_cam.name]))
print(
"Final thorax displacement:",
np.round(thorax_positions[-1] - thorax_positions[0], 3),
"mm",
)
sim.print_performance_report()
Running flybody turning controller: 100%|█| 20000/20000 [00:13<00:00, 1430.65i
Rendered frames: 500 Final thorax displacement: [-1.8000e-02 1.8226e+01 -9.8000e-02] mm PERFORMANCE PROFILE
| Stage | Time/step (us) | Percent (%) | Throughput (iters/s) | Throughput x realtime |
|---|---|---|---|---|
| Physics simulation advancement | 107 | 78 | 9316 | 0.93 |
| Rendering* | 30 | 22 | 33550 | 3.36 |
| TOTAL | 137 | 100 | 7292 | 0.73 |
* Note: 500 frames were rendered out of 20000 steps. Therefore, rendering time per image is 1192 us.
In [7]:
Copied!
sim.renderer.show_in_notebook()
sim.renderer.save_video(output_dir / "flybody_turning_controller.mp4")
sim.renderer.show_in_notebook()
sim.renderer.save_video(output_dir / "flybody_turning_controller.mp4")
flybody/body_cam |
In [8]:
Copied!
fig, axes = plt.subplots(
3, 1, figsize=(5, 7), tight_layout=True, height_ratios=[2, 1, 1]
)
axes[0].plot(
thorax_positions[:, 0] - thorax_positions[0, 0],
thorax_positions[:, 1] - thorax_positions[0, 1],
)
axes[0].set_aspect("equal", adjustable="box")
axes[0].set_xlabel("Forward displacement (mm)")
axes[0].set_ylabel("Lateral displacement (mm)")
axes[0].set_title("Thorax path")
axes[1].plot(time_grid, descending_signals[:, 0], label="Left")
axes[1].plot(time_grid, descending_signals[:, 1], label="Right")
axes[1].set_ylabel("Signal")
axes[1].set_title("Descending command")
axes[1].legend()
for leg_idx, leg in enumerate(preprogrammed_steps.legs):
axes[2].plot(time_grid, cpg_magnitudes[:, leg_idx], label=leg.upper())
axes[2].set_xlabel("Time (s)")
axes[2].set_ylabel("Magnitude")
axes[2].set_title("CPG magnitudes")
axes[2].legend(ncols=3, fontsize=8)
fig.savefig(output_dir / "flybody_turning_controller_diagnostics.png")
fig, axes = plt.subplots(
3, 1, figsize=(5, 7), tight_layout=True, height_ratios=[2, 1, 1]
)
axes[0].plot(
thorax_positions[:, 0] - thorax_positions[0, 0],
thorax_positions[:, 1] - thorax_positions[0, 1],
)
axes[0].set_aspect("equal", adjustable="box")
axes[0].set_xlabel("Forward displacement (mm)")
axes[0].set_ylabel("Lateral displacement (mm)")
axes[0].set_title("Thorax path")
axes[1].plot(time_grid, descending_signals[:, 0], label="Left")
axes[1].plot(time_grid, descending_signals[:, 1], label="Right")
axes[1].set_ylabel("Signal")
axes[1].set_title("Descending command")
axes[1].legend()
for leg_idx, leg in enumerate(preprogrammed_steps.legs):
axes[2].plot(time_grid, cpg_magnitudes[:, leg_idx], label=leg.upper())
axes[2].set_xlabel("Time (s)")
axes[2].set_ylabel("Magnitude")
axes[2].set_title("CPG magnitudes")
axes[2].legend(ncols=3, fontsize=8)
fig.savefig(output_dir / "flybody_turning_controller_diagnostics.png")
