diff --git a/figures/sabercat_results.py b/figures/sabercat_results.py
index caed8bd..f80ec25 100644
--- a/figures/sabercat_results.py
+++ b/figures/sabercat_results.py
@@ -114,7 +114,7 @@ def compute_traj_length(values: gtsam.Values):
                     wheel_type += "_planar"
                 has_imu = "imu" in sensors
 
-            error = jrl.computeATEPose3(runs["_gt"], values, align=args.align)
+            error = rose.jrl.computeATEPose3(runs["_gt"], values, align=args.align)
 
             data.append(
                 [
diff --git a/figures/sim.pdf b/figures/sim.pdf
index 1d08d30..c621dd5 100644
Binary files a/figures/sim.pdf and b/figures/sim.pdf differ
diff --git a/figures/surface.pdf b/figures/surface.pdf
index b2d4932..48c36d1 100644
Binary files a/figures/surface.pdf and b/figures/surface.pdf differ
diff --git a/python/rose/jrl.py b/python/rose/jrl.py
index 8ea685c..015a376 100644
--- a/python/rose/jrl.py
+++ b/python/rose/jrl.py
@@ -1,9 +1,12 @@
 import gtsam
 import jrl
 
-
 # ------------------------- Helpers for Wrapping GTSAM containers ------------------------- #
 from rose.rose_python import (
+    IMUBiasTag,
+    StereoPoint2Tag,
+    CombinedIMUTag,
+    StereoFactorPose3Point3Tag,
     makeRoseParser,
     makeRoseWriter,
     computeATEPose3,
@@ -27,7 +30,7 @@ def values2typedvalues(values: gtsam.Values) -> jrl.TypedValues:
         elif char in ["l", "v", "i"]:
             types[k] = jrl.Point3Tag
         elif char == "b":
-            types[k] = jrl.IMUBiasTag
+            types[k] = IMUBiasTag
         elif char == "s":
             types[k] = jrl.Point2Tag
         else:
diff --git a/python/rose/sim.py b/python/rose/sim.py
index 3bc2013..210ae60 100644
--- a/python/rose/sim.py
+++ b/python/rose/sim.py
@@ -130,9 +130,7 @@ def __init__(
         self.gravity = np.array([0, 0, 9.81])
 
         self.vw_last = x0.rotation().matrix() @ np.array([v, 0, 0])
-        self.state = State(
-            x0, self._m(x0), velocity=self.vw_last, w_intr=self.params.w_intr
-        )
+        self.state = State(x0, velocity=self.vw_last, w_intr=self.params.w_intr)
 
         self.measurements = [
             {
@@ -204,7 +202,6 @@ def _f(self, state, w, v):
         return (
             State(
                 pose,
-                self._m(pose),
                 R @ v_body,
                 bias_w,
                 bias_a,
@@ -239,12 +236,6 @@ def _setup_manifold(self, M, x, y, yaw):
         R = gtsam.Rot3(np.column_stack((x_basis, y_basis, z_basis)))
         return gtsam.Pose3(R, p0)
 
-    def _m(self, pose):
-        p = pose.translation()
-        R = pose.rotation().matrix()
-        m_new = R.T @ self.gradM2(p) @ R / np.linalg.norm(self.gradM(p))
-        return Manifold(np.array([m_new[0, 0], m_new[0, 1], m_new[1, 1]]))
-
     # ------------------------- Wheel Utilities ------------------------- #
     def _setup_wheels(self):
         # Setup Slippage
diff --git a/requirements.in b/requirements.in
index 7da2925..a4c9223 100644
--- a/requirements.in
+++ b/requirements.in
@@ -10,4 +10,5 @@ tqdm
 yourdfpy
 pandas
 seaborn
-tabulate
\ No newline at end of file
+tabulate
+sympy
\ No newline at end of file
diff --git a/scripts/optimize_cov.py b/scripts/optimize_cov.py
index d1a248c..7944936 100644
--- a/scripts/optimize_cov.py
+++ b/scripts/optimize_cov.py
@@ -8,8 +8,8 @@
 import jrl
 import numpy as np
 from rose.dataset import CamNoise, Dataset2JRL, Sensor, WheelNoise
-from rose.flat import FlatDataset
-from rose.kaist import KaistDataset
+import rose
+from rose.datasets import FlatDataset, KaistDataset
 from rose.rose_python import (
     PlanarPriorTag,
     WheelRoseIntrSlipTag,
@@ -211,7 +211,7 @@ def objective(x, arg=""):
                     0,
                     True,
                 )
-                ate_none = jrl.computeATEPose3(gt, sol, False, False)[0] / km
+                ate_none = rose.jrl.computeATEPose3(gt, sol, False, False)[0] / km
             except:
                 ate_none = 500
 
diff --git a/scripts/sim.py b/scripts/sim.py
index 3d727b2..ed579a9 100644
--- a/scripts/sim.py
+++ b/scripts/sim.py
@@ -69,7 +69,7 @@ def compute_traj_length(values: gtsam.Values):
     gt = dataset.get_ground_truth()
 
     writer = rose.rose_python.makeRoseWriter()
-    writer.writeResults(values2results(gt), "figures/data/gt.jrr", False)
+    writer.writeResults(values2results(gt), "figures/data/GT.jrr", False)
 
     traj = {"GT": dataset.traj[Sensor.GT], "Wheel": dataset.traj[Sensor.WHEEL]}
     data = [["Kind", "ATEt", "ATEr"]]
@@ -120,7 +120,7 @@ def compute_traj_length(values: gtsam.Values):
 
         # Compute error
         km = compute_traj_length(gt) / 1000
-        et, er = jrl.computeATEPose3(gt, sol, False)
+        et, er = rose.jrl.computeATEPose3(gt, sol, False)
         data.append([names[tag], et, er * 180 / np.pi])
         traj[names[tag]] = sol
         print(f"\tFinished {tag}, e: {graph.error(sol):.2E}...")