From 736ca4edd7063a9bfaed5c0be0d44f76cef5fbb4 Mon Sep 17 00:00:00 2001
From: = <stepjamuk@gmail.com>
Date: Sun, 7 Apr 2024 09:18:12 +0100
Subject: [PATCH] Kinematic moved to base

---
 mojo/elements/body.py    |  7 -------
 mojo/elements/element.py | 12 ++++++++++++
 mojo/elements/geom.py    |  7 -------
 3 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/mojo/elements/body.py b/mojo/elements/body.py
index d57cd8b..77d270c 100644
--- a/mojo/elements/body.py
+++ b/mojo/elements/body.py
@@ -59,13 +59,6 @@ def joints(self) -> list[joint.Joint]:
         joints = self.mjcf.find_all("joint") or []
         return [joint.Joint(self._mojo, mjcf) for mjcf in joints]
 
-    def set_kinematic(self, value: bool):
-        if value and not self.is_kinematic():
-            self.mjcf.add("freejoint")
-            self._mojo.mark_dirty()
-        if not value and self.is_kinematic() and self.mjcf.freejoint is not None:
-            self.mjcf.freejoint.remove()
-
     def set_euler(self, euler: np.ndarray):
         self.set_quaternion(
             quaternion.as_float_array(
diff --git a/mojo/elements/element.py b/mojo/elements/element.py
index 6da690b..507a497 100644
--- a/mojo/elements/element.py
+++ b/mojo/elements/element.py
@@ -61,6 +61,18 @@ def get_quaternion(self) -> np.ndarray:
         mujoco.mju_mat2Quat(quat, self._mojo.physics.bind(self.mjcf).xmat)
         return quat
 
+    def set_kinematic(self, value: bool):
+        if value and not self.is_kinematic():
+            self.mjcf.add("freejoint")
+            self._mojo.mark_dirty()
+        if (
+            not value
+            and self.is_kinematic()
+            and hasattr(self.mjcf, "freejoint")
+            and self.mjcf.freejoint is not None
+        ):
+            self.mjcf.freejoint.remove()
+
     def is_kinematic(self) -> bool:
         return _is_kinematic(self.mjcf)
 
diff --git a/mojo/elements/geom.py b/mojo/elements/geom.py
index 46758d3..f16f04c 100644
--- a/mojo/elements/geom.py
+++ b/mojo/elements/geom.py
@@ -153,13 +153,6 @@ def is_collidable(self) -> bool:
             and self._mojo.physics.bind(self.mjcf).conaffinity == 1
         )
 
-    def set_kinematic(self, value: bool):
-        if value and not self.is_kinematic():
-            self.mjcf.parent.add("freejoint")
-            self._mojo.mark_dirty()
-        if not value and self.is_kinematic() and self.mjcf.parent.freejoint is not None:
-            self.mjcf.parent.freejoint.remove()
-
     def has_collided(self, other: Geom = None):
         if other is not None and not other.is_kinematic() and not self.is_kinematic():
             warnings.warn("You are checking collisions of two non-kinematic bodies.")