ASTGCN & Baseline Models Repository

This repository contains implementations and experiments for spatial–temporal graph convolution models and several baseline methods. The primary focus is on ASTGCN (Attention-based Spatio-Temporal Graph Convolutional Network), but you can also run an MLP, a TGCN, and a ConvLSTM model.

Repo link: https://github.com/pauljac7700/ASTGCN-robot-pose-deviation.git

Note: This repository is developed for Python 3.10.14.

---

Overview

• ASTGCN
  Implements both single and multi-target versions with several ablations:

  • ASTGCN_multi: The main model with multiple output nodes for pose targets and residuals.
  • Ablation variants include:
    • ASTGCN_multi_no_attention
    • ASTGCN_multi_no_spatial
    • ASTGCN_multi_no_temporal

• MLP
  A baseline Multi-Layer Perceptron that runs on preprocessed data. Use the run_MLP_multi script with its configuration.

• TGCN
  A Temporal Graph Convolutional Network that uses the preprocessed data (generated by the ASTGCN prep script). It comes with its own configuration, training, and test scripts.

• ConvLSTM
  A Convolutional LSTM model that also runs on the preprocessed ASTGCN data. It includes its own prep, training, and test scripts.
  Important: Run the ConvLSTM prep script immediately after generating the ASTGCN data.

---

Installation

1. Clone the Repository:

   git clone https://github.com/pauljac7700/ASTGCN-robot-pose-deviation.git
   cd ASTGCN-robot-pose-deviation

2. Set Up Python Environment:

   Ensure you are using Python 3.10.14. For example, using conda:

   conda create -n astgcn_env python=3.10.14
   conda activate astgcn_env

3. Install Dependencies:

   Install the required packages (e.g., using pip):

   pip install -r requirements.txt

---

Data Preparation

1. Configuration:
   YAML configuration files (e.g., config_ASTGCN.yaml, config_TGCN.yaml, config_ConvLSTM.yaml) set dataset parameters, model hyperparameters, and training settings.

2. Preprocessing for ASTGCN:
   Run the ASTGCN data preparation script (e.g., prep_data_ASTGCN.py) to split your dataset into training, validation, and test sets. This script also computes scalers and saves the preprocessed data under the data/ folder.

3. Additional Preprocessing:

   • ConvLSTM: Run the ConvLSTM preprocessing script immediately after the ASTGCN data is prepared.
   • MLP: The MLP model uses the ASTGCN preprocessed data and is run via the run_MLP_multi script.

---

Models and How to Run

ASTGCN

• Training:
  Run the training script with the appropriate config:

  python train_ASTGCN_multi.py --config config_ASTGCN.yaml

• Testing:
  Evaluate the trained model using:

  python test_ASTGCN_multi.py --config config_ASTGCN.yaml

• Ablation Study:
  Use train_ASTGCN_multi_ablation.py and its corresponding test script for ablation experiments. The variant is selected by the ablation_model key in the config.

MLP

• Running the MLP:
  Execute the run_MLP_multi script along with its configuration file.

TGCN

• Training & Testing:
  TGCN comes with its own scripts (train_TGCN.py and test_TGCN.py). TGCN uses the ASTGCN preprocessed data as input.

ConvLSTM

• Preprocessing:
  First, run the ConvLSTM prep script (e.g., prep_ConvLSTM.py) right after generating ASTGCN data.
• Training & Testing:
  Use train_ConvLSTM.py and test_ConvLSTM.py along with the provided configuration file.

---

Configuration Files

The YAML configuration files include:

• Dataset Settings:

  • dataset_dimension: "6D" or "3D"
  • dataset_name, dataset_type
  • Data splitting ratios (test_size, val_size)

• Preprocessing Options:

  • prep_data_incl_past_residuals
  • scalers_file
  • adjacency_matrix_file

• Target and Residual Variables:

  • Target Pose Variables:
    For 6D: [x_t, y_t, z_t, rx_t, ry_t, rz_t]
    For 3D: [x_t, y_t, z_t]
  • Residual Variables:
    For 6D: [x_dif, y_dif, z_dif, rx_dif, ry_dif, rz_dif]
    For 3D: [x_dif, y_dif, z_dif]

• Model Hyperparameters:

  • nb_block, in_channels_6D, in_channels_3D, K, nb_chev_filter, nb_time_filter, time_strides, num_for_predict, len_input

• Training Settings:

  • Learning rate, weight decay, batch size, number of epochs, early stopping patience

• Device:

  • cuda or cpu

• Model Names:

  • E.g., ASTGCN_multi, ASTGCN_multi_no_attention, etc.

---

How to Use

1. Set Up Environment & Install Dependencies:
   Follow the installation steps above.

2. Prepare Data:
   Run the ASTGCN data preparation script to generate training, validation, and test sets. Run additional preprocessing (ConvLSTM, etc.) as needed.

3. Select and Configure Model:
   Choose the model (ASTGCN, MLP, TGCN, or ConvLSTM) you wish to run. Edit the corresponding YAML config file to set dataset parameters, model hyperparameters, and training settings.

4. Run Training:
   For example, to train ASTGCN:

   python train_ASTGCN_multi.py --config config_ASTGCN.yaml

5. Run Testing:
   To evaluate the model:

   python test_ASTGCN_multi.py --config config_ASTGCN.yaml

6. For TGCN and ConvLSTM:
   Use their dedicated training and testing scripts along with their configuration files.

7. For the MLP:
   Run the run_MLP_multi script with its configuration.

---

Graph Structure and Residual Extraction

• Graph 1, 3, 6, 7:
  Residual predictions are taken from the node at index num_joints+1.

• Graph 2:
  Residual predictions are extracted from multiple nodes starting at index num_joints.

• Graph 4 & 5:
  The graph splits the residual information into two nodes:

  • One for residual position (3 features)
  • One for residual orientation (3 features)
    In this case, if only one prediction step is used, the script squeezes the time dimension and flattens the two nodes into a single vector (resulting in a vector of 6 features per sample).

---

Contact

For any questions or issues, please contact:

Paul Jacobi
paul.jacobi@rwth-aachen.de