Refactor musculoskeletal (FlyMimic) model onto the procedural composition workflow

[sibocw]

Summary

Refactor the musculoskeletal (FlyMimic) body model (#270) so it is built through the same procedural composition workflow as NeuroMechFly/FlyBody, instead of loading a self-contained, pre-authored MJCF and wrapping it.

Concretely:

• Make MusculoskeletalFly a first-class BaseFly subclass, built mesh-by-mesh from extracted config files (rigging.yaml, mujoco_globals.yaml, etc.), exactly like NeuroMechFly.
• Add a new BaseFly.add_muscle() composition step that attaches the sites, spatial tendons, Hill-type muscle actuators, passive joint properties, and joint-lock equality constraints.
• Replace the fly-specific MusculoskeletalWorld (which currently bundles the fly and the scene) with a slim environment-only world that provides just the tethered/in-air mount + floor + lighting from the FlyMimic MJCF, into which the fly is attached like any other fly.

Motivation

Today the muscle model is the odd one out. NeuroMechFly and FlyBody both compose a body from meshes + YAML rigging via BaseFly and attach into composable worlds (FlatGroundWorld, terrain, etc.). MusculoskeletalFly instead subclasses BaseCompositionElement directly, loads a self-contained MJCF (which ships its own floor/light/camera), and is paired with a bespoke MusculoskeletalWorld that adopts the fly's root rather than attaching the fly into a scene.

This split has real costs:

• Two divergent code paths for "a fly," each needing separate maintenance, tests, and docs.
• A world that exists only to host one fly model — backwards relative to FlyGym's compose model, where worlds are scenes and flies attach into them.
• FlyGym features don't apply uniformly to the muscle fly: vision, colorize, tracking cameras, terrain, contact-sensor plumbing, etc. all have to be special-cased or are simply unavailable.

Unifying onto the procedural workflow gives one consistent API/workflow, lets the muscle fly reuse add_vision()/colorize()/worlds for free, and makes muscles a reusable composition primitive (add_muscle()) rather than something frozen into a static XML.

Background investigation (already done)

The bodies are not interchangeable with NeuroMechFly

FlyMimic is the same fly lineage but a distinct rig — you cannot overlay muscles onto a procedurally-built NeuroMechFly body:

• Leg topology differs. FlyMimic splits the front legs into separate LFTrochanter + LFFemur bodies (needed to actuate the trochanter joint); NeuroMechFly fuses them (lf_trochanterfemur).
• The 84 muscle-attachment sites + 15 spatial tendons are authored in FlyMimic's own body frames against that split topology.
• FlyMimic ships explicit per-body inertials (71 of them) and per-joint passive params (springref, custom ranges, stiffness=0.4, damping=0.02), whereas the standard BaseFly build lets MuJoCo derive inertia from geoms and applies uniform joint stiffness/damping.

Therefore add_muscle() belongs on BaseFly as a general capability, but the bundled muscle config only fits a fly whose body is the FlyMimic rig — "just add add_muscle to NeuroMechFly" cannot be literal.

Mesh investigation (FlyMimic meshes/stl/ vs NeuroMechFly meshes/fullsize/)

Body segments fall into three categories:

A. Identical, shared verbatim — Thorax and abdomen A1A2…A6: same vertex counts, Kabsch residual 0.00000 mm, ~0° rotation, ~0 translation (the differing md5 is just re-export/normals).

B. Same geometry, but a rigid transform baked into the vertices (origin/orientation redefined and compensated by the body pos/quat):

segment	baked rotation	baked translation	Kabsch residual
Rostrum	92.7°	[0.15, 0.01, -0.21] mm	0.0018 mm
Haustellum	32.3°	present	0.0010 mm

NMF keeps these meshes canonical and carries the pose in the MJCF body transform (NMF rostrum/haustellum body quat = 0°); FlyMimic bakes the transform into the mesh-local frame. This is the key "mesh origin treated differently" gotcha.

C. Entirely different meshes (resolution and/or topology) — Head (FM 51,882 vs NMF 252,018 verts), eyes, antennae, and all legs. The muscle-bearing left-front leg is ultra-high-res (LFFemur 50,532 verts) while other legs are coarse.

Mirroring: NMF ships left meshes only and builds the right side via a negative-Y mesh scale scale=(S, −S, S). FlyMimic ships explicit, independent L/R meshes that are not mirrors of each other (e.g. LFFemur 50,532 vs RFFemur 3,645 verts). The geom-origin convention is otherwise identical in both (mesh-origin at body-origin; only the floor geom carries a pos offset).

Consequence: ship FlyMimic's own complete mesh set, disable mirroring (mirror_left2right=False), and take body pos/quat verbatim from the MJCF. Do not attempt to reuse/deduplicate against NMF meshes — only thorax/abdomen would line up; rostrum/haustellum would be mis-posed and head/eyes/legs are different meshes entirely. This is precisely why the "full re-extraction" approach (below) is the right one.

Other confirmed facts

• The fly is tethered: the thorax is fixed to the world (no free joint); qpos is 14 leg DOFs. Maps cleanly onto an in-air/tethered environment world.
• There are exactly 15 actuators, all muscles (the motor default class is declared but unused).
• 7 right-front-leg joints are locked via <equality><joint ... polycoef="0 0 0 0 0">.
• default-pose keyframe: qpos="0.1745 0.4357 0.2614 -0.4081 -2.388 0.1702 1.862 0 0 0 0 0 0 0".

Proposed plan

Sequenced so the body composition is proven equivalent before touching world/demo integration.

Step 1 — One-time extraction script (scripts/dev/extract_musculoskeletal_configs.py)

Parse assets/model/musculoskeletal/best_combined_arm_damping_stiff_cvt3.xml → emit into assets/model/musculoskeletal/:

• rigging.yaml — per body: pos, quat, mass, plus explicit inertial (ipos, iquat, diaginertia). Must reference FlyMimic's own meshes; mirroring disabled.
• mujoco_globals.yaml — timestep=1e-4, gravity, angle=radian, size (njmax/nconmax/nuser_jnt), autolimits.
• muscle.yaml — 84 sites (parent body + pos), 15 spatial tendons (ordered site waypoints + width/rgba), 15 Hill actuators (tendon, ctrlrange, lengthrange, gainprm/biasprm/dynprm), the muscle default class, the 7 RF-leg equality locks, and per-joint passive params (springref, range, stiffness, damping, armature).
• skeleton.yaml (or equivalent) — kinematic tree (parent→child, joint name, axis) for rebuilding joints (FlyMimic topology ≠ NMF, e.g. split front-leg trochanter).
• Keep meshes in meshes/stl/.
• Preserve MJCF element order for bodies/joints/actuators (see Step 6 risk).

Step 2 — Extend BaseFly for explicit inertials

_add_one_body_and_geoms currently sets mass only. Add optional honoring of ipos/iquat/diaginertia from rigging.yaml. NeuroMechFly's config omits these → behavior unchanged.

Step 3 — Add BaseFly.add_muscle(muscle_config_path)

New composition step (peer of add_joints/add_actuators/add_leg_adhesion): emits sites, spatial tendons, the 15 muscle general actuators, the equality locks, and applies per-joint passive params. Populates jointdof_to_mjcfactuator_by_type[MUSCLE] exactly as today so the imitation env (which finds muscles by name in the compiled model) is unaffected.

Step 4 — Joint reconciliation (main design risk)

FlyMimic joints carry their own springref/range and a non-NMF topology, so they don't fit add_joints(skeleton, neutral_pose, stiffness=, damping=) cleanly. Either generalize add_joints to accept per-joint passive specs + a data-driven (non-anatomy.py) skeleton, or add a dedicated joint-building path that consumes the extracted skeleton.yaml/muscle.yaml directly. Recommendation: a small data-driven skeleton builder, to avoid bending NMF's anatomy classes around FlyMimic's body plan.

Step 5 — Rewrite MusculoskeletalFly + replace the world

• MusculoskeletalFly(BaseFly) with default paths to the new assets; __init__ → meshes → bodies/geoms → joints → add_muscle(). Keep muscle_names, eye-body add_vision.
• Replace MusculoskeletalWorld with an environment-only world (e.g. an in-air/tethered scene providing FlyMimic's mount + floor + lighting), into which the fly is attached via the normal add_fly() path — not a world that adopts the fly's root. Drop FlyMimic's floor/light/world-camera/skybox from the fly itself (the world supplies them).

Step 6 — Update the demo (flygym_demo/muscle_imitation/)

fly.py → standard composition (MusculoskeletalFly + environment world + Simulation); update env.py/record.py/train.py/__main__.py imports.
Risk: env.py, the mocap clips in data.py, and existing checkpoints under runs/ are keyed to joint/actuator/body names and order. The extractor must preserve MJCF order exactly, or clips/checkpoints break.

Step 7 — Equivalence validation (de-risks dynamics drift)

A test that compiles both the new procedural model and the original MJCF and asserts matching nbody/njnt/nu/ntendon/nsite/neq, body masses + inertials (within tol), joint ranges/springref, actuator gain/bias params, and tendon site routing. This is the objective green/red signal that the refactor is behavior-preserving — build it early, behind Steps 1–4, before integration.

Step 8 — GPU path, tests, docs

Rebuild check_mjwarp_compatibility / build_musculoskeletal_gpu_simulation on the new composition (or fold into generic GPUSimulation(world)); update tests/examples/test_muscle_imitation.py + add a MusculoskeletalFly composition test; refresh docs/tutorials/6_muscle_imitation.md and docstrings.

Biggest risks (surfaced)

1. Joint/skeleton reconciliation (Step 4) — FlyMimic's topology + per-joint passive params don't fit the anatomy-driven add_joints cleanly.
2. Order preservation (Step 6) — qpos/actuator/body ordering must match the original MJCF or trained checkpoints and mocap clips silently break.

Both are mitigated by an order-preserving extractor (Step 1) + the equivalence test (Step 7).

Relevant files

• src/flygym/compose/fly/musculoskeletal.py — current MusculoskeletalFly (to be rewritten as a BaseFly subclass)
• src/flygym/compose/fly/base_fly.py — BaseFly (add add_muscle(), explicit inertials)
• src/flygym/compose/fly/neuromechfly.py — reference for the procedural subclass pattern
• src/flygym/compose/world/musculoskeletal.py — MusculoskeletalWorld (to be replaced by an environment-only world)
• src/flygym/assets/model/musculoskeletal/best_combined_arm_damping_stiff_cvt3.xml — extraction source
• src/flygym_demo/muscle_imitation/{fly,env,data,record,train,__main__}.py — demo to update
• tests/examples/test_muscle_imitation.py, docs/tutorials/6_muscle_imitation.md

---

Investigation scripts used to derive the mesh findings are throwaway (run against the assets above with trimesh); they can be reproduced from the tables in this issue if needed.

[sibocw]

Deferred: not implemented in v2.1.0