diff --git a/unity/Editor/Components/MjMouseSpring.cs b/unity/Editor/Components/MjMouseSpring.cs
index b2264727..7e1c2462 100644
--- a/unity/Editor/Components/MjMouseSpring.cs
+++ b/unity/Editor/Components/MjMouseSpring.cs
@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#if UNITY_EDITOR
 using System;
 using UnityEditor;
 using UnityEngine;
@@ -23,7 +22,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 +101,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;
@@ -188,4 +191,3 @@ namespace Mujoco {
     }
   }
 }
-#endif
diff --git a/unity/Editor/Components/MjShapeComponentEditor.cs b/unity/Editor/Components/MjShapeComponentEditor.cs
index 99e24740..20281e61 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() {