Motor Policy Training
Relevant source files
The following files were used as context for generating this wiki page:
Motor Policy Training optimizes a neural model based on a fruit fly connectome to control simulated limb actuators. This training enables the fly model to physically operate vehicle controls, such as a steering wheel or pedals, within the CARLA simulation environment. The system uses measured connectivity from the MaleCNS dataset to constrain the model architecture.
Sources: README.md:3-8, scripts/train_wheel.py:46-52
Training Architecture
The training process uses a supervised learning approach called "Teacher Training" where an offline inverse kinematics (IK) solver provides expert demonstrations. The WheelPolicy model processes sensory observations and outputs joint commands for the fly's legs.
The diagram shows the flow from expert demonstrations to the optimization of the connectome-based motor policy.
Sources: scripts/train_wheel.py:48-52, scripts/train_wheel.py:101-105
Core Components
- WheelPolicy: A PyTorch-based neural network model that retains the measured neuron-pair connections. It includes trained gains and rate dynamics for approximately 165,122 neurons.
- WheelRig / PedalRig: MuJoCo-based simulation environments that provide the physical body, joints, and contact mechanics for the wheel and pedal interfaces.
- Sensory/Motor Populations: Specific subsets of neurons identified as
vnc_sensoryandvnc_motoract as the input and output interfaces for the model.
Sources: scripts/train_wheel.py:65-71, scripts/train_wheel.py:102-104, scripts/pedal_experiment.py:33-35
Training Workflow
The training workflow consists of demonstration gathering, model calibration, and optimization loops.
1. Expert Demonstration
An offline IK solver generates actions for a range of targets (e.g., wheel angles from -0.45 to 0.45 radians). These observations and actions are stored as training data.
Sources: scripts/train_wheel.py:72-79
2. Decoder Calibration
The policy calibrates its frozen decoder using sensory observations. This step occurs before the main optimization loop to stabilize the sensory mapping.
Sources: scripts/train_wheel.py:82-83
3. Optimization Loop
The system uses the Adam optimizer to minimize the mean squared error between the policy's predicted actions and the teacher's expert actions.
| Parameter | Value | Description |
|---|---|---|
| Learning Rate | 0.04 | Step size for gradient updates |
| Batch Size | 16 | Number of training samples per step |
| Optimizer | Adam | Optimization algorithm used |
| Loss Function | MSE | Square of prediction error divided by action scale |
| Training Steps | 600 | Number of iterations in the pilot run |
Sources: scripts/train_wheel.py:53-56, scripts/train_wheel.py:91-100
Evaluation and Validation
After training, the policy undergoes physical evaluation against 100 held-out targets. A trial is considered successful if the final error is within a specific threshold (e.g., 0.13 radians for steering).
The sequence shows the evaluation loop where the trained policy is tested on unseen targets without teacher assistance.
Sources: scripts/train_wheel.py:18-38, scripts/train_wheel.py:51-52
Data Structures
The system records training history and evaluation metrics to JSON files for analysis.
Evaluation Metrics
| Field | Type | Description |
|---|---|---|
initial_successes | int | Passed trials before training (baseline) |
trained_successes | int | Passed trials after training |
edge_gain_mean_abs_change | float | Magnitude of weight updates in the core model |
peak_gpu_allocated_gb | float | Maximum VRAM usage during training |
Sources: scripts/train_wheel.py:116-125
Observation Vector
The observation provided to the model is a concatenated vector containing:
- The target value (e.g., requested wheel angle).
- The current physical position of the control (e.g.,
rig.angle). - The specific joint positions (
qpos) relevant to the limb being trained.
Sources: scripts/train_wheel.py:10-11
Conclusion
Motor Policy Training bridges the gap between biological connectome data and physical robotic control. By training only the core edge gains and leaks while keeping sensory and motor interfaces fixed, the system demonstrates that connectome-constrained models can learn specific motor skills like stationary steering with high accuracy (100/100 targets passed in pilot tests).