Get/set matrix for sites

stepjam · Apr 5, 2024 · 69c3098 · 69c3098
1 parent 58e7594
commit 69c3098
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/mojo/elements/site.py b/mojo/elements/site.py
@@ -93,6 +93,16 @@ def get_quaternion(self) -> np.ndarray:
         mujoco.mju_mat2Quat(quat, self._mojo.physics.bind(self.mjcf).xmat)
         return quat
 
+    def set_matrix(self, matrix: np.ndarray):
+        assert matrix.shape == (3, 3)
+        self._mojo.physics.bind(self.mjcf).xmat = np.reshape(matrix, (9,))
+        quat = np.zeros(4)
+        mujoco.mju_mat2Quat(quat, self._mojo.physics.bind(self.mjcf).xmat)
+        self.mjcf.quat = quat
+
+    def get_matrix(self) -> np.ndarray:
+        return np.reshape(self._mojo.physics.bind(self.mjcf).xmat.copy(), (3, 3))
+
     def set_color(self, color: np.ndarray):
         color = np.array(color)
         if len(color) == 3:

diff --git a/tests/elements/test_site.py b/tests/elements/test_site.py
@@ -30,6 +30,12 @@ def test_get_set_quaternion(mojo: Mojo, site: Site):
     assert_array_equal(site.get_quaternion(), expected)
 
 
+def test_get_set_matrix(mojo: Mojo, site: Site):
+    expected = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+    site.set_matrix(expected)
+    assert_array_equal(site.get_matrix(), expected)
+
+
 def test_get_set_color(mojo: Mojo, site: Site):
     expected = np.array([0.8, 0.8, 0.8, 1.0], dtype=np.float32)
     site.set_color(expected)