Mini Tutorial for applying to a custom dataset/robot

Author: DanielChaseButterfield

README.md

@@ -32,16 +32,21 @@ git submodule update
 
 ### Replicating Paper Experiments
 
-To replicate the experiments referenced in our paper (like the "GRF Estimation" experiment shown below) or to access our trained model weights, see `paper/README.md`.
+We provide code for replicating the exact experiments in our paper and provide full model weights for every model referenced in our paper. See `paper/README.md` for more information.
 
 <img src="paper/website_images/figure5.png" alt="Parameter sizes and Ablation study" width="600">
 
 ### Applying to your Robot/Dataset
 
-Although in our paper, we only applied the MI-HGNN on quadruped robots for contact perception, it can also be applied to other multi-body dynamical systems and on other tasks/datasets. New URDF files can be added by following the instructions in `urdf_files/README.md`. Datasets can be found in the `src/mi_hgnn/datasets_py` directory, and model definitions and training code can be found in the `src/mi_hgnn/lightning_py` directory. We encourage you to extend the library for your own applications. Please reference [#Replicating-Paper-Experiments](#replicating-paper-experiments) for examples on how to train and evaluate models with our repository.
+Although our paper's scope was limited to application of MI-HGNN on quadruped robots for contact perception, it can easily be applied to other multi-body dynamical systems and on other tasks/datasets, following the steps below:
 
-After making changes, rebuild the library following the instructions in [#Installation](#installation). To make sure that your changes haven't
-broken critical functionality, run the test cases found in the `tests` directory.
+1. Add new URDF files for your robots by following the instructions in `urdf_files/README.md`. 
+2. Incorporate your custom dataset using our `FlexibleDataset` class and starter `CustomDatasetTemplate.py` file by following the instructions at `src/mi_hgnn/datasets_py/README.md`.
+3. After making your changes, rebuild the library following the instructions in [#Installation](#installation). To make sure that your changes haven't
+broken critical functionality, run the test cases with the command `python -m unittest discover tests/ -v`.
+4. Using the files in the `research` directory as an example, call our `train_model` and `evaluate_model` functions provided in `src/mi_hgnn/lightning_py/gnnLightning.py` with defined train, validation, and test sequences. 
+
+We've designed the library to be easily applicable to a variety of datasets and robots, and have provided a variety of customization options in training, dataset creation, and logging. We're excited to see everything you can do with the MI-HGNN!
 
 ### Simulated A1 Dataset
 
@@ -59,7 +64,7 @@ If you'd like to use this dataset, the recorded sequences can be found on [Dropb
 ---
 ### Contributing
 
-If you'd like to contribute to the repository, write sufficient and necessary test cases for your additions in the `tests` directory, and then open a pull request.
+We encourage you to extend the library for your own applications. If you'd like to contribute to the repository, write sufficient and necessary test cases for your additions in the `tests` directory, and then open a pull request. Reach out to us if you have any questions.
 
 ### Citation
 

src/mi_hgnn/datasets_py/CustomDatasetTemplate.py

@@ -0,0 +1,118 @@
+from .flexibleDataset import FlexibleDataset
+import scipy.io as sio
+from pathlib import Path
+import numpy as np
+
+class CustomDataset(FlexibleDataset):
+
+
+    # ========================= DOWNLOADING ==========================
+    def get_downloaded_dataset_file_name(self):
+        """
+        Type the name of the file extension of your dataset sequence files here!
+        """
+        return "data.<YOUR_EXTENSION_HERE>"
+
+    # ========================= PROCESSING ===========================
+    def process(self):
+        # Load the path to the downoaded file
+        path_to_file = Path(self.root, 'raw', 'data.<YOUR_EXTENSION_HERE>')
+
+        # TODO: Convert into a MATLAB data dictionary format here!
+        mat_data = None
+
+        # Make sure to save it at this location
+        sio.savemat(Path(self.root, 'processed', 'data.mat'), mat_data)
+
+        # TODO: Get the number of dataset entries in the file
+        dataset_entries = None
+
+        # Write a txt file to save the dataset length & and first sequence index
+        with open(str(Path(self.processed_dir, "info.txt")), "w") as f:
+            file_id, loc = self.get_file_id_and_loc()
+            f.write(str(dataset_entries) + " " + file_id)
+
+    # ============= DATA SORTING ORDER AND MAPPINGS ==================
+    def get_urdf_name_to_dataset_array_index(self) -> dict:
+        """
+        Implement this function to tell `FlexibleDataset` how 
+        the data returned by `load_data_at_dataset_seq()` corresponds
+        to the joints in the robot URDF file.
+
+        Traditionally a robot only has one base node, so it should get a value
+        of 0. Next, type the name of each leg joint in the URDF file, and add
+        the index of its value in the corresponding joint arrays returned by
+        load_data_at_dataset_seq(). Do the same for the joints in the URDF
+        representing a fixed foot, with the indices of their values in the foot 
+        position and foot velocity arrays.
+        """
+
+        return {
+            '<URDF_BASE_NODE>': 0,
+
+            '<URDF_JOINT_NODE>': 2,
+            '<URDF_JOINT_NODE2>': 0,
+            '<URDF_JOINT_NODE3>': 1,
+
+            '<URDF_FOOT_NODE>': 1,
+            '<URDF_FOOT_NODE2>': 0,
+            }
+
+    # ===================== DATASET PROPERTIES =======================
+    def get_expected_urdf_name(self):
+        return "<EXPECTED_URDF_NAME_HERE>"
+
+    # ======================== DATA LOADING ==========================
+    def load_data_at_dataset_seq(self, seq_num: int):
+        """
+        When this function is called, the .mat file data saved in process()
+        is available at self.mat_data.
+
+        For information on the expected format of these variables, see the
+        load_data_at_dataset_seq() function defition in flexibleDataset.py.
+        """
+
+        # TODO: Load the data as numpy arrays, and don't forget to incorporate self.history_length
+        # to load a history of measurments.
+        lin_acc = None
+        ang_vel = None
+        j_p = None
+        j_v = None
+        j_T = None
+        f_p = None
+        f_v = None
+        contact_labels = None
+        r_p = None
+        r_o = None
+        timestamps = None
+        # Note, if you don't have data for a specific return value, just return None, 
+        # and `FlexibleDataset` will know not to use it if it is not required.
+
+        return lin_acc, ang_vel, j_p, j_v, j_T, f_p, f_v, contact_labels, r_p, r_o, timestamps
+    
+# ================================================================
+# ===================== DATASET SEQUENCES ========================
+# ================================================================
+
+class CustomDataset_sequence1(CustomDataset):
+    """
+    To load a dataset sequence from Google, first upload the corresponding file on Google Drive, set "General Access"
+    to "Anyone with the link", and then copy the link. Paste the link, and then extract the string between the text of
+    '/file/d/' and '/view?usp=sharing'. Take this string, and paste it as the first return argument below.
+    """
+    def get_file_id_and_loc(self):
+        return "17h4kMUKMymG_GzTZTMHPgj-ImDKZMg3R", "Google"
+    
+class CustomDataset_sequence2(CustomDataset):
+    """
+    To load a dataset sequence from Dropbox, first you'll need to upload the corresponding file on Dropbox and  
+    generate a link for viewing. Make sure that access is given to anyone with the link, and that this permission won't
+    expire, doesn't require a password, and allows for downloading. Finally, copy and paste the link as the first return
+    argument below, but change the last number from 0 to 1 (this tells Dropbox to send the raw file, instead of a webpage).
+    """
+    def get_file_id_and_loc(self):
+        return "<Your_Link_Here>", "Dropbox"
+    
+"""
+Create classes for each of your sequences...
+"""

src/mi_hgnn/datasets_py/README.md

@@ -0,0 +1,63 @@
+# Implementing Custom Datasets
+
+We hope that many people use our MI-HGNN on a variety of datasets. We provide the `FlexibleDataset` class which provides many convenient features and can be inherited for use with custom datasets. Below is a short summary of its features:
+- Automatic download of relevant datasets from the Internet (from Google Drive or Dropbox).
+- Data sorting to match the order of joint, foot, and base nodes in the robot graph.
+- Wrapper for the `RobotGraph` class that generates the graph from the robot URDF file.
+- Easy customization with custom history lengths and a normalization parameter.
+- Provides custom get() function returns for training both an MLP and the MI-HGNN.
+- Option for easy evaluation on floating-base dynamics model, though our current implementation is specific for the simulated A1 robot in our paper, meaning changes will be necessary for proper results on your robot.
+
+However, `FlexibleDataset` currently only supports the following input data:
+- lin_acc (np.array) - IMU linear acceleration
+- ang_vel (np.array) - IMU angular velocity
+- j_p (np.array) - Joint positions
+- j_v (np.array) - Joint velocities
+- j_T (np.array) - Joint Torques
+- f_p (np.array) - Foot position
+- f_v (np.array) - Foot velocity
+- labels (np.array) - The Dataset labels (either Z direction GRFs, or contact states)
+- r_p (np.array) - Robot position (GT)
+- r_o (np.array) - Robot orientation (GT) as a quaternion, in the order (x, y, z, w)
+- timestamps (np.array) - Array containing the timestamps of the data
+
+Also note that not all of these are used depending on the applied model (MLP vs. MIHGNN vs Floating-Base Dynamics).
+
+If `FlexibleDataset` supports your input data, then you can easily use it by writing a simple dataset class that inherits from `FlexibleDataset`, similar to `LinTzuYaunDataset` or `QuadSDKDataset`. We've provided a template for you in the `CustomDatasetTemplate.py` file, which you can use to start. 
+
+## Using the Custom Dataset Template
+
+This section will explain how to edit the `CustomDatasetTemplate.py` file for use with your own dataset to take advantage of the features of the `FlexibleDataset` class. 
+
+First, open the file and rename the class to your liking.
+
+### Adding Dataset Sequences
+Next, scroll down to the bottom of the file where it says `DATASET SEQUENCES`. Add every sequence of your dataset as its own class, which will require you to upload the data either to Dropbox or Google. See `CustomDatasetTemplate.py` for details.
+
+This is a clean way for data loading, as it allows the user to later combine different sequences as they'd like with the `torch.utils.data.ConcatDataset` function (see `research/train_classification_sample_eff.py` for an example). Defining these classes also means that training an MI-HGNN model on a different computer doesn't require the user to manually download any datasets, as `FlexibleDataset` will do it for you.
+
+Also, when the files are downloaded, they will be renamed to the value provided by `get_downloaded_dataset_file_name()`. Overwrite this function so that the file extension is correct `.mat` for a Matlab file, `.bag` for a ROSbag file, etc.
+
+### Implementing Data Processing
+Now that you can load your dataset files, you need to implement processing. This step should be implemented in `process()`, and should convert the file from whatever format it is currently in into a `.mat` file for fast training speeds. You'll also need to provide code for extracting the number of dataset entries in this sequence, which will be saved into a .txt file for future use.
+
+Implement this function. You can see `quadSDKDataset.py` for an example of converting a ROSbag file into a .mat file.
+
+### Loading data for use with FlexibleDataset
+Now that data is loaded and processed, you can now implement the function for opening the .mat file and extracting the relevant dataset sequence.
+This should be done in `load_data_at_dataset_seq()`. The .mat file you saved in the last step will now be available at `self.mat_data` for easy access.
+Note that this function will also need to use the `self.history_length` parameter to support training with a history of measurements. See `CustomDatasetTemplate.py` for details, and see `LinTzuYaunDataset.py` for a proper implementation of this function.
+
+### Setting the proper URDF file
+Since its easy for the user to provide the wrong URDF file for a dataset sequence, `FlexibleDataset` checks that the URDF file provided by the user matches what the dataset expects. You can tell `FlexibleDataset` which URDF file should be used with this dataset by going to the URDF file and copying the name found at the top of the file, like pictured below:
+
+```
+<robot name="miniCheetah">
+```
+
+This name should be pasted into `get_expected_urdf_name()`.
+
+### Facilitating Data Sorting
+Finally, the last step is to tell `FlexibleDataset` what order your dataset data is in. For example, which index in the joint position array corresponds to a specific joint in the URDF file? To do this, you'll implement `get_urdf_name_to_dataset_array_index()`.
+
+After doing this, your dataset will work with our current codebase for training MLP and MI-HGNN models! You can now instantiate your dataset and use it in a similar manner to the datasets in the `research` directory. Happy Training!

urdf_files/README.md

@@ -6,20 +6,23 @@ Note, the ```urdfParser.py``` file will take the ```*.urdf``` file and generate
 
 ## Adding a new URDF file
 
-Before you add a new URDF file, you need to make sure that it has all of it's non-kinematic nodes pruned. If it doesn't you'll have to generate a new URDF file without them. See [Generating new URDF files](#generating-new-urdf-files) for more information.
+Before you add a new URDF file, you need to make sure that it has all of it's non-kinematic nodes pruned. If it doesn't, you'll have to generate a new URDF file without them. See [Generating new URDF files](#generating-new-urdf-files) for more information.
 
-Once you have a new URDF file, make sure to create a new folder for it and add the URDF. Add the repository that the URDF depends on to the folder as a git submodule. Finally, if you want to, add the PDF following the instructions in [Generating PDF](#generating-pdf)
+Once you have a new URDF file, make sure to create a new folder for it and add the URDF. Add the repository that the URDF depends on to the folder as a git submodule. Finally, if you desire, you can add the PDF for the URDF following the instructions in [Generating PDF](#generating-pdf).
 
 ### Generating new URDF files
 
 First, make sure to install ROS and the corresponding repository that comes along with the URDF file (which I'll call the URDF repository). Build the repository using ROS, and make sure it builds without errors. You may need to install more dependencies, depending on what the URDF repository says.
 
-Next, find the xacro file used to create the URDF. This will most likely be found in the URDF repository. Manually comment out any non-kinematic structures. Then run the following command to generate the new URDF file:
+Next, find the xacro file used to create the URDF. This will most likely be found in the URDF repository. Manually comment out any non-kinematic structures. Why is this necessary? Our MI-HGNN relies upon the input graph being morphology-informed, which we define as graph nodes representing kinematic joints and graph edges representing kinematic links. Thus, it assumes that the input graph is composed of all of the robots kinematic joints, and that no other fixed structures are included (like cameras or IMUs). Failing to remove any non-kinematic structures will cause unexpected effects, likely decreasing model performance. An easy way to make sure you have done this properly is by [generating a PDF](#generating-pdf) of the URDF file after you create it.
+
+Run the following command to generate the new URDF file:
 
 ```
 rosrun xacro xacro <xacro_path> > <new_urdf_path>
 ```
 
+
 ### Generating PDF
 To generate the pdf file that corresponds to the urdf file, install ROS 1, and run the following command: