diff --git a/doc/changelog.rst b/doc/changelog.rst index 30a419e3..c1eb2f57 100644 --- a/doc/changelog.rst +++ b/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` 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. diff --git a/doc/mjx.rst b/doc/mjx.rst index 8dc4bf66..28fb312d 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -245,6 +245,8 @@ The following features are **in development** and coming soon: - ``TENDON`` * - :ref:`Sensors ` - All except ``PLUGIN``, ``USER`` + * - Lights + - Positions and directions of lights The following features are **unsupported**: diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index d3e048d2..472356b0 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/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 diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py index 79aafa1a..0d094210 100644 --- a/mjx/mujoco/mjx/_src/device.py +++ b/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)) ), } diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index c6bbdbc1..f04eafe3 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index d1fbbe55..4b0c2f4f 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/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. diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py index 89f360cb..168637dc 100644 --- a/mjx/mujoco/mjx/_src/scan.py +++ b/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_j(typ_ids['j']) - if group_by == 'j' - else key_u(typ_ids['u'], typ_ids['j']) - ) + key = key_fn(typ_ids) 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 diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 632a225d..ce099c7e 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/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 - 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 + v_local_to_global = jax.vmap(support.local_to_global) # 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.""" diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index eb631ce8..0b25f71f 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/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(""" diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 8f5d27a2..0bdc0d50 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/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 diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 578f78b6..19085ff5 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/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 diff --git a/mjx/mujoco/mjx/test_data/pendula.xml b/mjx/mujoco/mjx/test_data/pendula.xml index 0dc476ab..f2a3fc8e 100644 --- a/mjx/mujoco/mjx/test_data/pendula.xml +++ b/mjx/mujoco/mjx/test_data/pendula.xml @@ -19,15 +19,22 @@ + + + + + + + - +