Add cameras to MJX. Fixes #1422

PiperOrigin-RevId: 610870732 Change-Id: Iaa39bd00f56046c595dc961281ac23446ca5bdf2
2024-02-27 14:16:16 -08:00
parent 478ce0c8d8
commit 2067e2083a
12 changed files with 191 additions and 36 deletions
@@ -19,6 +19,7 @@ from typing import Optional, Tuple, Union
 import jax
 from jax import numpy as jp
 import mujoco
+from mujoco.mjx._src import math
 from mujoco.mjx._src import scan
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import Data
@@ -186,3 +187,15 @@ def xfrc_accumulate(m: Model, d: Data) -> jax.Array:
      jp.arange(m.nbody),
  )
  return jp.sum(qfrc, axis=0)
+
+
+def local_to_global(
+    world_pos: jax.Array,
+    world_quat: jax.Array,
+    local_pos: jax.Array,
+    local_quat: jax.Array,
+) -> Tuple[jax.Array, jax.Array]:
+  """Converts local position/orientation to world frame."""
+  pos = world_pos + math.rotate(local_pos, world_quat)
+  mat = math.quat_to_mat(math.quat_mul(world_quat, local_quat))
+  return pos, mat