chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add cameras to MJX. Fixes #1422

Browse Source 
PiperOrigin-RevId: 610870732
Change-Id: Iaa39bd00f56046c595dc961281ac23446ca5bdf2
This commit is contained in:
Baruch Tabanpour
2024-02-27 14:16:16 -08:00
committed by Copybara-Service
parent 478ce0c8d8
commit 2067e2083a
12 changed files with 191 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+2 -1
View File
@@ -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
+2
View File
@@ -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
+1
View File
@@ -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
+8 -2
View File
@@ -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
+1
View File
@@ -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
+5 -2
View File
@@ -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
+31 -21
View File
@@ -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
mjx/mujoco/mjx/_src/smooth.py
+75 -8
View File
@@ -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."""
mjx/mujoco/mjx/_src/smooth_test.py
+4
View File
@@ -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
+13
View File
@@ -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
+41 -1
View File
@@ -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
+8 -1
View File
@@ -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/>