Skip to content

Commit

Permalink
HeterogeneousRobotGraph & Heterogeneous Dataset for Go1 Simulated Data
Browse files Browse the repository at this point in the history
  • Loading branch information
DanielChaseButterfield committed May 4, 2024
1 parent b3875b2 commit 5b9eea6
Show file tree
Hide file tree
Showing 10 changed files with 552 additions and 175 deletions.
13 changes: 8 additions & 5 deletions research/animations.py
Original file line number Diff line number Diff line change
@@ -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))
self.play(Write(text))
4 changes: 2 additions & 2 deletions research/devel.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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'

Expand Down
6 changes: 3 additions & 3 deletions research/evaluator.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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'

Expand Down
21 changes: 13 additions & 8 deletions research/urdfTest.py
Original file line number Diff line number Diff line change
@@ -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()
2 changes: 1 addition & 1 deletion src/grfgnn/__init__.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
from .urdfParser import RobotURDF
from .graphParser import RobotGraph, NormalRobotGraph, HeterogeneousRobotGraph
from .datasets import CerberusStreetDataset, CerberusTrackDataset, CerberusCampusDataset, Go1SimulatedDataset, FlexibleDataset
Loading

0 comments on commit 5b9eea6

Please sign in to comment.