Skip to content

Commit

Permalink
Refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
Nikita Chernyadev committed Apr 17, 2024
1 parent c85ec89 commit 0ea8af6
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 2 deletions.
2 changes: 1 addition & 1 deletion mojo/elements/body.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,5 +99,5 @@ def set_kinematic(self, value: bool):
self.mjcf.add("freejoint")
self._mojo.mark_dirty()
elif not value and self.is_kinematic():
self.mjcf.freejoint.remove()
self.remove_all_joints()
self._mojo.mark_dirty()
11 changes: 11 additions & 0 deletions mojo/elements/element.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,14 @@ def _is_kinematic(elem: mjcf.Element):
return has_freejoint or has_joints or _is_kinematic(elem.parent)


def _remove_all_joints(elem: mjcf.Element):
if hasattr(elem, "freejoint") and elem.freejoint is not None:
elem.freejoint.remove()
if hasattr(elem, "joint") and len(elem.joint) > 0:
for joint in elem.joint:
joint.remove()


def _find_freejoint(elem: mjcf.Element):
if elem.parent is None:
# Root of tree
Expand Down Expand Up @@ -73,6 +81,9 @@ def get_quaternion(self) -> np.ndarray:
def is_kinematic(self) -> bool:
return _is_kinematic(self.mjcf)

def remove_all_joints(self):
_remove_all_joints(self.mjcf)

@property
def id(self):
return self._mojo.physics.bind(self.mjcf).element_id
Expand Down
2 changes: 1 addition & 1 deletion mojo/elements/geom.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,5 +176,5 @@ def set_kinematic(self, value: bool):
self.mjcf.parent.add("freejoint")
self._mojo.mark_dirty()
elif not value and self.is_kinematic():
self.mjcf.parent.freejoint.remove()
self.parent.remove_all_joints()
self._mojo.mark_dirty()

0 comments on commit 0ea8af6

Please sign in to comment.