Research files changed

research/animations.py

@@ -0,0 +1,17 @@
+from manim import *
+from grfgnn import RobotURDF
+
+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)
+
+        # 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))

research/evaluator.py

@@ -6,7 +6,7 @@
 path_to_urdf = Path(Path('.').parent, 'urdf_files', 'A1', 'a1.urdf').absolute()
 path_to_cerberus_street = Path(Path('.').parent, 'datasets', 'cerberus_street').absolute()
 path_to_cerberus_track = Path(Path('.').parent, 'datasets', 'cerberus_track').absolute()
-path_to_checkpoint = Path(Path('.').parent, 'models', 'mlp-Glenn-Wallace', 'epoch=6-val_loss=5764.05.ckpt')
+path_to_checkpoint = Path(Path('.').parent, 'models', TODO, 'epoch=6-val_loss=5764.05.ckpt')
 
 # Load the A1 urdf
 A1_URDF = RobotURDF(path_to_urdf, 'package://a1_description/',
@@ -21,6 +21,6 @@
     path_to_cerberus_track,
     A1_URDF, model_type)
 
-visualize_model_outputs('mlp', path_to_checkpoint, path_to_urdf, 
+visualize_model_outputs(model_type, path_to_checkpoint, path_to_urdf, 
                         path_to_cerberus_track, 10000,
-                        Path('.', 'research', 'pred_new_mlp.png'))
+                        Path('.', 'research', 'pred_new_' + model_type + '.png'))