You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
# MI-HGNN for contact estimation/classification on various robots
2
-
This repository implements a Morphology-Informed Heterogeneous Graph Neural Network (MI-HGNN) for estimating contact information on the feet of a quadruped robot.
This repository implements a Morphology-Inspired Heterogeneous Graph Neural Network (MI-HGNN) for estimating contact information on the feet of a quadruped robot. For more details, see our publication "[MI-HGNN: Morphology-Informed Heterogeneous Graph Neural Network for Legged Robot Contact Perception](https://arxiv.org/abs/2409.11146)" and our [project page](https://lunarlab-gatech.github.io/Morphology-Informed-HGNN/).
3
3
4
-
Additionally, by providing a compatible URDF file, this software can convert a variety of robot structures to graph format for learning with the MI-HGNN. See [#Applying-MI-HGNN-to-your-own-robot
5
-
](#applying-mi-hgnn-to-your-own-robot) for more information.
4
+
Additionally, it can be applied to a variety of robot structures and datasets, as our software can convert compatible robot URDF files to graph format and provides a template for implementing custom datasets. See [#Applying-MI-HGNN-to-your-own-robot](#applying-mi-hgnn-to-your-own-robot) for more information.
For information on our method, see our [project page](https://lunarlab-gatech.github.io/Morphology-Informed-HGNN/) and [paper](https://arxiv.org/abs/2409.11146).
8
+
## Setup
9
+
---
10
10
11
-
## Installation
11
+
###Installation
12
12
To get started, setup a Conda Python environment with Python=3.11:
13
13
```
14
14
conda create -n mi-hgnn python=3.11
@@ -22,37 +22,62 @@ pip install .
22
22
23
23
Note, if you have any issues with setup, refer to `environment_files/README.md` so you can install the exact libraries we used.
24
24
25
-
## URDF Download
25
+
###URDF Download
26
26
The necessary URDF files are part of git submodules in this repository, so run the following commands to download them:
27
27
```
28
28
git submodule init
29
29
git submodule update
30
30
```
31
31
32
-
## Replicating Paper Experiments
32
+
## Usage
33
+
---
33
34
34
-
To replicate the experiments referenced in our paper or access our trained model weights, see `paper/README.md`.
35
+
### Replicating Paper Experiments
35
36
36
-
## Applying MI-HGNN to your own robot
37
+
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.
37
38
38
-
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. New URDF files can be added by following the instructions in `urdf_files/README.md`, and our software will automatically convert the URDF into a graph compatible for learning with the MI-HGNN.
39
+
<imgsrc="paper/website_images/figure5.png"alt="Parameter sizes and Ablation study"width="600">
39
40
40
-
##Editing and Contributing
41
+
### Applying to your Robot/Dataset
41
42
42
-
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.
43
+
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:
43
44
44
-
After making changes, rebuild the library following the instructions in [#Installation](#installation). To make sure that your changes haven't
45
-
broken critical functionality, run the test cases found in the `tests` directory.
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.
47
+
1. Add new URDF files for your robots by following the instructions in `urdf_files/README.md`. Our software will automatically convert the URDF into a graph compatible for learning with the MI-HGNN.
48
+
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`.
49
+
3. After making your changes, rebuild the library following the instructions in [#Installation](#installation). To make sure that your changes haven't
50
+
broken critical functionality, run the test cases with the command `python -m unittest discover tests/ -v`.
51
+
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.
48
52
49
-
## Citation
53
+
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!
54
+
55
+
56
+
### Simulated A1 Dataset
57
+
58
+
To evaluate the performance of our model on GRF estimation, we generated our own simulated GRF dataset, which we now contribute to the community as well. We recorded proprioceptive sensor data and the corresponding ground truth GRFs by operating an A1 robot in the [Quad-SDK](https://github.com/lunarlab-gatech/quad_sdk_fork) simulator. In total, our dataset comprises of 530,779 synchronized data samples with a variety of frictions, terrains, and speeds. All of the different sequences are outlined in the table below:
A visualization of the various data collection environments can be seen below.
63
+
64
+
65
+
66
+
If you'd like to use this dataset, the recorded sequences can be found on [Dropbox](https://www.dropbox.com/scl/fo/4iz1oobx71qoceu2jenie/AJPggD4yIAFXf5508wBz-hY?rlkey=4miys9ap0iaozgdelntms8lxb&st=0oz7kgyq&dl=0). See `paper/README.md` and Section V-B of our publication for specific details on this dataset and how to use it.
67
+
68
+
## Other Info
69
+
---
70
+
### Contributing
71
+
72
+
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.
73
+
74
+
### Citation
50
75
51
76
If you find our repository or our work useful, please cite the relevant publication:
author={Butterfield, Daniel and Garimella, Sandilya Sai and Cheng, Nai-Jen and Gan, Lu},
57
82
journal={arXiv preprint arXiv:2409.11146},
58
83
year={2024},
@@ -61,6 +86,6 @@ If you find our repository or our work useful, please cite the relevant publicat
61
86
}
62
87
```
63
88
64
-
## Contact / Issues
89
+
###Contact / Issues
65
90
66
91
For any issues with this repository, feel free to open an issue on GitHub. For other inquiries, please contact Daniel Butterfield ([email protected]) or the Lunar Lab (https://sites.gatech.edu/lunarlab/).
Copy file name to clipboardExpand all lines: paper/README.md
+11-1
Original file line number
Diff line number
Diff line change
@@ -99,11 +99,21 @@ Finally, Figure 5 is generated by running the `create_regression_plots.py` file,
99
99
100
100
### GRF Quad-SDK Dataset
101
101
102
-
For this experiment, we used a dataset that we generated ourselves using Quad-SDK and Gazebo. Our modified fork that we used can be found here: [quad-sdk-fork](https://github.com/lunarlab-gatech/quad_sdk_fork). We generated a total of 21 sequences. The following table relates the dataset sequence name (in code) to the corresponding parameters used for that sequence:
102
+
For this experiment, we generated a dataset using Quad-SDK and Gazebo. The dataset consists of synchronized proprioceptive sensor measurements for a simulated A1 robot at a maximum of 500 Hz, including joint angle, joint angular velocity, and joint torque from 12 joint encoders, base linear acceleration, base angular velocity from IMU, and GRFs for each leg in Z direction. It also includes the ground truth robot pose, represented as a translation and a quaternion.
103
+
104
+
We generated a total of 21 sequences. The following table relates the dataset sequence name (in code) to the corresponding parameters used for that sequence:
The dataset files can be found on [Dropbox](https://www.dropbox.com/scl/fo/4iz1oobx71qoceu2jenie/AJPggD4yIAFXf5508wBz-hY?rlkey=4miys9ap0iaozgdelntms8lxb&st=0oz7kgyq&dl=0), and the code that we used to generate this dataset can be found here: [quad-sdk-fork](https://github.com/lunarlab-gatech/quad_sdk_fork). To find the file corresponding to a sequence in the table above, notice that the files are named with the following convention below:
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:
4
+
- Automatic download of relevant datasets from the Internet (from Google Drive or Dropbox).
5
+
- Data sorting to match the order of joint, foot, and base nodes in the robot graph.
6
+
- Wrapper for the `RobotGraph` class that generates the graph from the robot URDF file.
7
+
- Easy customization with custom history lengths and a normalization parameter.
8
+
- Provides custom get() function returns for training both an MLP and the MI-HGNN.
9
+
- 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.
10
+
11
+
However, `FlexibleDataset` currently only supports the following input data:
12
+
- lin_acc (np.array) - IMU linear acceleration
13
+
- ang_vel (np.array) - IMU angular velocity
14
+
- j_p (np.array) - Joint positions
15
+
- j_v (np.array) - Joint velocities
16
+
- j_T (np.array) - Joint Torques
17
+
- f_p (np.array) - Foot position
18
+
- f_v (np.array) - Foot velocity
19
+
- labels (np.array) - The Dataset labels (either Z direction GRFs, or contact states)
20
+
- r_p (np.array) - Robot position (GT)
21
+
- r_o (np.array) - Robot orientation (GT) as a quaternion, in the order (x, y, z, w)
22
+
- timestamps (np.array) - Array containing the timestamps of the data
23
+
24
+
Also note that not all of these are used depending on the applied model (MLP vs. MIHGNN vs Floating-Base Dynamics).
25
+
26
+
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.
27
+
28
+
## Using the Custom Dataset Template
29
+
30
+
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.
31
+
32
+
First, open the file and rename the class to your liking.
33
+
34
+
### Adding Dataset Sequences
35
+
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.
36
+
37
+
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.
38
+
39
+
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).
40
+
41
+
### Implementing Data Processing
42
+
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.
43
+
44
+
Implement this function. You can see `quadSDKDataset.py` for an example of converting a ROSbag file into a .mat file.
45
+
46
+
### Implementing Data Loading
47
+
Now that data is loaded and processed, you can now implement the function for opening the .mat file and extracting the relevant dataset sequence.
48
+
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.
49
+
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.
50
+
51
+
### Setting the proper URDF file
52
+
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:
53
+
54
+
```
55
+
<robot name="miniCheetah">
56
+
```
57
+
58
+
This name should be pasted into `get_expected_urdf_name()`.
59
+
60
+
### Facilitating Data Sorting
61
+
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()`. See `CustomDatasetTemplate.py` for more details.
62
+
63
+
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 like in the examples in the `research` directory. Happy Training!
0 commit comments