From 184b145a3d91c6ea7c999335e205094cd57bf0ef Mon Sep 17 00:00:00 2001
From: Daniel Butterfield <danielcbutter@gmail.com>
Date: Tue, 16 Apr 2024 17:45:53 -0400
Subject: [PATCH] Research files changed

---
 research/animations.py | 17 +++++++++++++++++
 research/evaluator.py  |  6 +++---
 2 files changed, 20 insertions(+), 3 deletions(-)
 create mode 100644 research/animations.py

diff --git a/research/animations.py b/research/animations.py
new file mode 100644
index 0000000..3261e5c
--- /dev/null
+++ b/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))
\ No newline at end of file
diff --git a/research/evaluator.py b/research/evaluator.py
index 18eb790..b3ea451 100644
--- a/research/evaluator.py
+++ b/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'))
\ No newline at end of file
+                        Path('.', 'research', 'pred_new_' + model_type + '.png'))
\ No newline at end of file