mirror of
https://github.com/rndlabsoy/fly-brain-full.git
synced 2026-09-10 11:20:43 +02:00
erojasoficial-byte
05974fefe0
- fly_embodied.py: fly_name was only assigned inside the `if args.flight` block but used unconditionally by the proboscis joint lookup. Added fly_name = fly.name before the proboscis section. - benchmark.py: replaced Unicode checkmark/cross (U+2713/U+2717) with ASCII "OK"/"FAIL" to prevent encoding crash on Windows consoles.
1037 lines
46 KiB
Python
1037 lines
46 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Embodied Drosophila: Brain-body closed loop with interactive 3D viewer.
|
|
|
|
Connects the fly-brain connectome simulation (LIF neurons on GPU)
|
|
to a NeuroMechFly v2 biomechanical body (flygym + MuJoCo). Behaviors emerge
|
|
from spike propagation through the real connectome — no hand-coded rules.
|
|
|
|
Usage:
|
|
python fly_embodied.py # Auto-demo: cycles through stimuli
|
|
python fly_embodied.py --stimulus p9 # Start with P9 forward walking
|
|
python fly_embodied.py --no-auto # Manual only (keyboard)
|
|
python fly_embodied.py --visual # REAL VISION: compound eye → connectome
|
|
|
|
Keys (in MuJoCo viewer window):
|
|
1 = Sugar GRNs -> forward walking (via downstream P9/MN9)
|
|
2 = P9 direct -> forward walking
|
|
3 = LC4 looming -> escape response (via Giant Fiber)
|
|
4 = JO touch -> grooming (via aDN1)
|
|
5 = Bitter GRNs -> aversion
|
|
6 = Or56a olfactory -> repulsion
|
|
0 = No stimulus -> spontaneous basal activity
|
|
SPACE = Toggle auto-demo on/off
|
|
"""
|
|
|
|
import sys
|
|
import argparse
|
|
import numpy as np
|
|
import mujoco
|
|
import mujoco.viewer
|
|
|
|
from flygym import Fly
|
|
from flygym.simulation import SingleFlySimulation
|
|
from flygym.examples.locomotion import PreprogrammedSteps
|
|
from flygym.examples.locomotion.turning_controller import HybridTurningController
|
|
|
|
from brain_body_bridge import (
|
|
BrainEngine, DNRateDecoder, BrainBodyBridge, STIMULI, DN_GROUPS,
|
|
)
|
|
from visual_system import VisualSystem
|
|
from looming_arena import LoomingArena
|
|
from brain_monitor import BrainMonitorProcess
|
|
from somatosensory import SomatosensorySystem, VibrationSource
|
|
from gustatory import GustatorySystem, TasteZone
|
|
from olfactory import OlfactorySystem, OdorSource
|
|
from vocalization import WingSongSystem
|
|
from flight import FlightSystem, FlightState
|
|
|
|
try:
|
|
from consciousness import ConsciousnessDetector
|
|
except ImportError:
|
|
ConsciousnessDetector = None
|
|
|
|
|
|
# ============================================================================
|
|
# Auto-demo sequence: cycles through stimuli so the fly is always active
|
|
# ============================================================================
|
|
|
|
AUTO_DEMO_SEQUENCE = [
|
|
# (stimulus_name, duration_seconds, description)
|
|
('p9', 4.0, 'Forward walking (P9 neurons)'),
|
|
('lc4', 2.0, 'ESCAPE! (LC4 looming -> Giant Fiber)'),
|
|
(None, 2.0, 'Recovery (no stimulus)'),
|
|
('sugar', 4.0, 'Sugar detected (feeding approach)'),
|
|
(None, 1.5, 'Pause'),
|
|
('jo', 4.0, 'Antennal touch (JO -> grooming)'),
|
|
(None, 2.0, 'Recovery'),
|
|
('p9', 3.0, 'Walking again (P9)'),
|
|
('lc4', 1.5, 'ESCAPE! (looming threat)'),
|
|
('p9', 3.0, 'Resume walking (P9)'),
|
|
('bitter', 3.0, 'Bitter taste (aversion)'),
|
|
(None, 2.0, 'Pause'),
|
|
('or56a', 3.0, 'Bad smell (Or56a olfactory)'),
|
|
(None, 1.5, 'Recovery'),
|
|
]
|
|
|
|
|
|
# ============================================================================
|
|
# Grooming Controller
|
|
# ============================================================================
|
|
|
|
class GroomingController:
|
|
"""Generates front-leg oscillation for antennal grooming behavior."""
|
|
|
|
def __init__(self, preprogrammed_steps=None, freq_hz=4.0):
|
|
self.steps = preprogrammed_steps or PreprogrammedSteps()
|
|
self.freq = freq_hz
|
|
self.neutral = np.zeros(42)
|
|
for i, leg in enumerate(self.steps.legs):
|
|
self.neutral[i * 7:(i + 1) * 7] = self.steps.get_joint_angles(
|
|
leg, np.pi, 0.0
|
|
)
|
|
|
|
def get_action(self, time_s):
|
|
"""Return joints+adhesion action dict for grooming at given time."""
|
|
joints = self.neutral.copy()
|
|
phase = 2 * np.pi * self.freq * time_s
|
|
femur_offset = 0.3 * np.sin(phase)
|
|
tibia_offset = 0.4 * np.sin(phase + np.pi / 2)
|
|
for base in (0, 21): # LF and RF
|
|
joints[base + 3] += femur_offset
|
|
joints[base + 5] += tibia_offset
|
|
adhesion = np.array([0, 1, 1, 0, 1, 1])
|
|
return {"joints": joints, "adhesion": adhesion}
|
|
|
|
|
|
# ============================================================================
|
|
# Main Simulation
|
|
# ============================================================================
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Embodied Drosophila')
|
|
parser.add_argument('--no-viewer', action='store_true',
|
|
help='Run headless (no MuJoCo viewer)')
|
|
parser.add_argument('--no-brain', action='store_true',
|
|
help='Run body only with manual drive (no neural sim)')
|
|
parser.add_argument('--no-auto', action='store_true',
|
|
help='Disable auto-demo (manual keyboard only)')
|
|
parser.add_argument('--stimulus', type=str, default=None,
|
|
choices=list(STIMULI.keys()),
|
|
help='Initial stimulus to activate')
|
|
parser.add_argument('--duration', type=float, default=0.0,
|
|
help='Max sim duration in seconds (0 = unlimited)')
|
|
parser.add_argument('--visual', action='store_true',
|
|
help='Enable real vision: compound eye -> connectome '
|
|
'photoreceptors (uses LoomingArena)')
|
|
parser.add_argument('--approach-angle', type=float, default=0.0,
|
|
help='Ball approach angle in degrees '
|
|
'(0=front, +45=right, -45=left)')
|
|
parser.add_argument('--monitor', action='store_true',
|
|
help='Open brain monitor window (dorsal view)')
|
|
parser.add_argument('--somatosensory', action='store_true',
|
|
help='Enable touch (contact forces) and sound '
|
|
'(vibration sources) via JO neurons')
|
|
parser.add_argument('--gustatory', action='store_true',
|
|
help='Enable taste zones (sugar/bitter) on arena '
|
|
'floor via GRN neurons')
|
|
parser.add_argument('--olfactory', action='store_true',
|
|
help='Enable olfactory system: attractive (Or42b) '
|
|
'and repulsive (Or56a) odor sources')
|
|
parser.add_argument('--vocalize', action='store_true',
|
|
help='Enable wing song production (courtship/alarm) '
|
|
'driven by DN activity')
|
|
parser.add_argument('--flight', action='store_true',
|
|
help='Enable virtual flight via external forces '
|
|
'(GF triggers takeoff, xfrc_applied on Thorax)')
|
|
parser.add_argument('--consciousness', action='store_true',
|
|
help='Enable consciousness proxy measurement '
|
|
'(Phi/IIT, GWT, Self-Model, Perturbation)')
|
|
args = parser.parse_args()
|
|
|
|
# -- State --
|
|
active_stimulus = [args.stimulus or 'p9'] # Default: P9 forward walking
|
|
stim_changed = [True]
|
|
auto_demo_enabled = [not args.no_auto and args.stimulus is None
|
|
and not args.visual]
|
|
|
|
# Auto-demo state
|
|
demo_idx = [0]
|
|
demo_time_remaining = [AUTO_DEMO_SEQUENCE[0][1]]
|
|
|
|
# -- Keyboard mapping (GLFW keycodes) --
|
|
KEY_MAP = {
|
|
ord('1'): 'sugar',
|
|
ord('2'): 'p9',
|
|
ord('3'): 'lc4',
|
|
ord('4'): 'jo',
|
|
ord('5'): 'bitter',
|
|
ord('6'): 'or56a',
|
|
ord('0'): None,
|
|
}
|
|
GLFW_KEY_SPACE = 32
|
|
|
|
def key_callback(keycode):
|
|
if keycode == GLFW_KEY_SPACE:
|
|
auto_demo_enabled[0] = not auto_demo_enabled[0]
|
|
state = "ON" if auto_demo_enabled[0] else "OFF"
|
|
print(f"\n[Auto-demo] {state}")
|
|
if auto_demo_enabled[0]:
|
|
demo_idx[0] = 0
|
|
demo_time_remaining[0] = AUTO_DEMO_SEQUENCE[0][1]
|
|
return
|
|
|
|
if keycode in KEY_MAP:
|
|
auto_demo_enabled[0] = False # Manual key disables auto-demo
|
|
active_stimulus[0] = KEY_MAP[keycode]
|
|
stim_changed[0] = True
|
|
name = active_stimulus[0]
|
|
if name and name in STIMULI:
|
|
print(f"\n[Manual] {STIMULI[name]['description']}")
|
|
else:
|
|
print("\n[Manual] OFF -- spontaneous activity")
|
|
|
|
# ── Initialize brain ──
|
|
brain = None
|
|
if not args.no_brain:
|
|
print("Initializing brain (138,639 neurons on GPU)...")
|
|
brain = BrainEngine(device='cuda')
|
|
|
|
# ── Initialize visual system (if --visual) ──
|
|
visual = None
|
|
if args.visual and brain is not None:
|
|
print("Initializing visual system (photoreceptor mapping)...")
|
|
visual = VisualSystem(brain.flyid2i, brain.i2flyid)
|
|
|
|
# ── Initialize somatosensory system (if --somatosensory) ──
|
|
somato = None
|
|
vibration_sources = []
|
|
if args.somatosensory and brain is not None:
|
|
print("Initializing somatosensory system (JO touch + sound)...")
|
|
somato = SomatosensorySystem(brain.flyid2i)
|
|
# Default vibration sources
|
|
vibration_sources = [
|
|
VibrationSource(
|
|
position=[30.0, 20.0, 1.0],
|
|
frequency=200.0, amplitude=0.8,
|
|
label='courtship'),
|
|
VibrationSource(
|
|
position=[-20.0, -15.0, 1.0],
|
|
frequency=400.0, amplitude=0.6,
|
|
label='alarm'),
|
|
]
|
|
for vs in vibration_sources:
|
|
print(f" Vibration: '{vs.label}' at "
|
|
f"[{vs.position[0]:.0f},{vs.position[1]:.0f}]mm "
|
|
f"f={vs.frequency:.0f}Hz amp={vs.amplitude:.1f}")
|
|
|
|
# ── Initialize gustatory system (if --gustatory) ──
|
|
taste_zones = []
|
|
gusto = None
|
|
if args.gustatory and brain is not None:
|
|
print("Initializing gustatory system (sugar/bitter GRN zones)...")
|
|
taste_zones = [
|
|
TasteZone(center=[15.0, 5.0], radius=8.0,
|
|
taste='sugar', label='sugar_patch'),
|
|
TasteZone(center=[20.0, -10.0], radius=6.0,
|
|
taste='bitter', label='bitter_patch'),
|
|
]
|
|
gusto = GustatorySystem(brain.flyid2i, taste_zones)
|
|
|
|
# ── Initialize olfactory system (if --olfactory) ──
|
|
olfact = None
|
|
odor_sources = []
|
|
if args.olfactory and brain is not None:
|
|
print("Initializing olfactory system (Or42b attractive + Or56a repulsive)...")
|
|
odor_sources = [
|
|
OdorSource(
|
|
position=[25.0, 10.0, 1.0],
|
|
odor_type='attractive', amplitude=0.9, spread=25.0,
|
|
label='food'),
|
|
OdorSource(
|
|
position=[-15.0, -12.0, 1.0],
|
|
odor_type='repulsive', amplitude=0.8, spread=20.0,
|
|
label='geosmin'),
|
|
]
|
|
olfact = OlfactorySystem(brain.flyid2i)
|
|
for src in odor_sources:
|
|
print(f" Odor: '{src.label}' ({src.odor_type}) at "
|
|
f"[{src.position[0]:.0f},{src.position[1]:.0f}]mm "
|
|
f"amp={src.amplitude:.1f} spread={src.spread:.0f}mm")
|
|
|
|
# ── Initialize wing song system (if --vocalize) ──
|
|
wing_song = None
|
|
if args.vocalize and brain is not None:
|
|
print("Initializing wing song system (courtship/alarm via DN)...")
|
|
wing_song = WingSongSystem(self_hearing_gain=0.2)
|
|
print(f" Pulse=200Hz, Sine=160Hz, Alarm=400Hz (self-hearing=20%)")
|
|
|
|
# ── Initialize consciousness detection (if --consciousness) ──
|
|
consciousness = None
|
|
if args.consciousness and brain is not None:
|
|
if ConsciousnessDetector is not None:
|
|
consciousness = ConsciousnessDetector(brain)
|
|
else:
|
|
print("[WARN] consciousness.py not found, --consciousness ignored")
|
|
|
|
# ── Initialize flight system placeholder (if --flight) ──
|
|
flight_sys = None
|
|
thorax_body_id = -1
|
|
qpos_adr = -1
|
|
dof_adr = -1
|
|
proboscis_jnt_id = -1
|
|
|
|
# ── Initialize body ──
|
|
print("Initializing body (NeuroMechFly v2 + MuJoCo)...")
|
|
contact_sensors = [
|
|
f"{leg}{seg}"
|
|
for leg in ["LF", "LM", "LH", "RF", "RM", "RH"]
|
|
for seg in ["Tibia", "Tarsus1", "Tarsus2",
|
|
"Tarsus3", "Tarsus4", "Tarsus5"]
|
|
]
|
|
|
|
fly = Fly(
|
|
enable_adhesion=True,
|
|
draw_adhesion=False,
|
|
contact_sensor_placements=contact_sensors,
|
|
enable_vision=args.visual,
|
|
)
|
|
|
|
# ── Add proboscis joint (Rostrum has no joint in stock NeuroMechFly) ──
|
|
rostrum_body = fly.model.find("body", "Rostrum")
|
|
if rostrum_body is not None:
|
|
rostrum_body.add(
|
|
"joint",
|
|
name="joint_Proboscis",
|
|
type="hinge",
|
|
axis=[0, 1, 0], # pitch axis: rotates proboscis down
|
|
range=[-0.1, 1.2], # retracted to fully extended (radians)
|
|
stiffness=50.0, # spring pulls it back (retracted at rest)
|
|
damping=5.0,
|
|
)
|
|
print("[Proboscis] Added hinge joint to Rostrum body")
|
|
else:
|
|
print("[Proboscis] WARNING: Rostrum body not found in fly model")
|
|
|
|
# Use LoomingArena when visual mode is enabled
|
|
arena_kwargs = {}
|
|
if args.visual:
|
|
arena_ground = 500 if args.flight else 100
|
|
arena_start = 120.0 if args.flight else 80.0
|
|
arena_kwargs['arena'] = LoomingArena(
|
|
ball_radius=6.0,
|
|
approach_speed=15.0,
|
|
start_distance=arena_start,
|
|
ball_height=1.5,
|
|
approach_angle=args.approach_angle,
|
|
taste_zones=taste_zones,
|
|
odor_sources=odor_sources,
|
|
ground_size=arena_ground,
|
|
)
|
|
angle_str = f" angle={args.approach_angle}°" if args.approach_angle != 0 else ""
|
|
print(f"[Visual] LoomingArena: r=6mm sphere from {arena_start:.0f}mm at 15mm/s{angle_str}")
|
|
|
|
sim = HybridTurningController(
|
|
fly=fly,
|
|
timestep=1e-4,
|
|
seed=0,
|
|
**arena_kwargs,
|
|
)
|
|
|
|
# ── Disable flygym's internal vision rendering BEFORE reset ──
|
|
# Cameras and retina are already initialized from Fly.__init__().
|
|
# We must prevent flygym from calling dm_control's physics.render()
|
|
# during reset/step — it creates an offscreen GL context that
|
|
# conflicts with the MuJoCo passive viewer on Windows.
|
|
if args.visual:
|
|
fly.enable_vision = False
|
|
|
|
# ── Initialize bridge ──
|
|
decoder = DNRateDecoder(window_ms=50.0, dt_ms=0.1, max_rate=200.0)
|
|
bridge = BrainBodyBridge(decoder, escape_threshold=0.3,
|
|
groom_threshold=0.02)
|
|
groom_ctrl = GroomingController()
|
|
|
|
# ── Register lateralized populations for directional escape ──
|
|
if visual is not None and brain is not None:
|
|
lplc2_idx = visual.get_lplc2_indices(brain.flyid2i)
|
|
lc4_idx = visual.get_lc4_indices(brain.flyid2i)
|
|
for name, indices in {**lplc2_idx, **lc4_idx}.items():
|
|
brain.register_population(name, indices)
|
|
decoder.register_population(name)
|
|
|
|
# ── Register JO populations for monitoring ──
|
|
if somato is not None and brain is not None:
|
|
if len(somato.touch_idx_left) > 0:
|
|
brain.register_population('JO_touch_L', somato.touch_idx_left)
|
|
decoder.register_population('JO_touch_L')
|
|
if len(somato.touch_idx_right) > 0:
|
|
brain.register_population('JO_touch_R', somato.touch_idx_right)
|
|
decoder.register_population('JO_touch_R')
|
|
if len(somato.sound_idx_left) > 0:
|
|
brain.register_population('JO_sound_L', somato.sound_idx_left)
|
|
decoder.register_population('JO_sound_L')
|
|
if len(somato.sound_idx_right) > 0:
|
|
brain.register_population('JO_sound_R', somato.sound_idx_right)
|
|
decoder.register_population('JO_sound_R')
|
|
|
|
# ── Reset simulation ──
|
|
obs, info = sim.reset(seed=0)
|
|
print(f"Fly spawned at {obs['fly'][0]} mm")
|
|
|
|
# ── Post-reset: initialize flight system with model data ──
|
|
if args.flight and brain is not None:
|
|
model_ptr = sim.physics.model.ptr
|
|
fly_name = fly.name
|
|
|
|
# Find Thorax body id
|
|
thorax_name = f"{fly_name}/Thorax"
|
|
thorax_body_id = mujoco.mj_name2id(
|
|
model_ptr, mujoco.mjtObj.mjOBJ_BODY, thorax_name)
|
|
if thorax_body_id < 0:
|
|
thorax_body_id = mujoco.mj_name2id(
|
|
model_ptr, mujoco.mjtObj.mjOBJ_BODY, "Thorax")
|
|
|
|
# Total mass: sum ONLY fly bodies (prefix = fly_name/)
|
|
fly_mass = 0.0
|
|
for bid in range(model_ptr.nbody):
|
|
bname = mujoco.mj_id2name(
|
|
model_ptr, mujoco.mjtObj.mjOBJ_BODY, bid)
|
|
if bname and bname.startswith(f"{fly_name}/"):
|
|
fly_mass += float(model_ptr.body_mass[bid])
|
|
|
|
# Gravity: flygym already uses mm units, so opt.gravity is in mm/s²
|
|
gravity_mm = float(abs(model_ptr.opt.gravity[2]))
|
|
|
|
flight_sys = FlightSystem(
|
|
total_mass=fly_mass,
|
|
gravity=gravity_mm,
|
|
)
|
|
# Find free joint for orientation override during flight
|
|
freejoint_id = -1
|
|
for jid in range(model_ptr.njnt):
|
|
if model_ptr.jnt_type[jid] == mujoco.mjtJoint.mjJNT_FREE:
|
|
jbody = model_ptr.jnt_bodyid[jid]
|
|
jname = mujoco.mj_id2name(
|
|
model_ptr, mujoco.mjtObj.mjOBJ_BODY, jbody)
|
|
if jname and jname.startswith(f"{fly_name}/"):
|
|
freejoint_id = jid
|
|
break
|
|
if freejoint_id >= 0:
|
|
qpos_adr = int(model_ptr.jnt_qposadr[freejoint_id])
|
|
dof_adr = int(model_ptr.jnt_dofadr[freejoint_id])
|
|
else:
|
|
qpos_adr = dof_adr = -1
|
|
|
|
print(f"[Flight] mass={fly_mass*1e6:.1f}mg gravity={gravity_mm:.0f}mm/s² "
|
|
f"mg={flight_sys.mg:.6f}mN thorax_id={thorax_body_id} "
|
|
f"freejoint={freejoint_id}")
|
|
|
|
# ── Find proboscis joint (added dynamically to Rostrum) ──
|
|
model_ptr = sim.physics.model.ptr
|
|
fly_name = fly.name
|
|
for jname_candidate in [f"{fly_name}/joint_Proboscis", "joint_Proboscis"]:
|
|
proboscis_jnt_id = mujoco.mj_name2id(
|
|
model_ptr, mujoco.mjtObj.mjOBJ_JOINT, jname_candidate)
|
|
if proboscis_jnt_id >= 0:
|
|
break
|
|
if proboscis_jnt_id >= 0:
|
|
proboscis_qadr = int(model_ptr.jnt_qposadr[proboscis_jnt_id])
|
|
print(f"[Proboscis] joint_id={proboscis_jnt_id} qpos_adr={proboscis_qadr}")
|
|
else:
|
|
proboscis_qadr = -1
|
|
print("[Proboscis] WARNING: joint not found after compilation")
|
|
|
|
# ── Set up manual vision rendering ──
|
|
eye_renderer = None
|
|
retina = None
|
|
geom_hide_ids = []
|
|
eye_cam_ids = {}
|
|
|
|
if visual is not None:
|
|
retina = fly.retina
|
|
fly_name = fly.name
|
|
|
|
model_ptr = sim.physics.model.ptr
|
|
|
|
# Create MuJoCo-native renderer (coexists with passive viewer)
|
|
eye_renderer = mujoco.Renderer(model_ptr, height=512, width=450)
|
|
|
|
# Locate eye cameras
|
|
for side in ["L", "R"]:
|
|
cam_name = f"{fly_name}/{side}Eye_cam"
|
|
cid = mujoco.mj_name2id(
|
|
model_ptr, mujoco.mjtObj.mjOBJ_CAMERA, cam_name)
|
|
eye_cam_ids[side] = cid
|
|
print(f"[Vision] Camera '{cam_name}' -> id={cid}")
|
|
|
|
# Locate geoms to hide during eye rendering (avoid self-occlusion)
|
|
for geom_name in getattr(fly, '_geoms_to_hide', []):
|
|
full_name = f"{fly_name}/{geom_name}"
|
|
gid = mujoco.mj_name2id(
|
|
model_ptr, mujoco.mjtObj.mjOBJ_GEOM, full_name)
|
|
if gid >= 0:
|
|
geom_hide_ids.append(gid)
|
|
print(f"[Vision] {len(geom_hide_ids)} geoms hidden during eye render")
|
|
|
|
# ── Launch MuJoCo viewer (clean, no UI panels) ──
|
|
viewer = None
|
|
if not args.no_viewer:
|
|
print("Launching MuJoCo viewer...")
|
|
viewer = mujoco.viewer.launch_passive(
|
|
sim.physics.model.ptr, sim.physics.data.ptr,
|
|
key_callback=key_callback,
|
|
show_left_ui=False,
|
|
show_right_ui=False,
|
|
)
|
|
# Configure viewer options and camera
|
|
if viewer is not None:
|
|
viewer.opt.label = mujoco.mjtLabel.mjLABEL_SITE
|
|
# Hide fly's default sites (groups 0-2), show only arena labels (group 4)
|
|
for g in range(3):
|
|
viewer.opt.sitegroup[g] = 0
|
|
viewer.opt.sitegroup[4] = 1
|
|
if viewer is not None and thorax_body_id >= 0:
|
|
viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
|
viewer.cam.trackbodyid = thorax_body_id
|
|
viewer.cam.distance = 40.0
|
|
viewer.cam.azimuth = -120.0
|
|
viewer.cam.elevation = -25.0
|
|
|
|
# ── Set initial stimulus ──
|
|
if brain is not None:
|
|
brain.set_stimulus(active_stimulus[0])
|
|
stim_changed[0] = False
|
|
stim_desc = STIMULI.get(active_stimulus[0], {}).get(
|
|
'description', active_stimulus[0] or 'none')
|
|
print(f"Initial stimulus: {stim_desc}")
|
|
|
|
# ── Launch brain monitor (separate process) ──
|
|
monitor = None
|
|
if args.monitor:
|
|
print("Launching brain monitor...")
|
|
monitor = BrainMonitorProcess()
|
|
monitor.start()
|
|
|
|
# ── Timing constants ──
|
|
MONITOR_INTERVAL = 500 # send data every 500 body steps (~50ms sim)
|
|
BRAIN_RATIO = 100 # 1 brain step per 100 body steps (10Hz neural update)
|
|
VISION_RATIO = 1000 # process vision every 1000 body steps (= 100ms, 10Hz)
|
|
STEPS_PER_FRAME = 167 # body steps per viewer frame (~60fps at 1e-4 timestep)
|
|
STATUS_INTERVAL = 10000 # status print every 1.0s sim time
|
|
|
|
body_step = 0
|
|
prev_mode = 'walking'
|
|
cached_visual = (None, None) # cached ALL visual layer (indices, rates)
|
|
last_vision_obs = None # last vision obs for diagnostics
|
|
physics_errors = 0 # consecutive physics error counter
|
|
|
|
import time as _time
|
|
_frame_target = 1.0 / 60.0 # 16.7ms per frame
|
|
_next_viewer_sync = _time.perf_counter()
|
|
_fps_counter = 0
|
|
_fps_timer = _time.perf_counter()
|
|
_measured_fps = 0.0
|
|
|
|
print()
|
|
print("=" * 70)
|
|
print(" EMBODIED DROSOPHILA -- Interactive Brain-Body Simulation")
|
|
if visual is not None:
|
|
print(" *** REAL VISION ACTIVE: compound eye -> connectome ***")
|
|
if somato is not None:
|
|
print(" *** SOMATOSENSORY ACTIVE: touch + sound -> JO neurons ***")
|
|
if gusto is not None:
|
|
print(" *** GUSTATORY ACTIVE: sugar/bitter zones -> GRN neurons ***")
|
|
if olfact is not None:
|
|
print(" *** OLFACTORY ACTIVE: food/danger odors -> ORN neurons ***")
|
|
if wing_song is not None:
|
|
print(" *** VOCALIZATION ACTIVE: wing song -> JO self-hearing ***")
|
|
if flight_sys is not None:
|
|
print(f" *** FLIGHT ACTIVE: GF > {flight_sys.takeoff_thresh} triggers virtual takeoff ***")
|
|
print("=" * 70)
|
|
if auto_demo_enabled[0]:
|
|
print(" MODE: Auto-demo (cycles through stimuli automatically)")
|
|
print(" Press SPACE to toggle auto-demo on/off")
|
|
else:
|
|
print(" MODE: Manual (use keyboard to change stimulus)")
|
|
if visual is not None:
|
|
print(" VISION: Photoreceptors receive real visual input from flygym")
|
|
print(" A dark sphere approaches — escape should emerge naturally!")
|
|
print(" Keys: 1=sugar 2=P9 3=looming 4=grooming 5=bitter 6=olfactory")
|
|
print(" 0=off SPACE=toggle auto | Close viewer to exit")
|
|
print("=" * 70)
|
|
print()
|
|
|
|
# ── Main loop ──
|
|
try:
|
|
while True:
|
|
# Check exit conditions
|
|
if viewer is not None:
|
|
if not viewer.is_running():
|
|
break
|
|
elif args.duration > 0:
|
|
if body_step * sim.timestep >= args.duration:
|
|
break
|
|
|
|
# ── Auto-demo: advance sequence ──
|
|
if auto_demo_enabled[0] and brain is not None:
|
|
demo_time_remaining[0] -= sim.timestep
|
|
if demo_time_remaining[0] <= 0:
|
|
demo_idx[0] = (demo_idx[0] + 1) % len(AUTO_DEMO_SEQUENCE)
|
|
stim_name, duration, desc = AUTO_DEMO_SEQUENCE[demo_idx[0]]
|
|
demo_time_remaining[0] = duration
|
|
active_stimulus[0] = stim_name
|
|
stim_changed[0] = True
|
|
print(f"\n >>> [{desc}] "
|
|
f"({stim_name or 'none'}, {duration:.1f}s)")
|
|
|
|
# ── Update stimulus ──
|
|
if stim_changed[0] and brain is not None:
|
|
brain.set_stimulus(active_stimulus[0])
|
|
stim_changed[0] = False
|
|
# Re-apply cached visual rates (set_stimulus zeroes all rates)
|
|
if cached_visual[0] is not None:
|
|
brain.set_visual_rates(*cached_visual)
|
|
# Re-apply somatosensory rates
|
|
if somato is not None:
|
|
jo_idx, jo_rates = somato.get_rates()
|
|
brain.set_sensory_rates(jo_idx, jo_rates)
|
|
# Re-apply gustatory rates
|
|
if gusto is not None:
|
|
grn_idx, grn_rates = gusto.get_rates()
|
|
brain.set_sensory_rates(grn_idx, grn_rates)
|
|
# Re-apply olfactory rates
|
|
if olfact is not None:
|
|
or_idx, or_rates = olfact.get_rates()
|
|
brain.set_sensory_rates(or_idx, or_rates)
|
|
|
|
# ── Visual processing (every VISION_RATIO body steps) ──
|
|
# Uses mujoco.Renderer (not dm_control physics.render)
|
|
# to avoid GL context conflict with the passive viewer.
|
|
if visual is not None and body_step % VISION_RATIO == 0:
|
|
model_ptr = sim.physics.model.ptr
|
|
data_ptr = sim.physics.data.ptr
|
|
|
|
# Hide self-geoms (eyes, antennae, coxae) to avoid occlusion
|
|
saved_alpha = []
|
|
for gid in geom_hide_ids:
|
|
saved_alpha.append(model_ptr.geom_rgba[gid, 3].copy())
|
|
model_ptr.geom_rgba[gid, 3] = 0.0
|
|
|
|
# Render both eyes with MuJoCo native renderer
|
|
readouts = []
|
|
for side in ["L", "R"]:
|
|
cid = eye_cam_ids.get(side, -1)
|
|
if cid < 0:
|
|
readouts.append(np.zeros((721, 2), dtype=np.float32))
|
|
continue
|
|
eye_renderer.update_scene(data_ptr, camera=cid)
|
|
raw_img = eye_renderer.render()
|
|
fish_img = retina.correct_fisheye(raw_img)
|
|
hex_pxls = retina.raw_image_to_hex_pxls(fish_img)
|
|
readouts.append(hex_pxls)
|
|
|
|
# Restore self-geoms
|
|
for i, gid in enumerate(geom_hide_ids):
|
|
model_ptr.geom_rgba[gid, 3] = saved_alpha[i]
|
|
|
|
vision_obs = np.array(readouts, dtype=np.float32)
|
|
last_vision_obs = vision_obs
|
|
|
|
# Inject ALL visual layers: R1-R8, L1, L2, Mi1, Tm1, Tm2, T2
|
|
# T2 -> LC4 -> GF propagates through pure connectome
|
|
vis_idx, vis_rates = visual.process_visual_layers(vision_obs)
|
|
cached_visual = (vis_idx, vis_rates)
|
|
brain.set_visual_rates(vis_idx, vis_rates)
|
|
|
|
# ── Somatosensory processing (every brain step) ──
|
|
if somato is not None and body_step % BRAIN_RATIO == 0:
|
|
# Touch: read contact forces from MuJoCo
|
|
contact_forces = obs.get('contact_forces', np.zeros((36, 3)))
|
|
somato.process_contact(contact_forces)
|
|
|
|
# Sound: compute vibration from fly position and heading
|
|
fly_pos = obs['fly'][0] # position in mm
|
|
fly_orient = obs.get('fly_orientation', np.zeros(3))
|
|
# fly_orientation is body X-axis (forward) in world frame
|
|
fly_heading = float(np.arctan2(fly_orient[1], fly_orient[0]))
|
|
somato.process_vibration(fly_pos, fly_heading, vibration_sources)
|
|
|
|
# Inject JO rates into brain
|
|
jo_idx, jo_rates = somato.get_rates()
|
|
brain.set_sensory_rates(jo_idx, jo_rates)
|
|
|
|
# Update bridge with somatosensory state
|
|
bridge.tactile_force = somato.max_contact_force
|
|
bridge.sound_orientation_bias = somato.orientation_bias
|
|
|
|
# ── Gustatory processing (every brain step) ──
|
|
if gusto is not None and body_step % BRAIN_RATIO == 0:
|
|
end_effectors = obs.get('end_effectors', np.zeros((6, 3)))
|
|
gusto.process(end_effectors)
|
|
|
|
# Inject GRN rates into brain
|
|
grn_idx, grn_rates = gusto.get_rates()
|
|
brain.set_sensory_rates(grn_idx, grn_rates)
|
|
|
|
# Update bridge with gustatory state
|
|
bridge.bitter_active = gusto.bitter_active
|
|
|
|
# ── Olfactory processing (every brain step) ──
|
|
if olfact is not None and body_step % BRAIN_RATIO == 0:
|
|
fly_pos = obs['fly'][0]
|
|
fly_orient = obs.get('fly_orientation', np.zeros(3))
|
|
fly_heading = float(np.arctan2(fly_orient[1], fly_orient[0]))
|
|
olfact.process(fly_pos, fly_heading, odor_sources)
|
|
|
|
# Inject ORN rates into brain
|
|
or_idx, or_rates = olfact.get_rates()
|
|
brain.set_sensory_rates(or_idx, or_rates)
|
|
|
|
# Update bridge with olfactory state
|
|
bridge.olfactory_attraction_bias = olfact.attraction_bias
|
|
bridge.olfactory_repulsive = olfact.is_repulsive_escape
|
|
bridge.olfactory_repulsion_bias = olfact.repulsion_bias
|
|
|
|
# ── Wing song processing (every brain step, silent during flight) ──
|
|
if wing_song is not None and body_step % BRAIN_RATIO == 0:
|
|
if bridge.mode != 'flight':
|
|
fly_pos_ws = obs['fly'][0]
|
|
wing_song.process(decoder, fly_pos_ws, BRAIN_RATIO * sim.timestep)
|
|
elif wing_song.is_singing:
|
|
# Flying fly doesn't sing — wings are for flight
|
|
wing_song.active_song = None
|
|
wing_song.wing_freq = 0.0
|
|
wing_song.wing_amp = 0.0
|
|
|
|
# Add wing song vibration to somatosensory input
|
|
if somato is not None and wing_song.is_singing:
|
|
wing_sources = wing_song.get_vibration_sources()
|
|
all_vib = vibration_sources + wing_sources
|
|
fly_orient_ws = obs.get('fly_orientation', np.zeros(3))
|
|
heading_ws = float(np.arctan2(fly_orient_ws[1], fly_orient_ws[0]))
|
|
somato.process_vibration(fly_pos_ws, heading_ws, all_vib)
|
|
# Re-inject updated JO rates
|
|
jo_idx, jo_rates = somato.get_rates()
|
|
brain.set_sensory_rates(jo_idx, jo_rates)
|
|
|
|
# ── Flight processing (every brain step) ──
|
|
if flight_sys is not None and body_step % BRAIN_RATIO == 0:
|
|
fly_pos_fl = obs['fly'][0]
|
|
# fly_orientation IS the forward direction vector (body X-axis in world)
|
|
fly_fwd = obs.get('fly_orientation', np.array([1.0, 0.0, 0.0]))
|
|
flight_sys.update(
|
|
decoder, fly_pos_fl, fly_fwd,
|
|
BRAIN_RATIO * sim.timestep)
|
|
bridge.flight_active = flight_sys.is_airborne
|
|
|
|
# ── Brain step (1 per BRAIN_RATIO body steps) ──
|
|
if brain is not None and body_step % BRAIN_RATIO == 0:
|
|
brain.step()
|
|
dn_spikes = brain.get_dn_spikes()
|
|
pop_spikes = brain.get_population_spikes() if brain.populations else None
|
|
decoder.update(dn_spikes, pop_spikes)
|
|
if consciousness is not None:
|
|
consciousness.update(body_step, bridge.mode)
|
|
|
|
# ── Per-eye T2 fallback for directional escape ──
|
|
if visual is not None and cached_visual[0] is not None:
|
|
vis_eye = visual._T2_eye
|
|
vis_rates_arr = cached_visual[1]
|
|
if vis_rates_arr is not None and len(vis_rates_arr) > 0:
|
|
mask_L = vis_eye == 0
|
|
mask_R = vis_eye == 1
|
|
t2_left = float(vis_rates_arr[mask_L].mean()) if mask_L.any() else 0.0
|
|
t2_right = float(vis_rates_arr[mask_R].mean()) if mask_R.any() else 0.0
|
|
bridge.visual_threat_bias = (
|
|
(t2_right - t2_left) / (t2_left + t2_right + 1e-6))
|
|
|
|
# ── Compute drive from DN rates ──
|
|
drive = bridge.compute_drive(dt=BRAIN_RATIO * sim.timestep)
|
|
|
|
# ── Mode transitions (throttle prints to avoid I/O spam) ──
|
|
if bridge.mode != prev_mode:
|
|
if (bridge.mode == 'flight' or prev_mode == 'flight'
|
|
or body_step % 100 == 0):
|
|
print(f" >> Behavior: {prev_mode} -> {bridge.mode}")
|
|
if bridge.mode == 'flight' and flight_sys is not None:
|
|
h = flight_sys._escape_heading
|
|
y = flight_sys._escape_yaw
|
|
print(f" [Flight] escape heading=[{h[0]:.2f},{h[1]:.2f}] "
|
|
f"yaw={np.degrees(y):.1f}°")
|
|
prev_mode = bridge.mode
|
|
|
|
# ── Apply flight forces (before physics step) ──
|
|
if flight_sys is not None:
|
|
data_ptr = sim.physics.data.ptr
|
|
if flight_sys.is_airborne:
|
|
data_ptr.xfrc_applied[thorax_body_id] = flight_sys.force_torque
|
|
else:
|
|
# Clear residual forces when grounded
|
|
data_ptr.xfrc_applied[thorax_body_id] = 0.0
|
|
|
|
# ── Body step ──
|
|
try:
|
|
if bridge.mode == 'flight':
|
|
# Flight: legs frozen in neutral pose, adhesion OFF
|
|
flight_action = {
|
|
"joints": groom_ctrl.neutral.copy(),
|
|
"adhesion": np.zeros(6),
|
|
}
|
|
obs, reward, terminated, truncated, info = \
|
|
SingleFlySimulation.step(sim, flight_action)
|
|
elif bridge.mode == 'grooming':
|
|
groom_action = groom_ctrl.get_action(
|
|
body_step * sim.timestep)
|
|
obs, reward, terminated, truncated, info = \
|
|
SingleFlySimulation.step(sim, groom_action)
|
|
else:
|
|
obs, reward, terminated, truncated, info = sim.step(drive)
|
|
physics_errors = 0 # reset on success
|
|
|
|
# ── Proboscis extension during feeding ──
|
|
if proboscis_qadr >= 0 and bridge.mode == 'feeding':
|
|
sim.physics.data.ptr.qpos[proboscis_qadr] = 1.0
|
|
|
|
# ── Orientation override during flight (post-step) ──
|
|
# Directly set the free joint quaternion to prevent spinning.
|
|
# Torques can't reliably control orientation on an articulated
|
|
# body with 100+ joints. This guarantees rigid heading lock.
|
|
if (flight_sys is not None and flight_sys.is_airborne
|
|
and qpos_adr >= 0):
|
|
data_ptr = sim.physics.data.ptr
|
|
desired_quat = flight_sys.get_desired_quat()
|
|
data_ptr.qpos[qpos_adr + 3:qpos_adr + 7] = desired_quat
|
|
data_ptr.qvel[dof_adr + 3:dof_adr + 6] = 0.0
|
|
# During landing: damp upward velocity to help descent
|
|
if flight_sys.state == FlightState.LANDING:
|
|
vz = data_ptr.qvel[dof_adr + 2]
|
|
if vz > 0: # going up — damp it
|
|
data_ptr.qvel[dof_adr + 2] = vz * 0.95
|
|
# Also damp horizontal for clean landing
|
|
data_ptr.qvel[dof_adr + 0] *= 0.98
|
|
data_ptr.qvel[dof_adr + 1] *= 0.98
|
|
|
|
except Exception as e:
|
|
physics_errors += 1
|
|
if physics_errors >= 50:
|
|
print(f" Physics unstable ({physics_errors} errors): {e}")
|
|
break
|
|
continue # skip this step, try to recover
|
|
|
|
body_step += 1
|
|
|
|
# ── Sync viewer at wall-clock 60fps ──
|
|
if viewer is not None and body_step % STEPS_PER_FRAME == 0:
|
|
_now = _time.perf_counter()
|
|
# Pace simulation to real-time
|
|
_sleep = _next_viewer_sync - _now
|
|
if _sleep > 0.001:
|
|
_time.sleep(_sleep)
|
|
viewer.sync()
|
|
# Prevent accumulated time debt when falling behind
|
|
_next_viewer_sync = max(
|
|
_next_viewer_sync, _now) + _frame_target
|
|
# FPS measurement
|
|
_fps_counter += 1
|
|
if _now - _fps_timer >= 1.0:
|
|
_measured_fps = _fps_counter / (_now - _fps_timer)
|
|
_fps_counter = 0
|
|
_fps_timer = _now
|
|
|
|
# ── Status print ──
|
|
if body_step % STATUS_INTERVAL == 0:
|
|
t_sim = body_step * sim.timestep
|
|
pos = obs['fly'][0]
|
|
stim_label = active_stimulus[0] or 'none'
|
|
auto_tag = " [auto]" if auto_demo_enabled[0] else ""
|
|
pos_z = pos[2] if len(pos) > 2 else 0.0
|
|
flight_tag = ""
|
|
if flight_sys is not None and flight_sys.is_airborne:
|
|
flight_tag = f" | {flight_sys.get_status_str()}"
|
|
fps_tag = f" {_measured_fps:.0f}fps" if _measured_fps > 0 else ""
|
|
status_line = (
|
|
f" t={t_sim:.1f}s{fps_tag} stim={stim_label:>6s}{auto_tag} "
|
|
f"pos=[{pos[0]:.1f}, {pos[1]:.1f}, {pos_z:.1f}] mm "
|
|
f"{bridge.get_status_str()}{flight_tag}"
|
|
)
|
|
# Add visual circuit monitoring when vision is active
|
|
if visual is not None:
|
|
d = decoder
|
|
gf_rate = np.mean([
|
|
d.get_normalized('GF_1'),
|
|
d.get_normalized('GF_2')])
|
|
ball_x = getattr(
|
|
arena_kwargs.get('arena'), 'ball_pos',
|
|
np.array([0]))[0]
|
|
# LC4 spike monitoring (downstream from lamina)
|
|
lc4_info = ""
|
|
if brain is not None:
|
|
lc4_idx = brain.stim_indices.get('lc4', [])
|
|
if lc4_idx:
|
|
spk = brain.state[2]
|
|
lc4_spikes = spk[0, lc4_idx].sum().item()
|
|
lc4_info = f" LC4spk={lc4_spikes:.0f}"
|
|
# Vision diagnostics (retina values are [0,1])
|
|
vis_info = ""
|
|
if last_vision_obs is not None:
|
|
bL = np.mean(last_vision_obs[0])
|
|
bR = np.mean(last_vision_obs[1])
|
|
n_dark_L = np.sum(
|
|
np.mean(last_vision_obs[0], axis=1) < 0.25)
|
|
n_dark_R = np.sum(
|
|
np.mean(last_vision_obs[1], axis=1) < 0.25)
|
|
vis_info = (
|
|
f" bright=[{bL:.2f},{bR:.2f}]"
|
|
f" dark_omm=[{n_dark_L},{n_dark_R}]")
|
|
# Lamina neuron counts
|
|
lamina_info = f" L1={visual._n_L1} L2={visual._n_L2} T2={visual._n_T2}"
|
|
# LPLC2/LC4 laterality
|
|
lplc2_info = ""
|
|
lplc2_L = decoder.get_pop_rate('LPLC2_left')
|
|
lplc2_R = decoder.get_pop_rate('LPLC2_right')
|
|
if lplc2_L > 0 or lplc2_R > 0:
|
|
lplc2_info = f" LPLC2=[{lplc2_L:.1f},{lplc2_R:.1f}]"
|
|
status_line += (
|
|
f" | ball_x={ball_x:.0f}"
|
|
f" GF={gf_rate:.3f}{lc4_info}"
|
|
f"{lamina_info}{lplc2_info}{vis_info}")
|
|
# Add somatosensory monitoring
|
|
if somato is not None:
|
|
somato_str = somato.get_status_str()
|
|
if somato_str != "JO=silent":
|
|
status_line += f" | {somato_str}"
|
|
# Add gustatory monitoring
|
|
if gusto is not None:
|
|
gusto_str = gusto.get_status_str()
|
|
if gusto_str:
|
|
status_line += f" | {gusto_str}"
|
|
# Add olfactory monitoring
|
|
if olfact is not None:
|
|
olf_str = olfact.get_status_str()
|
|
if olf_str:
|
|
status_line += f" | {olf_str}"
|
|
# Add wing song monitoring
|
|
if wing_song is not None:
|
|
ws_str = wing_song.get_status_str()
|
|
if ws_str:
|
|
status_line += f" | {ws_str}"
|
|
# Add consciousness monitoring
|
|
if consciousness is not None:
|
|
status_line += f" | {consciousness.get_status_str()}"
|
|
print(status_line)
|
|
|
|
# ── Send data to brain monitor ──
|
|
if monitor is not None and body_step % MONITOR_INTERVAL == 0:
|
|
d = decoder
|
|
mon_data = {
|
|
't_sim': body_step * sim.timestep,
|
|
'mode': bridge.mode,
|
|
'drive': [bridge.left_drive, bridge.right_drive],
|
|
'stimulus': active_stimulus[0] or 'none',
|
|
'dn_forward': d.get_group_rate('forward'),
|
|
'dn_escape': d.get_group_rate('escape'),
|
|
'dn_groom': d.get_group_rate('groom'),
|
|
'dn_backward': d.get_group_rate('backward'),
|
|
'dn_feed': d.get_group_rate('feed'),
|
|
'dn_turn_L': d.get_group_rate('turn_L'),
|
|
'dn_turn_R': d.get_group_rate('turn_R'),
|
|
'threat_asym': bridge.threat_asym,
|
|
}
|
|
# Somatosensory data (when available)
|
|
if somato is not None:
|
|
mon_data['jo_contact'] = somato.touch_level
|
|
mon_data['jo_sound'] = somato.sound_level
|
|
mon_data['jo_touch_L'] = (
|
|
somato.touch_rate_left / somato.TOUCH_MAX_RATE)
|
|
mon_data['jo_touch_R'] = (
|
|
somato.touch_rate_right / somato.TOUCH_MAX_RATE)
|
|
mon_data['jo_sound_L'] = (
|
|
somato.sound_rate_left / somato.SOUND_MAX_RATE)
|
|
mon_data['jo_sound_R'] = (
|
|
somato.sound_rate_right / somato.SOUND_MAX_RATE)
|
|
mon_data['contact_force'] = somato.max_contact_force
|
|
mon_data['sound_bias'] = somato.orientation_bias
|
|
# Gustatory data (when available)
|
|
if gusto is not None:
|
|
mon_data['sugar_level'] = gusto.sugar_level
|
|
mon_data['bitter_level'] = gusto.bitter_level
|
|
# Wing song data (when available)
|
|
if wing_song is not None:
|
|
mon_data['wing_freq'] = wing_song.wing_freq
|
|
mon_data['wing_amp'] = wing_song.wing_amp
|
|
mon_data['wing_song'] = wing_song.active_song or 'silent'
|
|
mon_data['wing_level'] = wing_song.song_level
|
|
# Olfactory data (when available)
|
|
if olfact is not None:
|
|
mon_data['or_att_L'] = (
|
|
olfact.attractive_rate_left / olfact.ATTRACTIVE_MAX_RATE)
|
|
mon_data['or_att_R'] = (
|
|
olfact.attractive_rate_right / olfact.ATTRACTIVE_MAX_RATE)
|
|
mon_data['or_rep_L'] = (
|
|
olfact.repulsive_rate_left / olfact.REPULSIVE_MAX_RATE)
|
|
mon_data['or_rep_R'] = (
|
|
olfact.repulsive_rate_right / olfact.REPULSIVE_MAX_RATE)
|
|
mon_data['or_attractive'] = olfact.attractive_level
|
|
mon_data['or_repulsive'] = olfact.repulsive_level
|
|
# Flight data (when available)
|
|
if flight_sys is not None:
|
|
mon_data['flight_level'] = flight_sys.flight_level
|
|
mon_data['flight_state'] = flight_sys.state.name.lower()
|
|
mon_data['flight_alt'] = flight_sys.altitude
|
|
mon_data['flight_wing_freq'] = flight_sys.wing_freq
|
|
# Visual data (when available)
|
|
if visual is not None:
|
|
mon_data['lplc2_left'] = d.get_pop_rate('LPLC2_left')
|
|
mon_data['lplc2_right'] = d.get_pop_rate('LPLC2_right')
|
|
mon_data['lc4_left'] = d.get_pop_rate('LC4_left')
|
|
mon_data['lc4_right'] = d.get_pop_rate('LC4_right')
|
|
if last_vision_obs is not None:
|
|
mon_data['bright_left'] = float(
|
|
np.mean(last_vision_obs[0]))
|
|
mon_data['bright_right'] = float(
|
|
np.mean(last_vision_obs[1]))
|
|
mon_data['dark_omm_left'] = int(np.sum(
|
|
np.mean(last_vision_obs[0], axis=1) < 0.25))
|
|
mon_data['dark_omm_right'] = int(np.sum(
|
|
np.mean(last_vision_obs[1], axis=1) < 0.25))
|
|
if cached_visual[1] is not None and hasattr(visual, '_T2_eye'):
|
|
vis_eye = visual._T2_eye
|
|
vis_r = cached_visual[1]
|
|
mask_L = vis_eye == 0
|
|
mask_R = vis_eye == 1
|
|
mon_data['t2_left'] = float(
|
|
vis_r[mask_L].mean() / 120.0) if mask_L.any() else 0.0
|
|
mon_data['t2_right'] = float(
|
|
vis_r[mask_R].mean() / 120.0) if mask_R.any() else 0.0
|
|
ball_pos = getattr(
|
|
arena_kwargs.get('arena'), 'ball_pos', None)
|
|
if ball_pos is not None:
|
|
mon_data['ball_x'] = float(ball_pos[0])
|
|
if consciousness is not None:
|
|
mon_data.update(consciousness.get_monitor_data())
|
|
monitor.send(mon_data)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nStopped by user.")
|
|
|
|
finally:
|
|
if consciousness is not None:
|
|
consciousness.save_session()
|
|
if brain is not None:
|
|
brain.save_plastic_weights()
|
|
if monitor is not None:
|
|
monitor.stop()
|
|
if viewer is not None:
|
|
viewer.close()
|
|
print(f"\nSimulation ended after {body_step} steps "
|
|
f"({body_step * sim.timestep:.2f}s sim time).")
|
|
if body_step > 0:
|
|
pos = obs['fly'][0]
|
|
print(f"Final position: [{pos[0]:.1f}, {pos[1]:.1f}, "
|
|
f"{pos[2]:.1f}] mm")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|