Add cameras to MJX. Fixes #1422

PiperOrigin-RevId: 610870732
Change-Id: Iaa39bd00f56046c595dc961281ac23446ca5bdf2

doc/changelog.rst

@@ -31,10 +31,11 @@ MJX
 5. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices.
 6. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`.
 7. Added support for :ref:`mjOption.impratio<mjOption>` to ``mjx.Model``.
+8. Added support for cameras in ``mjx.Model`` and ``mjx.Data``. Fixes :github:issue:`1422`.
 
 Python bindings
 ^^^^^^^^^^^^^^^
-7. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members
+9. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members
    of the ``mjContact`` struct, and all array members of the ``mjrContext`` struct.
 
 

doc/mjx.rst

@@ -245,6 +245,8 @@ The following features are **in development** and coming soon:
      - ``TENDON``
    * - :ref:`Sensors <mjtSensor>`
      - All except ``PLUGIN``, ``USER``
+   * - Lights
+     - Positions and directions of lights
 
 The following features are **unsupported**:
 

mjx/mujoco/mjx/__init__.py

@@ -36,6 +36,7 @@ from mujoco.mjx._src.io import put_data
 from mujoco.mjx._src.io import put_model
 from mujoco.mjx._src.passive import passive
 from mujoco.mjx._src.ray import ray
+from mujoco.mjx._src.smooth import camlight
 from mujoco.mjx._src.smooth import com_pos
 from mujoco.mjx._src.smooth import com_vel
 from mujoco.mjx._src.smooth import crb

mjx/mujoco/mjx/_src/device.py

@@ -69,9 +69,12 @@ _TRANSFORMS = {
     (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
+    (types.Data, 'cam_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
+    (types.Model, 'cam_mat0'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Contact, 'frame'): (
         lambda x: x.reshape(x.shape[:-1] + (3, 3))  # pylint: disable=g-long-lambda
-        if x is not None and x.shape[0] else jp.zeros((0, 3, 3))
+        if x is not None and x.shape[0]
+        else jp.zeros((0, 3, 3))
     ),
 }
 
@@ -80,9 +83,12 @@ _INVERSE_TRANSFORMS = {
     (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
+    (types.Data, 'cam_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
+    (types.Model, 'cam_mat0'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Contact, 'frame'): (
         lambda x: x.reshape(x.shape[:-2] + (9,))  # pylint: disable=g-long-lambda
-        if x is not None and x.shape[0] else jp.zeros((0, 9))
+        if x is not None and x.shape[0]
+        else jp.zeros((0, 9))
     ),
 }
 

mjx/mujoco/mjx/_src/forward.py

@@ -65,6 +65,7 @@ def fwd_position(m: Model, d: Data) -> Data:
   # TODO(robotics-simulation): tendon
   d = smooth.kinematics(m, d)
   d = smooth.com_pos(m, d)
+  d = smooth.camlight(m, d)
   d = smooth.crb(m, d)
   d = smooth.factor_m(m, d)
   d = collision_driver.collision(m, d)

mjx/mujoco/mjx/_src/io.py

@@ -128,6 +128,7 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
       for f in types.Model.fields()
       if f.type is jax.Array
   }
+  device_fields['cam_mat0'] = device_fields['cam_mat0'].reshape((-1, 3, 3))
   device_fields.update(mesh.get(m))
   device_fields = jax.device_put(device_fields, device=device)
 
@@ -185,6 +186,8 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=float),
       site_xpos=jp.zeros((m.nsite, 3), dtype=float),
       site_xmat=jp.zeros((m.nsite, 3, 3), dtype=float),
+      cam_xpos=jp.zeros((m.ncam, 3), dtype=float),
+      cam_xmat=jp.zeros((m.ncam, 3, 3), dtype=float),
       subtree_com=zero_nbody_3,
       cdof=zero_nv_6,
       cinert=zero_nbody_10,
@@ -301,7 +304,7 @@ def get_data_into(
 
       value = getattr(d_i, field.name)
 
-      if field.name in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'):
+      if field.name in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'):
         value = value.reshape((-1, 9))
 
       if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
@@ -368,7 +371,7 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
       if f.type is jax.Array
   }
 
-  for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'):
+  for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'):
     fields[fname] = fields[fname].reshape((-1, 3, 3))
 
   # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.

mjx/mujoco/mjx/_src/scan.py

@@ -140,6 +140,7 @@ def _check_input(m: Model, args: Any, in_types: str) -> None:
       'u': m.nu,
       'a': m.na,
       's': m.nsite,
+      'c': m.ncam,
   }
   for idx, (arg, typ) in enumerate(zip(args, in_types)):
     if len(arg) != size[typ]:
@@ -193,6 +194,7 @@ def flat(
       'v': split according to degrees of freedom (len(qvel))
       'u': split according to actuators
       'a': split according to actuator activations
+      'c': split according to camera
     out_types: string specifying the types the output dimension matches
     *args: the input arguments corresponding to ``in_types``
     group_by: the type to group by, either joints or actuators
@@ -205,24 +207,14 @@ def flat(
   """
   _check_input(m, args, in_types)
 
-  if group_by not in {'j', 'u'}:
+  if group_by not in {'j', 'u', 'c'}:
     raise NotImplementedError(f'group by type "{group_by}" not implemented.')
 
-  def key_j(ids):
+  def key_j(type_ids):
     if any(t in 'jqv' for t in in_types + out_types):
-      return tuple(m.jnt_type[ids])
+      return tuple(m.jnt_type[type_ids['j']])
     return ()
 
-  def key_u(ids_u, ids_j):
-    return (
-        m.actuator_biastype[ids_u],
-        m.actuator_gaintype[ids_u],
-        m.actuator_dyntype[ids_u],
-        m.actuator_trntype[ids_u],
-        m.jnt_type[ids_j],
-        m.actuator_trnid[ids_u, 1] == -1,  # key by refsite being present
-    )
-
   def type_ids_j(m, i):
     return {
         'b': i,
@@ -231,6 +223,17 @@ def flat(
         'q': np.nonzero(_q_bodyid(m) == i)[0],
     }
 
+  def key_u(type_ids):
+    ids_u, ids_j = type_ids['u'], type_ids['j']
+    return (
+        m.actuator_biastype[ids_u],
+        m.actuator_gaintype[ids_u],
+        m.actuator_dyntype[ids_u],
+        m.actuator_trntype[ids_u],
+        m.jnt_type[ids_j],
+        m.actuator_trnid[ids_u, 1] == -1,  # key by refsite being present
+    )
+
   def type_ids_u(m, i):
     typ_ids = {
         'u': i,
@@ -256,19 +259,26 @@ def flat(
 
     return typ_ids
 
+  def key_c(type_ids):
+    return m.cam_mode[type_ids['c']], m.cam_targetbodyid[type_ids['c']] >= 0
+
+  def type_ids_c(unused_m, i):
+    return {
+        'c': i,
+    }
+
+  type_ids_fn = {'j': type_ids_j, 'u': type_ids_u, 'c': type_ids_c}[group_by]
+  key_fn = {'j': key_j, 'u': key_u, 'c': key_c}[group_by]
+
   # build up a grouping of type take-ids in body/actuator order
   key_typ_ids, order = {}, []
   all_types = set(in_types + out_types)
-  n_items = {'j': m.nbody, 'u': m.nu}[group_by]
+  n_items = {'j': m.nbody, 'u': m.nu, 'c': m.ncam}[group_by]
   for i in np.arange(n_items, dtype=np.int32):
-    typ_ids = type_ids_j(m, i) if group_by == 'j' else type_ids_u(m, i)
+    typ_ids = type_ids_fn(m, i)
 
     # create grouping key
-    key = (
+    key = key_fn(typ_ids)
-        key_j(typ_ids['j'])
-        if group_by == 'j'
-        else key_u(typ_ids['u'], typ_ids['j'])
-    )
     order.append((key, typ_ids))
 
     # add ids per type to the corresponding group
@@ -307,7 +317,7 @@ def flat(
   # concatenate back to a single tree and drop the grouping dimension
   f_ret_is_seq = isinstance(ys[0], (list, tuple))
   ys = ys if f_ret_is_seq else [[y] for y in ys]
-  flat_ = {'j': 'b', 'u': 'uaj'}[group_by]
+  flat_ = {'j': 'b', 'u': 'uaj', 'c': 'c'}[group_by]
   ys = [
       [v if typ in flat_ else jp.concatenate(v) for v, typ in zip(y, out_types)]
       for y in ys

mjx/mujoco/mjx/_src/smooth.py

@@ -21,6 +21,7 @@ from mujoco.mjx._src import math
 from mujoco.mjx._src import scan
 from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
+from mujoco.mjx._src.types import CamLightType
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import JointType
@@ -96,25 +97,21 @@ def kinematics(m: Model, d: Data) -> Data:
       m.body_quat,
   )
 
-  @jax.vmap
+  v_local_to_global = jax.vmap(support.local_to_global)
-  def local_to_global(pos1, quat1, pos2, quat2):
-    pos = pos1 + math.rotate(pos2, quat1)
-    mat = math.quat_to_mat(math.quat_mul(quat1, quat2))
-    return pos, mat
 
   # TODO(erikfrey): confirm that quats are more performant for mjx than mats
-  xipos, ximat = local_to_global(xpos, xquat, m.body_ipos, m.body_iquat)
+  xipos, ximat = v_local_to_global(xpos, xquat, m.body_ipos, m.body_iquat)
   d = d.replace(qpos=qpos, xanchor=xanchor, xaxis=xaxis, xpos=xpos)
   d = d.replace(xquat=xquat, xmat=xmat, xipos=xipos, ximat=ximat)
 
   if m.ngeom:
-    geom_xpos, geom_xmat = local_to_global(
+    geom_xpos, geom_xmat = v_local_to_global(
         xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat
     )
     d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat)
 
   if m.nsite:
-    site_xpos, site_xmat = local_to_global(
+    site_xpos, site_xmat = v_local_to_global(
         xpos[m.site_bodyid], xquat[m.site_bodyid], m.site_pos, m.site_quat
     )
     d = d.replace(site_xpos=site_xpos, site_xmat=site_xmat)
@@ -199,6 +196,76 @@ def com_pos(m: Model, d: Data) -> Data:
   return d
 
 
+def camlight(m: Model, d: Data) -> Data:
+  """Computes camera and light positions and orientations."""
+  if m.ncam == 0:
+    return d.replace(cam_xpos=jp.zeros((0, 3)), cam_xmat=jp.zeros((0, 3, 3)))
+
+  # use target body only if target body is specified
+  is_target_cam = (m.cam_mode == CamLightType.TARGETBODY) | (
+      m.cam_mode == CamLightType.TARGETBODYCOM
+  )
+  cam_mode = np.where(
+      is_target_cam & (m.cam_targetbodyid < 0), CamLightType.FIXED, m.cam_mode
+  )
+
+  cam_xpos, cam_xmat = jax.vmap(support.local_to_global)(
+      d.xpos[m.cam_bodyid], d.xquat[m.cam_bodyid], m.cam_pos, m.cam_quat
+  )
+
+  def fn(
+      camid,
+      cam_mode,
+      cam_xpos,
+      cam_xmat,
+      body_xpos,
+      subtree_com,
+      target_body_xpos,
+      target_subtree_com,
+  ):
+    if cam_mode == CamLightType.TRACK:
+      cam_xmat = m.cam_mat0[camid]
+      cam_xpos = body_xpos + m.cam_pos0[camid]
+    elif cam_mode == CamLightType.TRACKCOM:
+      cam_xmat = m.cam_mat0[camid]
+      cam_xpos = subtree_com + m.cam_poscom0[camid]
+    elif cam_mode in (CamLightType.TARGETBODY, CamLightType.TARGETBODYCOM):
+      # get position to look at
+      pos = target_body_xpos
+      if cam_mode == CamLightType.TARGETBODYCOM:
+        pos = target_subtree_com
+      # zaxis = -desired camera direction, in global frame
+      mat_3 = math.normalize(cam_xpos - pos)
+      # xaxis: orthogonal to zaxis and to (0,0,1)
+      mat_1 = math.normalize(jp.cross(jp.array([0.0, 0.0, 1.0]), mat_3))
+      mat_2 = math.normalize(jp.cross(mat_3, mat_1))
+      cam_xmat = jp.array([mat_1, mat_2, mat_3]).T
+    return cam_xpos, cam_xmat
+
+  cam_xpos, cam_xmat = scan.flat(
+      m,
+      fn,
+      'c' * 8,
+      'cc',
+      jp.arange(m.ncam),
+      cam_mode,
+      cam_xpos,
+      cam_xmat,
+      d.xpos[m.cam_bodyid],
+      d.subtree_com[m.cam_bodyid],
+      d.xpos[m.cam_targetbodyid],
+      d.subtree_com[m.cam_targetbodyid],
+      group_by='c',
+  )
+
+  d = d.replace(
+      cam_xpos=cam_xpos,
+      cam_xmat=cam_xmat,
+  )
+
+  return d
+
+
 def crb(m: Model, d: Data) -> Data:
   """Runs composite rigid body inertia algorithm."""
 

mjx/mujoco/mjx/_src/smooth_test.py

@@ -94,6 +94,10 @@ class SmoothTest(absltest.TestCase):
     dx = jax.jit(mjx.transmission)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'actuator_length')
     _assert_attr_eq(d, dx, 'actuator_moment')
+    # camlight
+    dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'cam_xpos')
+    _assert_eq(d.cam_xmat.reshape((-1, 3, 3)), dx.cam_xmat, 'cam_xmat')
 
   def test_disable_gravity(self):
     m = mujoco.MjModel.from_xml_string("""

mjx/mujoco/mjx/_src/support.py

@@ -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

mjx/mujoco/mjx/_src/types.py

@@ -213,6 +213,24 @@ class BiasType(enum.IntEnum):
   # unsupported: MUSCLE, USER
 
 
+class CamLightType(enum.IntEnum):
+  """Type of camera light.
+
+  Attributes:
+    FIXED: pos and rot fixed in body
+    TRACK: pos tracks body, rot fixed in global
+    TRACKCOM: pos tracks subtree com, rot fixed in body
+    TARGETBODY: pos fixed in body, rot tracks target body
+    TARGETBODYCOM: pos fixed in body, rot tracks target subtree com
+  """
+
+  FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED
+  TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK
+  TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM
+  TARGETBODY = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY
+  TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM
+
+
 class Option(PyTreeNode):
   """Physics options.
 
@@ -283,6 +301,7 @@ class Model(PyTreeNode):
     njnt: number of joints
     ngeom: number of geoms
     nsite: number of sites
+    ncam: number of cameras
     nmesh: number of meshes
     nmeshvert: number of vertices in all meshes
     nmeshface: number of triangular faces in all meshes
@@ -360,6 +379,14 @@ class Model(PyTreeNode):
     site_bodyid: id of site's body                            (nsite,)
     site_pos: local position offset rel. to body              (nsite, 3)
     site_quat: local orientation offset rel. to body          (nsite, 4)
+    cam_mode:  camera tracking mode (mjtCamLight)             (ncam,)
+    cam_bodyid:  id of camera's body                          (ncam,)
+    cam_targetbodyid:  id of targeted body; -1: none          (ncam,)
+    cam_pos:  position rel. to body frame                     (ncam, 3)
+    cam_quat:  orientation rel. to body frame                 (ncam, 4)
+    cam_poscom0:  global position rel. to sub-com in qpos0    (ncam, 3)
+    cam_pos0:  global position rel. to body in qpos0          (ncam, 3)
+    cam_mat0:  global orientation in qpos0                    (ncam, 9)
     mat_rgba: rgba                                            (nmat, 4)
     mesh_vertadr: first vertex address                        (nmesh x 1)
     mesh_faceadr: first face address                          (nmesh x 1)
@@ -416,6 +443,7 @@ class Model(PyTreeNode):
   njnt: int
   ngeom: int
   nsite: int
+  ncam: int
   nmesh: int
   nmeshvert: int
   nmeshface: int
@@ -493,6 +521,14 @@ class Model(PyTreeNode):
   site_bodyid: np.ndarray
   site_pos: jax.Array
   site_quat: jax.Array
+  cam_mode: np.ndarray
+  cam_bodyid: np.ndarray
+  cam_targetbodyid: np.ndarray
+  cam_pos: jax.Array
+  cam_quat: jax.Array
+  cam_poscom0: jax.Array
+  cam_pos0: jax.Array
+  cam_mat0: jax.Array
   mesh_vertadr: np.ndarray
   mesh_faceadr: np.ndarray
   mesh_vert: np.ndarray
@@ -588,7 +624,7 @@ class Contact(PyTreeNode):
 
 
 class Data(PyTreeNode):
-  """Dynamic state that updates each step.
+  r"""Dynamic state that updates each step.\
 
   Attributes:
     solver_niter: number of solver iterations, per island         (mjNISLAND,)
@@ -614,6 +650,8 @@ class Data(PyTreeNode):
     geom_xmat: Cartesian geom orientation                         (ngeom, 3, 3)
     site_xpos: Cartesian site position                            (nsite, 3)
     site_xmat: Cartesian site orientation                         (nsite, 9)
+    cam_xpos: Cartesian camera position                           (ncam, 3)
+    cam_xmat: Cartesian camera orientation                        (ncam, 9)
     subtree_com: center of mass of each subtree                   (nbody, 3)
     cdof: com-based motion axis of each dof                       (nv, 6)
     cinert: com-based body inertia and mass                       (nbody, 10)
@@ -673,6 +711,8 @@ class Data(PyTreeNode):
   geom_xmat: jax.Array
   site_xpos: jax.Array
   site_xmat: jax.Array
+  cam_xpos: jax.Array
+  cam_xmat: jax.Array
   subtree_com: jax.Array
   cdof: jax.Array
   cinert: jax.Array

mjx/mujoco/mjx/test_data/pendula.xml

@@ -19,15 +19,22 @@
 
   <worldbody>
     <site name="origin"/>
+    <camera pos="0 1.5 0.8" xyaxes="-1 0 0 0 -0.2 .8"/>
+    <camera mode="targetbodycom" pos="1 -1 2"/>
+    <camera mode="targetbodycom" target="body" pos="0 -1 2"/>
+    <camera mode="targetbody" target="body" pos="0 -1 -2"/>
 
     <!-- a single free body -->
     <body pos="0 0 0">
+      <camera mode="fixed" pos="0 0.1 0.1" euler="2.7 0 0"/>
+      <camera pos="-3 0 1" xyaxes="0 -1 0 1 0 2" mode="trackcom"/>
+      <camera pos="-3 0 1" xyaxes="0 -1 0 1 0 2" mode="track"/>
       <freejoint/>
       <geom/>
     </body>
 
     <!-- a single ball joint with a limit -->
-    <body pos="0.5 0 0">
+    <body name="body" pos="0.5 0 0">
       <joint name="joint1" type="ball" range="0 35"/>
       <site name="s1" pos="1.3 0 0"/>
       <geom/>