Training LSTMs and ANNs to perform tasks with musculoskeletal models.
Environments include monkey model performing cycling.
Please cite the following paper if using uSim/nuSim in your work:
Link to corresponding paper (https://www.biorxiv.org/content/10.1101/2024.02.02.578628v1)
We highly recommend a linux system for easy installation.
First you will need to install Mujoco (older version). Please make sure that Anaconda as well as git are also installed on your system.
Download the library using this link: https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz
Create a hidden folder in your root directory called .mujoco as such (replacing the path with the path on your computer):
mkdir /home/username/.mujoco
Extract the downloaded library into the newly created hidden folder:
tar -xvf mujoco210-linux-x86_64.tar.gz -C ~/.mujoco/
Open your .bashrc file in your root/home directory:
nano .bashrc
Once in the .bashrc file, add the following line replacing the path with your true home directory:
export LD_LIBRARY_PATH=/home/user-name/.mujoco/mujoco210/bin
If your system has an nvidia GPU, add this line as well:
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia
Save, close, then source the .bashrc using the following command:
source ~/.bashrc
Reboot your system to ensure changes are made
Create a new environment using conda:
conda env create --name nusim --file=requirements.yml
Activate the conda environment:
conda activate nusim
If facing errors, installing the following libraries may help
If you are on linux (and may apply to Mac as well), there will likely be additional packages necessary. Here is a list of possible packages:
If facing errors, adding these packages may help.
Lastly, within the conda environment there are additional packages necessary to ensure the training can run:
To train the controller, run the following in terminal:
python append_musculo_targets.py
python find_init_pose.py --config configs/configs.txt
python main.py --config configs/configs.txt
This will save the controller in the ./checkpoint file with training iterations. The highest reward should reach >= 55000 for kinematic accuracy.
The episode reward with iterations should look like this:
(There may be slight variations due to random seed but trend should look similar)
Inputs:
Musculoskeletal Model:
(The path to the musculoskeletal_model.xml can also be specified in the configs.txt file with musculoskeletal_model_path param, if not using the above default path)
Experimental Kinematics:
Save the experimental kinematics in â./kinematics_data/kinematics.pklâ as a Python dict object with the following format:
dict{
<'marker_names'> : <['marker_name_1', ..., 'marker_name_n']>,
<'train'> : <dict_train>,
<'test'> : <dict_test>
}
{
<key> : <value>,
<key>: <value>,
.
.
.
<key>: <value>
}
<key: int=""> is the integer index of the corresponding condition. (starts from 0 for the first condition for both the training and testing conditions) </key:>
<value: numpy.ndarray=""> contains the kinematics for the corresponding condition with shape: [num_markers/targets, num_coordinates = 3, timepoints]. </value:>
num_markers are the number of experimental markers/bodies that are recorded. The order of the num_markers must correspond to the order in which the marker_names are listed. For example, if âmarker_namesâ = [âhandâ, âelbowâ], num_marker= 0 should contain the experimental kinematics for hand and num_marker=1 should contain the experimental kinematics for elbow.
num_coordinates are the x [-->], y[â] and z[ out of page] coordinates. Values of NaN for any coordinate will keep that coordinate locked.
An example for saving the experimental kinematics for the cycling task is given in ./exp_kin_cycling/saving_exp_kin.ipynb
(The path to kinematics.pkl file can also be specified using kinematics_path param in configs.txt file)
Neural Data (optional):
<'train'> : <dict_train>,
<'test'> : <dict_test>
}
<key: int=""> is the integer index of the corresponding condition as in the kinematics file.</key:>
<value: numpy.ndarray=""> is the numpy array that contains recorded firing rates with the following shape: [timepoints, num_neurons]. num_neurons are the total number of recorded neurons.</value:>
Note: If this step is omitted, various post-processing analyses which require recorded neural data such as CCA, will not run. nuSim training will also not proceed.
(nusim_data_path can also be specified in the configs.txt file)
Stimulus Data (optional):
Provide any experimental stimulus data in â./stimulus_data/stimulus_data.pklâ as a Python dict object.
dict{
<'train'> : <dict_train>,
<'test'> : <dict_test>
}
<key: int=""> is the integer index of the corresponding condition as in the kinematics file.</key:>
<value: numpy.ndarray=""> is the numpy array that contains recorded stimulus data with the following shape: [timepoints, num_features]. num_features are the corresponding features in that stimulus.</value:>
Initial Pose (optional):
Save the initial pose (containing the qpos and qvel) as numpy arrays in â./inital_pose/â as qpos.npy and qvel.npy with shape [nq, ]. nq is the number of joints in the xml model.
This step is optional. If omitted, the default initial pose for xml model will be used for CMA-ES and IK.
(initial_pose_path can also be specified in the configs.txt file)
Specifications:
Provide the parameters for various modules using the â./configs/configs.txtâ file. The details of each parameter/specification is given in the configs.txt file.
General Usage:
Inverse Kinematics:
Run:
python append_musculo_targets.py
This will append targets to the musculoskeletal xml file that will follow the preprocessed markers kinematics during simulation.
a. Run the following command in the terminal:
python find_init_pose.py --config configs/configs.txt --visualize True
This will use inverse kinematics (IK) to find the initial pose for the xml model to match the initial timepoint of the target kinematics.
If you see the output, âInitial Pose found and savedâ, skip 1b.
b. Run:
python find_init_pose_ik_cma.py --config configs/configs.txt --visualize True
This willl use CMA-ES optimization with IK to find a good initial pose for the xml model.
If you see, âInitial Pose found and saved using CMA-ES and Inverse Kinematicsâ, proceed to the next step.
Otherwise, provide a good inital pose for the xml model that preferably starts nearer to the inital marker/target position.
Run
python main --config configs/configs.txt --visualize True --mode test
This will visualize the target trajectories using a randomly initialized uSim controller network. Make sure target trajectories look as desired. Otherwise, change the kinematics preprocessing parameters (e.g. trajectory_scaling, center) in the ./configs/configs.txt file.
Run:
python visualize_trajectories_ik.py --config configs/configs.txt --visualize True
This will visualize the xml model following/tracking the training target trajectories. Before proceeding, make sure that the target trajectories are feasible and lie within the bounds of the xml model. Otherwise, adjust the target trajectories using the kinematics preprocessing parameters in configs.txt file.
This will also save the generated sensory feedback in â./test_data/sensory_feedback_ik.pklâ as Python dict object:
<key: int=""> corresponds to the integer index of the corresponding training condition</key:>
<value: numpy.ndarray=""> with shape: [timepoints, num_of_state_feedback_variables]</value:>
This can be used to get Proprioception for training neural networks.
Training the uSim Controller using DRL:
(Make sure DRL/SAC related parameters are specified correctly in the configs.txt file)
python main.py --config configs/configs.txt
python main.py --config configs/configs.txt --load_saved_nets_for_training True
Testing the uSim Controller:
To test the trained uSim controller, run:
python main.py --config configs/configs.txt --mode test --visualize True
This will visualize the xml model performing movements for training and testing conditions using the trained uSim controller.
This will also save the files used for post training analyses.
Post Training Analyses:
After training, the following modules are used for various analyses. All these modules are in â./Analysisâ
To visualize the kinematics for the training and testing conditions, see visualize_kinematics.ipynb
To visualize the uSim controllerâs population trajectories in PCA subspace, run:
python collective_pca.py
see CCA.ipynb
see LRA.ipynb
see procrustes.ipynb
Clone the fixed-point-finder in ./Analysis, https://github.com/mattgolub/fixed-point-finder
Run
python find_fp.py
The fixed point analysis is based on the original implementation: https://github.com/mattgolub/fixed-point-finder. Refer to the github repo for further information.
See and run jpca_nusim.m
Note: jPCA analysis is based on MM Churchlandâs original implementation. Please see it for further details (https://www.dropbox.com/scl/fo/duf5zbwcibsux467c6oc9/AIN-ZiFsy2Huyh8h7VMdL7g?rlkey=3o5axmq5hirel4cij7g64jc0r&e=1&dl=0)
Important for jPCA analysis:
Make sure that ./Analyses/jPCA_ForDistribution is included in the MATLAB path alongwith all sub-directories
Make sure that usim test_data folder is included in the MATLAB path. test_data folder is where the jpca data is saved during usim test
Perturbation Analyses:
Selective Feedback Elimination (SFE):
Specify the part of the sensory feedback to be eliminated in ./SAC/perturbation_specs.py using sf_elim variable. Run:
python main --config configs/configs.txt --mode SFE --visualize True
Sensory Perturbation:
Specify the perturbation vector to be added to the selected sensory feedback in ./SAC/perturbation_specs.py, e.g. muscle_lengths_pert. Run:
python main.py --config configs/configs.txt --mode sensory_pert --visualize True
Neural Perturbation:
The neural perturbation will add the given perturbation to the nodes of the uSim/nuSim controllerâs RNN.
Specify the neural perturbation vector in perturbation_specs.py using neural_pert variable. Run:
python main.py --config configs/configs.txt --mode neural_pert --visualize True
Change Musculoskeletal Properties:
To test the trained uSim controller under changed musculoskeletal properties:
Go to the folder â./musculoskeletal_model/â. Copy and paste the xml model âmusculo_targets.xmlâ. Rename the copied model as âmusculo_targets_pert.xmlâ.
Change the desired musculoskeletal properties in xml model âmusculo_targets_pert.xmlâ.
Run:
python main.py --config configs/configs.txt --mode musculo_properties --visualize True
All the above perturbation analyses will change the post training analyses files in place. To run the post training analyses after perturbation see Post Training Analyses section.