From 5b9eea6c20c5d23fb9d539980065ea4df29fdee6 Mon Sep 17 00:00:00 2001 From: Daniel Butterfield Date: Sat, 4 May 2024 19:00:46 -0400 Subject: [PATCH] HeterogeneousRobotGraph & Heterogeneous Dataset for Go1 Simulated Data --- research/animations.py | 13 +- research/devel.py | 4 +- research/evaluator.py | 6 +- research/urdfTest.py | 21 +- src/grfgnn/__init__.py | 2 +- src/grfgnn/datasets.py | 182 ++++++++++---- src/grfgnn/gnnLightning.py | 4 +- src/grfgnn/{urdfParser.py => graphParser.py} | 237 ++++++++++++++---- tests/testDatasets.py | 125 +++++++-- .../{testUrdfParser.py => testGraphParser.py} | 133 +++++++--- 10 files changed, 552 insertions(+), 175 deletions(-) rename src/grfgnn/{urdfParser.py => graphParser.py} (70%) rename tests/{testUrdfParser.py => testGraphParser.py} (67%) diff --git a/research/animations.py b/research/animations.py index 3261e5c..8aab7ae 100644 --- a/research/animations.py +++ b/research/animations.py @@ -1,17 +1,20 @@ from manim import * -from grfgnn import RobotURDF +from grfgnn import NormalRobotGraph + class CreateURDF(Scene): + def construct(self): # Load the A1 urdf - path_to_urdf = Path(Path('.').parent, 'urdf_files', 'A1', 'a1.urdf').absolute() - A1_URDF = RobotURDF(path_to_urdf, 'package://a1_description/', - 'unitree_ros/robots/a1_description', True) + path_to_urdf = Path(Path('.').parent, 'urdf_files', 'A1', + 'a1.urdf').absolute() + A1_URDF = NormalRobotGraph(path_to_urdf, 'package://a1_description/', + 'unitree_ros/robots/a1_description', True) # Create a Circle and give it text text = Text('base').scale(2) rect_1 = RoundedRectangle(corner_radius=0.5) self.play(Create(rect_1)) - self.play(Write(text)) \ No newline at end of file + self.play(Write(text)) diff --git a/research/devel.py b/research/devel.py index 18b35f4..b3cb876 100644 --- a/research/devel.py +++ b/research/devel.py @@ -1,5 +1,5 @@ from grfgnn.datasets import visualize_graph, CerberusTrackDataset, CerberusStreetDataset, CerberusCampusDataset, Go1SimulatedDataset -from grfgnn.urdfParser import RobotURDF +from grfgnn.graphParser import NormalRobotGraph from pathlib import Path from grfgnn.gnnLightning import evaluate_model_and_visualize import os @@ -20,7 +20,7 @@ def dataset_create(): Path('.').parent, 'datasets', 'xiong_simulated').absolute() # Load URDF files - GO1_URDF = RobotURDF(path_to_go1_urdf, 'package://go1_description/', + GO1_URDF = NormalRobotGraph(path_to_go1_urdf, 'package://go1_description/', 'unitree_ros/robots/go1_description', True) model_type = 'gnn' diff --git a/research/evaluator.py b/research/evaluator.py index 76b1681..8fc1ec9 100644 --- a/research/evaluator.py +++ b/research/evaluator.py @@ -1,5 +1,5 @@ from grfgnn.datasets import visualize_graph, CerberusTrackDataset, CerberusStreetDataset, CerberusCampusDataset -from grfgnn.urdfParser import RobotURDF +from grfgnn.graphParser import NormalRobotGraph from pathlib import Path from grfgnn.gnnLightning import evaluate_model_and_visualize @@ -16,9 +16,9 @@ #path_to_checkpoint = Path(Path('.').parent, 'models', TODO, 'epoch=6-val_loss=5764.05.ckpt') # Load URDF files -A1_URDF = RobotURDF(path_to_a1_urdf, 'package://a1_description/', +A1_URDF = NormalRobotGraph(path_to_a1_urdf, 'package://a1_description/', 'unitree_ros/robots/a1_description', True) -GO1_URDF = RobotURDF(path_to_go1_urdf, 'package://go1_description/', +GO1_URDF = NormalRobotGraph(path_to_go1_urdf, 'package://go1_description/', 'unitree_ros/robots/go1_description', True) model_type = 'mlp' diff --git a/research/urdfTest.py b/research/urdfTest.py index d8b7b4c..88aa014 100644 --- a/research/urdfTest.py +++ b/research/urdfTest.py @@ -1,12 +1,17 @@ -from grfgnn import RobotURDF +from grfgnn import NormalRobotGraph, HeterogeneousRobotGraph from pathlib import Path -path_to_a1_urdf = Path(Path('.').parent, 'urdf_files', 'A1', 'a1.urdf').absolute() -A1_URDF = RobotURDF(path_to_a1_urdf, 'package://a1_description/', - 'unitree_ros/robots/a1_description', True) +# path_to_a1_urdf = Path(Path('.').parent, 'urdf_files', 'A1', +# 'a1.urdf').absolute() +# A1_URDF = NormalRobotGraph(path_to_a1_urdf, 'package://a1_description/', +# 'unitree_ros/robots/a1_description', True) -path_to_go1_urdf = Path(Path('.').parent, 'urdf_files', 'Go1', 'go1.urdf').absolute() -GO1_URDF = RobotURDF(path_to_go1_urdf, 'package://go1_description/', - 'unitree_ros/robots/go1_description', True) +path_to_go1_urdf = Path(Path('.').parent, 'urdf_files', 'Go1', + 'go1.urdf').absolute() +# GO1_URDF = NormalRobotGraph(path_to_go1_urdf, 'package://go1_description/', +# 'unitree_ros/robots/go1_description', True) -print(GO1_URDF.robot_urdf.name) +GO1_HETERO_GRAPH = HeterogeneousRobotGraph( + path_to_go1_urdf, 'package://go1_description/', + 'unitree_ros/robots/go1_description', True) +GO1_HETERO_GRAPH.display_graph_info() diff --git a/src/grfgnn/__init__.py b/src/grfgnn/__init__.py index 4f9ef9a..444405f 100644 --- a/src/grfgnn/__init__.py +++ b/src/grfgnn/__init__.py @@ -1,2 +1,2 @@ -from .urdfParser import RobotURDF +from .graphParser import RobotGraph, NormalRobotGraph, HeterogeneousRobotGraph from .datasets import CerberusStreetDataset, CerberusTrackDataset, CerberusCampusDataset, Go1SimulatedDataset, FlexibleDataset diff --git a/src/grfgnn/datasets.py b/src/grfgnn/datasets.py index 05d149f..63d50fe 100644 --- a/src/grfgnn/datasets.py +++ b/src/grfgnn/datasets.py @@ -1,7 +1,7 @@ import torch import torch_geometric -from torch_geometric.data import Data, Dataset -from .urdfParser import RobotURDF +from torch_geometric.data import Data, Dataset, HeteroData +from .graphParser import RobotGraph, NormalRobotGraph, HeterogeneousRobotGraph import networkx import numpy as np import matplotlib.pyplot as plt @@ -31,21 +31,22 @@ class FlexibleDataset(Dataset): def __init__(self, root: str, - robotURDF: RobotURDF, + robotURDF: RobotGraph, data_format: str = 'gnn', transform=None, pre_transform=None, pre_filter=None): """ Parameters: - data_format (str): Either 'gnn' or 'mlp'. Determines how the - get() method returns data so that the dataset can be - trained with a gnn or an mlp. + data_format (str): Either 'gnn', 'mlp', or 'heterogeneous_gnn'. + Determines how the get() method returns data. """ # Check for valid data format self.data_format = data_format - if self.data_format != 'mlp' and self.data_format != 'gnn': - raise ValueError("Parameter 'data_format' must be 'gnn' or 'mlp'.") + if self.data_format != 'mlp' and self.data_format != 'gnn' and self.data_format != 'heterogeneous_gnn': + raise ValueError( + "Parameter 'data_format' must be 'gnn', 'mlp', or 'heterogeneous_gnn'." + ) # Setup the directories for raw and processed data self.root = root @@ -73,11 +74,6 @@ def __init__(self, # Save the URDF info self.URDF = robotURDF - # Open up the urdf file and get the edge matrix - self.edge_matrix = self.URDF.get_edge_index_matrix() - self.edge_matrix_tensor = torch.tensor(self.edge_matrix, - dtype=torch.long) - # Get node name to index mapping self.urdf_name_to_index = self.URDF.get_node_name_to_index_dict() @@ -109,14 +105,14 @@ def get_ground_truth_label_indices(self): should match the output of the node at index 3, and so on. Returns: - ground_truth_indices (list[int]) - List of the indices of + ground_truth_graph_indices (list[int]) - List of the indices of the graph nodes that should output the ground truth labels at the same corresponding index. """ - if self.get_ground_truth_indices is None: + if self.get_ground_truth_graph_indices is None: raise self.notImplementedError else: - return self.ground_truth_indices + return self.ground_truth_graph_indices def get_start_and_end_seq_ids(self): """ @@ -160,6 +156,12 @@ def get_helper_mlp(self, idx): """ raise self.notImplementedError + def get_helper_heterogeneous_gnn(self, idx): + """ + Get a dataset entry if we are using a Heterogeneous GNN model. + """ + raise self.notImplementedError + def get(self, idx): # Bounds check if idx < 0 or idx >= self.length: @@ -170,6 +172,8 @@ def get(self, idx): return self.get_helper_gnn(idx) elif self.data_format == 'mlp': return self.get_helper_mlp(idx) + elif self.data_format == 'heterogeneous_gnn': + return self.get_helper_heterogeneous_gnn(idx) class Go1SimulatedDataset(FlexibleDataset): @@ -179,7 +183,7 @@ class Go1SimulatedDataset(FlexibleDataset): def __init__(self, root: str, - robotURDF: RobotURDF, + robotURDF: HeterogeneousRobotGraph, data_format: str = 'gnn', transform=None, pre_transform=None, @@ -191,36 +195,43 @@ def __init__(self, pre_transform, pre_filter) # Map urdf names to array indexes - # TODO: BELOW NOT VERIFIED self.urdf_to_dataset_index = { - 'FL_hip_joint': 0, - 'FL_thigh_joint': 1, - 'FL_calf_joint': 2, - 'FR_hip_joint': 3, - 'FR_thigh_joint': 4, - 'FR_calf_joint': 5, - 'RL_hip_joint': 6, - 'RL_thigh_joint': 7, - 'RL_calf_joint': 8, - 'RR_hip_joint': 9, - 'RR_thigh_joint': 10, - 'RR_calf_joint': 11, - 'FL_foot_fixed': 0, - 'FR_foot_fixed': 1, - 'RL_foot_fixed': 2, - 'RR_foot_fixed': 3 + 'FR_hip_joint': 0, + 'FR_thigh_joint': 1, + 'FR_calf_joint': 2, + 'FL_hip_joint': 3, + 'FL_thigh_joint': 4, + 'FL_calf_joint': 5, + 'RR_hip_joint': 6, + 'RR_thigh_joint': 7, + 'RR_calf_joint': 8, + 'RL_hip_joint': 9, + 'RL_thigh_joint': 10, + 'RL_calf_joint': 11, + 'FR_foot_fixed': 0, + 'FL_foot_fixed': 1, + 'RR_foot_fixed': 2, + 'RL_foot_fixed': 3, } # Define the names and indicies that contain ground truth labels self.ground_truth_urdf_names = [ - 'FL_foot_fixed', 'FR_foot_fixed', 'RL_foot_fixed', 'RR_foot_fixed' + 'FR_foot_fixed', + 'FL_foot_fixed', + 'RR_foot_fixed', + 'RL_foot_fixed', ] - self.ground_truth_indices = [] + self.ground_truth_graph_indices = [] for urdf_name in self.ground_truth_urdf_names: - self.ground_truth_indices.append( + self.ground_truth_graph_indices.append( self.urdf_name_to_index[urdf_name]) + self.ground_truth_array_indices = [] + for urdf_name in self.ground_truth_urdf_names: + self.ground_truth_array_indices.append( + self.urdf_to_dataset_index[urdf_name]) + # Define the nodes that should recieve features self.nodes_for_attributes = [ 'FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint', 'FR_hip_joint', @@ -299,8 +310,6 @@ def process(self): f.write('\n') curr_seq_num += 1 - print(curr_seq_num) - # Write a txt file to save the dataset length & and first sequence index length = curr_seq_num with open(os.path.join(self.processed_dir, "info.txt"), "w") as f: @@ -330,7 +339,12 @@ def load_data_at_ros_seq(self, ros_seq: int): for i in range(0, len(line)): torques.append(float(line[i])) - return positions, velocities, torques, grfs + # Extract the ground truth labels + ground_truth_labels = [] + for val in self.ground_truth_array_indices: + ground_truth_labels.append(grfs[val]) + + return positions, velocities, torques, ground_truth_labels def get_helper_gnn(self, idx): # Load the rosbag information @@ -370,9 +384,56 @@ def get_helper_mlp(self, idx): x = torch.tensor((positions + velocities + torques), dtype=torch.float) # Create the ground truth lables - y = torch.tensor(ground_truth_labels, dtype=torch.float) + y = torch.tensor(ground_truth_labels[self.ground_truth_graph_indices], + dtype=torch.float) return x, y + def get_helper_heterogeneous_gnn(self, idx): + # Create the Heterogeneous Data object + data = HeteroData() + + # Get the edge matrices + bj, jb, jj, fj, jf = self.URDF.get_edge_index_matrices() + data['base', 'connect', 'joint'].edge_index = bj + data['joint', 'connect', 'base'].edge_index = jb + data['joint', 'connect', 'joint'].edge_index = jj + data['foot', 'connect', 'joint'].edge_index = fj + data['joint', 'connect', 'foot'].edge_index = jf + + # Load the rosbag information + positions, velocities, torques, ground_truth_labels = self.load_data_at_ros_seq( + self.first_index + idx) + + # Save the labels and number of nodes + data.y = torch.tensor(ground_truth_labels, dtype=torch.float) + data.num_nodes = self.URDF.get_num_nodes() + + # Create the feature matrices + number_nodes = self.URDF.get_num_of_each_node_type() + base_x = torch.ones((number_nodes[0], 0), dtype=torch.float) + joint_x = torch.ones((number_nodes[1], 3), dtype=torch.float) + foot_x = torch.ones((number_nodes[2], 0), dtype=torch.float) + + # For each joint specified + for urdf_node_name in self.nodes_for_attributes: + + # Find the index of this particular node + node_index = self.urdf_name_to_index[urdf_node_name] + + # Get the msg array index + msg_ind = self.urdf_to_dataset_index[urdf_node_name] + + # Add the features to x matrix + joint_x[node_index, 0] = positions[msg_ind] + joint_x[node_index, 1] = velocities[msg_ind] + joint_x[node_index, 2] = torques[msg_ind] + + # Save the matrices into the HeteroData object + data['base'].x = base_x # [num_papers, num_features_paper] + data['joint'].x = joint_x # [num_authors, num_features_author] + data['foot'].x = foot_x # [num_institutions, num_features_institution] + return data + class CerberusDataset(FlexibleDataset): """ @@ -386,7 +447,7 @@ class CerberusDataset(FlexibleDataset): def __init__(self, root: str, - robotURDF: RobotURDF, + robotURDF: NormalRobotGraph, data_format: str = 'gnn', transform=None, pre_transform=None, @@ -425,6 +486,11 @@ def __init__(self, 'RR_foot_fixed': 'RR_foot' } + # Open up the urdf file and get the edge matrix + self.edge_matrix = self.URDF.get_edge_index_matrix() + self.edge_matrix_tensor = torch.tensor(self.edge_matrix, + dtype=torch.long) + # Define the names and indicies that contain ground truth labels self.ground_truth_urdf_names = [ 'FL_foot_fixed', 'FR_foot_fixed', 'RL_foot_fixed', 'RR_foot_fixed' @@ -434,11 +500,16 @@ def __init__(self, for urdf_name in self.ground_truth_urdf_names: self.ground_truth_ros_names.append(self.urdf_to_ros_map[urdf_name]) - self.ground_truth_indices = [] + self.ground_truth_graph_indices = [] for urdf_name in self.ground_truth_urdf_names: - self.ground_truth_indices.append( + self.ground_truth_graph_indices.append( self.urdf_name_to_index[urdf_name]) + self.ground_truth_array_indices = [] + for ros_name in self.ground_truth_ros_names: + self.ground_truth_array_indices.append( + self.get_index_of_ros_name(ros_name)) + # Define the nodes that should recieve features self.nodes_for_attributes = [ 'FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint', 'FR_hip_joint', @@ -530,10 +601,10 @@ def load_data_at_ros_seq(self, ros_seq: int): efforts.append(float(line[i])) # Get the ground truth force labels + # TODO: Write these methods for these functions ground_truth_labels = [] - for name in self.ground_truth_ros_names: - ground_truth_labels.append( - efforts[self.get_index_of_ros_name(name)]) + for val in self.ground_truth_array_indices: + ground_truth_labels.append(efforts[val]) return positions, velocities, efforts, ground_truth_labels @@ -542,7 +613,7 @@ def get_helper_gnn(self, idx): positions, velocities, efforts, ground_truth_labels = self.load_data_at_ros_seq( self.first_index + idx) - # Create a note feature matrix + # Create a node feature matrix x = torch.ones((self.URDF.get_num_nodes(), 2), dtype=torch.float) # For each joint specified @@ -644,8 +715,8 @@ def get_start_and_end_seq_ids(self): return 119, 174560 -def visualize_graph(graph: Data, - urdf: RobotURDF, +def visualize_graph(pytorch_graph: Data, + robot_graph: NormalRobotGraph, fig_save_path: Path = None, draw_edges: bool = False): """ @@ -654,14 +725,15 @@ def visualize_graph(graph: Data, """ # Write the features onto the names - node_labels = urdf.get_node_index_to_name_dict() - for i in range(0, len(graph.x)): + node_labels = robot_graph.get_node_index_to_name_dict() + for i in range(0, len(pytorch_graph.x)): label = node_labels[i] - label += ": " + str(graph.x[i].numpy()) + label += ": " + str(pytorch_graph.x[i].numpy()) node_labels[i] = label # Convert to networkx graph - nx_graph = torch_geometric.utils.to_networkx(graph, to_undirected=True) + nx_graph = torch_geometric.utils.to_networkx(pytorch_graph, + to_undirected=True) # Draw the graph spring_layout = networkx.spring_layout(nx_graph) @@ -675,7 +747,7 @@ def visualize_graph(graph: Data, networkx.draw_networkx_edge_labels( nx_graph, pos=spring_layout, - edge_labels=urdf.get_edge_connections_to_name_dict(), + edge_labels=robot_graph.get_edge_connections_to_name_dict(), rotate=False, font_size=7) diff --git a/src/grfgnn/gnnLightning.py b/src/grfgnn/gnnLightning.py index 3363db6..09493ca 100644 --- a/src/grfgnn/gnnLightning.py +++ b/src/grfgnn/gnnLightning.py @@ -5,7 +5,7 @@ from torch_geometric.nn.models import GCN from lightning.pytorch.loggers import WandbLogger from .datasets import CerberusStreetDataset, CerberusTrackDataset, CerberusDataset -from .urdfParser import RobotURDF +from .graphParser import NormalRobotGraph from torch_geometric.loader import DataLoader from lightning.pytorch.callbacks import ModelCheckpoint from pathlib import Path @@ -231,7 +231,7 @@ def train_model(path_to_urdf: Path, path_to_cerberus_street: Path, path_to_cerberus_track: Path, model_type: str = 'gnn'): # Load the A1 urdf - A1_URDF = RobotURDF(path_to_urdf, 'package://a1_description/', + A1_URDF = NormalRobotGraph(path_to_urdf, 'package://a1_description/', 'unitree_ros/robots/a1_description', True) # Initalize the datasets diff --git a/src/grfgnn/urdfParser.py b/src/grfgnn/graphParser.py similarity index 70% rename from src/grfgnn/urdfParser.py rename to src/grfgnn/graphParser.py index b1fecfc..05276ba 100644 --- a/src/grfgnn/urdfParser.py +++ b/src/grfgnn/graphParser.py @@ -8,7 +8,7 @@ class InvalidURDFException(Exception): pass -class RobotURDF(): +class RobotGraph(): class Node(): """ @@ -22,6 +22,13 @@ def __init__(self, name: str, edge_parent: str, self.edge_parent: str = edge_parent self.edge_children: list[str] = edge_children + @staticmethod + def get_list_of_node_types(): + """ + Returns a list of all the node types. + """ + return ['base', 'joint', 'foot'] + def get_node_type(self): """ Returns the name of the node type for a heterogenous graph, @@ -30,13 +37,14 @@ def get_node_type(self): has a parent, it's a 'foot' node. If it only has children, it's a 'base' node. """ + node_types = self.get_list_of_node_types() if self.edge_parent is not None and len(self.edge_children) > 0: - return 'joint' + return node_types[1] elif self.edge_parent is not None and len(self.edge_children) == 0: - return 'foot' + return node_types[2] elif self.edge_parent is None and len(self.edge_children) > 0: - return 'base' + return node_types[0] else: raise Exception( "Every node should have a child or parent edge.") @@ -58,10 +66,10 @@ def __init__(self, urdf_to_desc_path: str, swap_nodes_and_edges=False): """ - Constructor for RobotURDF class. + Constructor for RobotGraph class. Args: - urdf_path (Path): The absolute path from this file (urdfParser.py) + urdf_path (Path): The absolute path from this file (graphParser.py) to the desired urdf file to load. ros_builtin_path (str): The path ROS uses in the urdf file to navigate to the urdf description directory. An example looks like this: @@ -70,7 +78,7 @@ def __init__(self, urdf_to_desc_path (str): The relative path from the urdf file to the urdf description directory. This directory typically contains folders like "meshes" and "urdf". - swap_nodes_and_edges (bool): If False, links in the URDF will be represented + swap_nodes_and_edges (bool): If False, links in the urdf file will be represented as nodes, and joints will be represented as edges. If True, links will be represented as edges, and joints will be represented as nodes. This will drop links that aren't connected to two joints. Default=False. @@ -131,7 +139,7 @@ def __init__(self, self.Edge(name=link.name + "_to_" + edge_child, parent=edge_parent, child=edge_child)) - + # Create nodes from the Joints self.nodes = [] for joint in self.robot_urdf.joints: @@ -231,9 +239,70 @@ def get_node_name_to_index_dict(self): index. Returns: - (dict[str, int]): A dictinoary that maps node name + (dict[str, int]): A dictionary that maps node name to index. """ + + raise NotImplementedError + + def get_num_nodes(self): + """ + Return the number of nodes in the graph. + + Returns: + (int): Number of node nodes in graph. + """ + return len(self.nodes) + + def get_node_from_name(self, node_name): + """ + Helper function that returns a node based + on the name provided. If no node with the + name is found, None is returned. + + Parameters: + node_name (str): Name of the node to return. + """ + + for node in self.nodes: + if node.name == node_name: + return node + else: + return None + + def display_graph_info(self): + """ + Helper function that displays information about the robot + graph on the command line in a readable format. + """ + + print("============ Displaying Robot Links: ============") + print("Link name: Mass - Inertia") + for link in self.robot_urdf.links: + print(f"{link.name}") + print("") + + print("============ Displaying Example Link: ============") + ex_link = self.robot_urdf.links[7] + print("Name: ", ex_link.name) + print("Mass: ", ex_link.inertial.mass) + print("Inertia: \n", ex_link.inertial.inertia) + print("Origin of Inertials: \n", ex_link.inertial.origin) + print("") + + print("============ Displaying Robot Joints: ============") + for joint in self.robot_urdf.joints: + print('{} -> {} <- {}'.format(joint.parent, joint.name, + joint.child)) + print("") + + +class NormalRobotGraph(RobotGraph): + """ + Graph where all of the nodes are homogeneous. + """ + + def get_node_name_to_index_dict(self): node_names = [] for node in self.nodes: node_names.append(node.name) @@ -242,7 +311,7 @@ def get_node_name_to_index_dict(self): def get_node_index_to_name_dict(self): """ Return a dictionary that maps the node index to its - name + name. Returns: (dict[int, str]): A dictionary that maps node index @@ -255,7 +324,7 @@ def get_node_index_to_name_dict(self): node_dict = dict(zip(range(len(self.nodes)), node_names)) return node_dict - def get_edge_index_matrix(self): + def get_edge_index_matrix(self, heterogenous=False): """ Return the edge connectivity matrix, which defines each edge connection to each node. This matrix is the 'edge_index' matrix passed to the PyTorch @@ -267,7 +336,7 @@ def get_edge_index_matrix(self): be put into the matrix twice, as each edge is bidirectional For example, a connection between nodes 0 and 1 will be found in the edge_index matrix as [[0, 1], [1, 0]]. - + """ node_dict = self.get_node_name_to_index_dict() @@ -285,15 +354,6 @@ def get_edge_index_matrix(self): return edge_matrix - def get_num_nodes(self): - """ - Return the number of nodes in the URDF file. - - Returns: - (int): Number of node nodes in URDF. - """ - return len(self.nodes) - def get_edge_connections_to_name_dict(self): """ Return a dictionary that maps a tuple of two node indices @@ -325,7 +385,7 @@ def get_edge_name_to_connections_dict(self): Returns: edge_dict (dict[str, np.array)]): - This dictionary takes the name of an edge in the URDF + This dictionary takes the name of an edge in the graph as input. It returns an np.array (N, 2), where N is the number of connections this name supports. """ @@ -341,28 +401,121 @@ def get_edge_name_to_connections_dict(self): return edge_dict - def display_URDF_info(self): + +class HeterogeneousRobotGraph(RobotGraph): + """ + Heterogeneous graph, where the nodes are different types. + """ + + def _get_nodes_organized_by_type(self): """ - Helper function that displays information about the robot - URDF on the command line in a readable format. + Organizes each node into a type with nodes of its same + type. + + Returns: + list[list[Node]]: The outer list holds as many lists + as there are types, and each inner list holds + all of the nodes of that particular type. """ + types = RobotGraph.Node.get_list_of_node_types() + nodes = [] + for type in types: + nodes_of_type = [] + for node in self.nodes: + if node.get_node_type() == type: + nodes_of_type.append(node.name) + nodes.append(nodes_of_type) + return nodes - print("============ Displaying Robot Links: ============") - print("Link name: Mass - Inertia") - for link in self.robot_urdf.links: - print(f"{link.name}") - print("") + def get_node_name_to_index_dict(self): + """ + Because this graph is heterogenous, the nodes index + depends of the type of the node. This also means that + multiple nodes can have the same index. + """ - print("============ Displaying Example Link: ============") - ex_link = self.robot_urdf.links[7] - print("Name: ", ex_link.name) - print("Mass: ", ex_link.inertial.mass) - print("Inertia: \n", ex_link.inertial.inertia) - print("Origin of Inertials: \n", ex_link.inertial.origin) - print("") + all_lists = [] + for nodes_of_type in self._get_nodes_organized_by_type(): + all_lists = all_lists + list( + zip(nodes_of_type, range(len(nodes_of_type)))) - print("============ Displaying Robot Joints: ============") - for joint in self.robot_urdf.joints: - print('{} -> {} <- {}'.format(joint.parent, joint.name, - joint.child)) - print("") + return dict(all_lists) + + def get_num_of_each_node_type(self): + """ + Returns the numbers of each node type. + + Returns: + list[int]: List of number of each node type + """ + nodes = self._get_nodes_organized_by_type() + numbers = [] + for node_list in nodes: + numbers.append(len(node_list)) + return numbers + + def get_edge_index_matrices(self): + """ + Return the edge connectivity matrices. + + Returns: + (list[np.array]): Multiple matrices, as outlined below: + data['base', 'connect', 'joint'].edge_index -> [2, X] + data['joint', 'connect', 'base'].edge_index -> [2, X] + data['joint', 'connect', 'joint'].edge_index -> [2, Y] + data['foot', 'connect', 'joint'].edge_index -> [2, Z] + data['joint', 'connect', 'foot'].edge_index -> [2, Z] + """ + + def add_to_matrix(matrix, vector): + """ + Helper function for appending to a numpy + matrix along the 1st axis. + """ + if matrix is None: + matrix = vector + else: + matrix = np.concatenate((matrix, vector), axis=1) + return matrix + + # Get the name to index dictionary + node_dict = self.get_node_name_to_index_dict() + + # Define all of the edge matrices + base_to_joint_matrix = None + joint_to_joint_matrix = None + foot_to_joint_matrix = None + + # Iterate through edges + for edge in self.edges: + + # Get the nodes for the parent and the child + parent_node: RobotGraph.Node = self.get_node_from_name(edge.parent) + child_node: RobotGraph.Node = self.get_node_from_name(edge.child) + + # Get their types and indices + parent_type = parent_node.get_node_type() + child_type = child_node.get_node_type() + p_index = node_dict[edge.parent] + c_index = node_dict[edge.child] + + # Add their info to the corresponding matrix + if parent_type == child_type and parent_type == 'joint': + edge_vector = np.array([[p_index, c_index], [c_index, + p_index]]) + joint_to_joint_matrix = add_to_matrix(joint_to_joint_matrix, edge_vector) + elif parent_type == 'base' and child_type == 'joint': + edge_vector = np.array([[p_index], [c_index]]) + base_to_joint_matrix = add_to_matrix(base_to_joint_matrix, edge_vector) + elif parent_type == 'joint' and child_type == 'foot': + edge_vector = np.array([[c_index], [p_index]]) + foot_to_joint_matrix = add_to_matrix(foot_to_joint_matrix, edge_vector) + else: + raise Exception("Not possible") + + # Create the last two matrices + joint_to_base_matrix = base_to_joint_matrix[[1, 0]] + joint_to_foot_matrix = foot_to_joint_matrix[[1, 0]] + + return base_to_joint_matrix, joint_to_base_matrix, joint_to_joint_matrix, \ + foot_to_joint_matrix, joint_to_foot_matrix diff --git a/tests/testDatasets.py b/tests/testDatasets.py index 9e2f15c..63cc7cf 100644 --- a/tests/testDatasets.py +++ b/tests/testDatasets.py @@ -1,11 +1,13 @@ import unittest from pathlib import Path -from grfgnn import RobotURDF, CerberusStreetDataset, CerberusTrackDataset, CerberusCampusDataset, Go1SimulatedDataset, FlexibleDataset +from grfgnn import NormalRobotGraph, HeterogeneousRobotGraph, CerberusStreetDataset, CerberusTrackDataset, CerberusCampusDataset, Go1SimulatedDataset, FlexibleDataset from rosbags.highlevel import AnyReader import os +from torch_geometric.data import HeteroData +import numpy as np -class TestDatasets(unittest.TestCase): +class TestCerberusDatasets(unittest.TestCase): """ Test that the Cerberus Dataset classes download, process, and read the info for creating graph datasets. We @@ -19,10 +21,14 @@ def setUp(self): Path(__file__).cwd(), 'urdf_files', 'A1', 'a1.urdf').absolute() self.go1_path = Path( Path(__file__).cwd(), 'urdf_files', 'Go1', 'go1.urdf').absolute() - self.A1_URDF = RobotURDF(self.a1_path, 'package://a1_description/', - 'unitree_ros/robots/a1_description', True) - self.GO1_URDF = RobotURDF(self.go1_path, 'package://go1_description/', - 'unitree_ros/robots/go1_description', True) + self.A1_URDF = NormalRobotGraph(self.a1_path, + 'package://a1_description/', + 'unitree_ros/robots/a1_description', + True) + self.GO1_URDF = NormalRobotGraph(self.go1_path, + 'package://go1_description/', + 'unitree_ros/robots/go1_description', + True) # Create the street dataset self.dataset_street_path = Path( @@ -203,8 +209,9 @@ def test_error_messages(self): with self.assertRaises(ValueError) as error: temp = CerberusStreetDataset(self.dataset_street_path, self.GO1_URDF, 'fake') - self.assertEqual("Parameter 'data_format' must be 'gnn' or 'mlp'.", - str(error.exception)) + self.assertEqual( + "Parameter 'data_format' must be 'gnn', 'mlp', or 'heterogeneous_gnn'.", + str(error.exception)) # Make sure we get an error when we initialize a class that isn't meant # to be initialized. @@ -223,28 +230,40 @@ def test_get_expected_urdf_name(self): self.assertEqual("a1", self.dataset_street.get_expected_urdf_name()) self.assertEqual("a1", self.dataset_track.get_expected_urdf_name()) - @unittest.skipIf( - os.getenv('RUN_LOCAL_TESTS', 'False') != 'True', - "Not running tests on datasets that can't be downloaded") - def test_go1_simulated_dataset(self): - """ - Make sure the data from the 100 rosbags are correctly given sequence - numbers for the Go1 Simulated dataset. - """ + # TODO: Add test cases for get() method + # TODO: Add test cases for new dataset class methods + +@unittest.skipIf( + os.getenv('RUN_LOCAL_TESTS', 'False') != 'True', + "Not running tests on datasets that can't be downloaded") +class TestGo1SimulatedDataset(unittest.TestCase): + """ + Test that the Go1 Simulated dataset successfully + processes and reads the info for creating graph datasets. + """ + + def setUp(self): # Set up the Go1 Simulated dataset path_to_go1_urdf = Path( Path('.').parent, 'urdf_files', 'Go1', 'go1.urdf').absolute() path_to_xiong_simulated = Path( Path('.').parent, 'datasets', 'xiong_simulated').absolute() - GO1_URDF = RobotURDF(path_to_go1_urdf, 'package://go1_description/', - 'unitree_ros/robots/go1_description', True) + self.GO1_graph = HeterogeneousRobotGraph( + path_to_go1_urdf, 'package://go1_description/', + 'unitree_ros/robots/go1_description', True) model_type = 'gnn' - go1_sim_dataset = Go1SimulatedDataset(path_to_xiong_simulated, - GO1_URDF, model_type) + self.go1_sim_dataset = Go1SimulatedDataset(path_to_xiong_simulated, + self.GO1_graph, model_type) + + def test_load_data_at_ros_seq(self): + """ + Make sure the data from the 100 rosbags are correctly given sequence + numbers for the Go1 Simulated dataset. + """ # Make sure data is loaded properly - p, v, t, gt = go1_sim_dataset.load_data_at_ros_seq(181916) + p, v, t, gt = self.go1_sim_dataset.load_data_at_ros_seq(181916) des_p = [ 0.09840758, 1.2600079, -2.0489328, -0.10457229, 0.7454401, -1.8136898, 0.04816602, 0.91943306, -1.9169945, -0.0340704, @@ -265,8 +284,68 @@ def test_go1_simulated_dataset(self): self.assertSequenceEqual(t, des_t) self.assertSequenceEqual(gt, des_gt) - # TODO: Add test cases for get() method - # TODO: Add test cases for new dataset class methods + def test_get_helper_heterogeneous_gnn(self): + # Get the HeteroData graph + heteroData: HeteroData = self.go1_sim_dataset.get_helper_heterogeneous_gnn( + 181916) + # Get the desired edge matrices + bj, jb, jj, fj, jf = self.GO1_graph.get_edge_index_matrices() + + # Make sure they match + np.testing.assert_array_equal( + heteroData['base', 'connect', 'joint'].edge_index, bj) + np.testing.assert_array_equal( + heteroData['joint', 'connect', 'base'].edge_index, jb) + np.testing.assert_array_equal( + heteroData['joint', 'connect', 'joint'].edge_index, jj) + np.testing.assert_array_equal( + heteroData['foot', 'connect', 'joint'].edge_index, fj) + np.testing.assert_array_equal( + heteroData['joint', 'connect', 'foot'].edge_index, jf) + + # Check the labels + labels_des = [-0.0063045006, 55.528183, 83.40855, -0.006935571] + np.testing.assert_array_equal(heteroData.y.numpy(), + np.array(labels_des, dtype=np.float32)) + + # Check the number of nodes + number_des = self.GO1_graph.get_num_nodes() + self.assertEqual(heteroData.num_nodes, number_des) + + des_p = [ + 0.09840758, 1.2600079, -2.0489328, -0.10457229, 0.7454401, + -1.8136898, 0.04816602, 0.91943306, -1.9169945, -0.0340704, + 1.2415031, -2.0689785 + ] + des_v = [ + 0.6230268, -5.9481835, -5.3682613, -0.26720884, 6.455389, + -1.3378538, -0.00086710247, 6.2338834, -0.5447279, 0.6295713, + -4.582517, -7.406303 + ] + des_t = [ + -0.4899872, -3.0470033, 0.363765, 5.630082, 5.196147, 11.241279, + -1.8939179, 4.083751, 16.447073, -1.3105631, -0.7087057, 0.13933142 + ] + + # Check the node attributes + base_x = np.array([[]], dtype=np.float32) + joint_x = np.array([[0.09840758, 0.6230268, -0.4899872], + [1.2600079, -5.9481835, -3.0470033], + [-2.0489328, -5.3682613, 0.363765], + [-0.10457229, -0.26720884, 5.630082], + [0.7454401, 6.455389, 5.196147], + [-1.8136898, -1.3378538, 11.241279], + [0.04816602, -0.00086710247, -1.8939179], + [0.91943306, 6.2338834, 4.083751], + [-1.9169945, -0.5447279, 16.447073], + [-0.0340704, 0.6295713, -1.3105631], + [1.2415031, -4.582517, -0.7087057], + [-2.0689785, -7.406303, 0.13933142]], + dtype=np.float32) + foot_x = np.array([[], [], [], []], dtype=np.float32) + np.testing.assert_array_equal(heteroData['base'].x.numpy(), base_x) + np.testing.assert_array_equal(heteroData['joint'].x.numpy(), joint_x) + np.testing.assert_array_equal(heteroData['foot'].x.numpy(), foot_x) if __name__ == '__main__': diff --git a/tests/testUrdfParser.py b/tests/testGraphParser.py similarity index 67% rename from tests/testUrdfParser.py rename to tests/testGraphParser.py index fd70f4a..6a5cc93 100644 --- a/tests/testUrdfParser.py +++ b/tests/testGraphParser.py @@ -1,25 +1,27 @@ from pathlib import Path import unittest -from grfgnn import RobotURDF +from grfgnn import RobotGraph, NormalRobotGraph, HeterogeneousRobotGraph from pathlib import Path import copy import pandas as pd import numpy as np import os import urchin +import numpy -class TestRobotURDF(unittest.TestCase): +class TestNormalRobotGraph(unittest.TestCase): def setUp(self): self.hyq_path = Path( Path(__file__).cwd(), 'urdf_files', 'HyQ', 'hyq.urdf').absolute() - self.HyQ_URDF = RobotURDF(self.hyq_path, 'package://hyq_description/', - 'hyq-description', False) - self.HyQ_URDF_swapped = RobotURDF(self.hyq_path, - 'package://hyq_description/', - 'hyq-description', True) + self.HyQ_URDF = NormalRobotGraph(self.hyq_path, + 'package://hyq_description/', + 'hyq-description', False) + self.HyQ_URDF_swapped = NormalRobotGraph(self.hyq_path, + 'package://hyq_description/', + 'hyq-description', True) def test_constructor(self): """ @@ -69,26 +71,26 @@ def test_constructor(self): # Additionally, links with multiple children joints get one # edge for each child. desired_edges = [ - RobotURDF.Edge('trunk_to_lf_haa_joint', "floating_base", - "lf_haa_joint"), - RobotURDF.Edge('trunk_to_lh_haa_joint', "floating_base", - "lh_haa_joint"), - RobotURDF.Edge('trunk_to_rf_haa_joint', "floating_base", - "rf_haa_joint"), - RobotURDF.Edge('trunk_to_rh_haa_joint', "floating_base", - "rh_haa_joint"), - RobotURDF.Edge('lf_hipassembly', "lf_haa_joint", "lf_hfe_joint"), - RobotURDF.Edge('lf_upperleg', "lf_hfe_joint", "lf_kfe_joint"), - RobotURDF.Edge('lf_lowerleg', "lf_kfe_joint", "lf_foot_joint"), - RobotURDF.Edge('rf_hipassembly', "rf_haa_joint", "rf_hfe_joint"), - RobotURDF.Edge('rf_upperleg', "rf_hfe_joint", "rf_kfe_joint"), - RobotURDF.Edge('rf_lowerleg', "rf_kfe_joint", "rf_foot_joint"), - RobotURDF.Edge('lh_hipassembly', "lh_haa_joint", "lh_hfe_joint"), - RobotURDF.Edge('lh_upperleg', "lh_hfe_joint", "lh_kfe_joint"), - RobotURDF.Edge('lh_lowerleg', "lh_kfe_joint", "lh_foot_joint"), - RobotURDF.Edge('rh_hipassembly', "rh_haa_joint", "rh_hfe_joint"), - RobotURDF.Edge('rh_upperleg', "rh_hfe_joint", "rh_kfe_joint"), - RobotURDF.Edge('rh_lowerleg', "rh_kfe_joint", "rh_foot_joint") + RobotGraph.Edge('trunk_to_lf_haa_joint', "floating_base", + "lf_haa_joint"), + RobotGraph.Edge('trunk_to_lh_haa_joint', "floating_base", + "lh_haa_joint"), + RobotGraph.Edge('trunk_to_rf_haa_joint', "floating_base", + "rf_haa_joint"), + RobotGraph.Edge('trunk_to_rh_haa_joint', "floating_base", + "rh_haa_joint"), + RobotGraph.Edge('lf_hipassembly', "lf_haa_joint", "lf_hfe_joint"), + RobotGraph.Edge('lf_upperleg', "lf_hfe_joint", "lf_kfe_joint"), + RobotGraph.Edge('lf_lowerleg', "lf_kfe_joint", "lf_foot_joint"), + RobotGraph.Edge('rf_hipassembly', "rf_haa_joint", "rf_hfe_joint"), + RobotGraph.Edge('rf_upperleg', "rf_hfe_joint", "rf_kfe_joint"), + RobotGraph.Edge('rf_lowerleg', "rf_kfe_joint", "rf_foot_joint"), + RobotGraph.Edge('lh_hipassembly', "lh_haa_joint", "lh_hfe_joint"), + RobotGraph.Edge('lh_upperleg', "lh_hfe_joint", "lh_kfe_joint"), + RobotGraph.Edge('lh_lowerleg', "lh_kfe_joint", "lh_foot_joint"), + RobotGraph.Edge('rh_hipassembly', "rh_haa_joint", "rh_hfe_joint"), + RobotGraph.Edge('rh_upperleg', "rh_hfe_joint", "rh_kfe_joint"), + RobotGraph.Edge('rh_lowerleg', "rh_kfe_joint", "rh_foot_joint") ] for i, edge in enumerate(self.HyQ_URDF_swapped.edges): match_found = False @@ -134,11 +136,6 @@ def test_constructor(self): for i, node in enumerate(self.HyQ_URDF_swapped.nodes): for j, node_des in enumerate(edge_names_copy): if (node.name == node_des): - print("\nNode name: ", node.name) - print("Node Des: ", node_des) - print("desired type: ", des_node_type[j]) - print("Node Parent: ", node.edge_parent) - print("Node children: ", node.edge_children) self.assertEqual(node.get_node_type(), des_node_type[j]) num_matches += 1 break @@ -178,8 +175,8 @@ def test_create_updated_urdf_file(self): self.assertFalse(os.path.exists(hyq_path_updated)) # Rebuild it - RobotURDF(self.hyq_path, 'package://hyq_description/', - 'hyq-description', False) + RobotGraph(self.hyq_path, 'package://hyq_description/', + 'hyq-description', False) self.assertTrue(os.path.exists(hyq_path_updated)) def test_get_node_name_to_index_dict(self): @@ -284,5 +281,73 @@ def test_get_edge_name_to_connections_dict(self): seen_arrays.add(array_tuple) +class TestHeterogeneousRobotGraph(unittest.TestCase): + + def setUp(self): + self.path_to_go1_urdf = Path( + Path('.').parent, 'urdf_files', 'Go1', 'go1.urdf').absolute() + + self.GO1_HETERO_GRAPH = HeterogeneousRobotGraph( + self.path_to_go1_urdf, 'package://go1_description/', + 'unitree_ros/robots/go1_description', True) + + def test_get_node_name_to_index_dict(self): + """ + Test that the dictionary properly assigns indices to the nodes. + """ + + dict_actual = self.GO1_HETERO_GRAPH.get_node_name_to_index_dict() + dict_desired = { + 'floating_base': 0, + 'FR_hip_joint': 0, + 'FR_thigh_joint': 1, + 'FR_calf_joint': 2, + 'FL_hip_joint': 3, + 'FL_thigh_joint': 4, + 'FL_calf_joint': 5, + 'RR_hip_joint': 6, + 'RR_thigh_joint': 7, + 'RR_calf_joint': 8, + 'RL_hip_joint': 9, + 'RL_thigh_joint': 10, + 'RL_calf_joint': 11, + 'FR_foot_fixed': 0, + 'FL_foot_fixed': 1, + 'RR_foot_fixed': 2, + 'RL_foot_fixed': 3 + } + self.assertDictEqual(dict_actual, dict_desired) + + def test_get_num_of_each_node_type(self): + """ + Test that we can properly count the number of each + type of node. + """ + + number_actual = self.GO1_HETERO_GRAPH.get_num_of_each_node_type() + number_desired = [1, 12, 4] + self.assertSequenceEqual(number_actual, number_desired) + + def test_get_edge_index_matrices(self): + """ + Test that we construct the correct matrices for a + heterogeneous graph. + """ + + bj, jb, jj, fj, jf = self.GO1_HETERO_GRAPH.get_edge_index_matrices() + bj_des = np.array([[0, 0, 0, 0], [0, 3, 6, 9]]) + jb_des = np.array([[0, 3, 6, 9], [0, 0, 0, 0]]) + jj_des = np.array([[0, 1, 1, 2, 3, 4, 4, 5, 6, 7, 7, 8, 9, 10, 10, 11], + [1, 0, 2, 1, 4, 3, 5, 4, 7, 6, 8, 7, 10, 9, 11, + 10]]) + fj_des = np.array([[0, 1, 2, 3], [2, 5, 8, 11]]) + jf_des = np.array([[2, 5, 8, 11], [0, 1, 2, 3]]) + numpy.testing.assert_array_equal(bj, bj_des) + numpy.testing.assert_array_equal(jb, jb_des) + numpy.testing.assert_array_equal(jj, jj_des) + numpy.testing.assert_array_equal(fj, fj_des) + numpy.testing.assert_array_equal(jf, jf_des) + + if __name__ == '__main__': unittest.main()