Added Inertias (I) and Center Of Mass (M)

Toggle functionality with I and M

Author: rohit-kumar-j

mujoco_viewer/mujoco_viewer.py

@@ -24,6 +24,8 @@ def __init__(self, model, data):
         self._contacts = False
         self._joints = False
         self._wire_frame = False
+        self._inertias = False
+        self._com = False
         self._render_every_frame = True
         self._image_idx = 0
         self._image_path = "/tmp/frame_%07d.png"
@@ -126,6 +128,14 @@ def _key_callback(self, window, key, scancode, action, mods):
                 self.model.geom_rgba[:, 3] /= 5.0
             else:
                 self.model.geom_rgba[:, 3] *= 5.0
+        # Display inertia
+        elif key == glfw.KEY_I:
+            self._inertias = not self._inertias
+            self.vopt.flags[mujoco.mjtVisFlag.mjVIS_INERTIA] = self._inertias
+        # Display center of mass
+        elif key == glfw.KEY_M:
+            self._com = not self._com
+            self.vopt.flags[mujoco.mjtVisFlag.mjVIS_COM] = self._com
         # Wireframe Rendering
         elif key == glfw.KEY_W:
             self._wire_frame = not self._wire_frame
@@ -381,6 +391,14 @@ def add_overlay(gridpos, text1, text2):
             topleft,
             "[J]oints",
             "On" if self._joints else "Off")
+        add_overlay(
+            topleft,
+            "[I]nertia",
+            "On" if self._inertias else "Off")
+        add_overlay(
+            topleft,
+            "Center of [M]ass",
+            "On" if self._com else "Off")
         add_overlay(
             topleft,
             "T[r]ansparent",