diff --git a/unity/Editor/Components/MjMouseSpring.cs b/unity/Editor/Components/MjMouseSpring.cs
index b2264727..c502e551 100644
--- a/unity/Editor/Components/MjMouseSpring.cs
+++ b/unity/Editor/Components/MjMouseSpring.cs
@@ -23,7 +23,7 @@ namespace Mujoco {
   // and left mouse drag will apply a force on the body.  Holding shift down will change the applied
   // force direction between World XZ plane and Y[camera-up].
 
-  [CustomEditor(typeof(MjBody))]
+  [CustomEditor(typeof(MjComponent), true)]
   public class MjMouseSpring : Editor {
     private bool _lastShiftKeyState = false;
 
@@ -102,7 +102,11 @@ namespace Mujoco {
         return;
       }
 
-      MjBody body = target as MjBody;
+      var targetObject = target as MjComponent;
+      MjBody body = targetObject.GetComponentInParent<MjBody>();
+      if(!body) 
+        return;
+
       Vector3 bodyPosition =
           body != null ? body.transform.position : Vector3.zero;
       var scene = MjScene.Instance;
diff --git a/unity/Editor/Components/MjShapeComponentEditor.cs b/unity/Editor/Components/MjShapeComponentEditor.cs
index 41b8ba22..3d0d687c 100644
--- a/unity/Editor/Components/MjShapeComponentEditor.cs
+++ b/unity/Editor/Components/MjShapeComponentEditor.cs
@@ -22,7 +22,7 @@ namespace Mujoco {
 
 [CustomEditor(typeof(MjShapeComponent), true)]
 [CanEditMultipleObjects]
-public class MjShapeComponentEditor : Editor {
+public class MjShapeComponentEditor : MjMouseSpring {
 
   public override void OnInspectorGUI() {