Add put_model, put_data, and get_data to MJX.
These io functions replace device_put, device_get_into. These new functions correctly copy intermediate fields such as efc_J, which the old functions could not. PiperOrigin-RevId: 586712273 Change-Id: Ia3a4104af17db2ebd41b41a2ca501c5de9ac5d05
This commit is contained in:
Erik Frey
committed by
Copybara-Service
Copybara-Service
parent
4c8ffb44bf
commit
67fa7c1dc2
@@ -27,6 +27,9 @@ MJX
|
||||
^^^
|
||||
|
||||
- Added ``site_xpos`` and ``site_xmat`` to MJX.
|
||||
- Added ``put_data``, ``put_model``, ``get_data`` to replace ``device_put`` and ``device_get_into``, which will be
|
||||
deprecated. These new functions correctly translate fields that are the result of intermediate calculations such as
|
||||
``efc_J``.
|
||||
|
||||
Bug fixes
|
||||
^^^^^^^^^
|
||||
|
||||
@@ -21,7 +21,10 @@ from mujoco.mjx._src.device import device_get_into
|
||||
from mujoco.mjx._src.device import device_put
|
||||
from mujoco.mjx._src.forward import forward
|
||||
from mujoco.mjx._src.forward import step
|
||||
from mujoco.mjx._src.io import get_data
|
||||
from mujoco.mjx._src.io import make_data
|
||||
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.smooth import com_pos
|
||||
from mujoco.mjx._src.smooth import com_vel
|
||||
|
||||
@@ -40,7 +40,6 @@ from mujoco.mjx._src.types import DisableBit
|
||||
from mujoco.mjx._src.types import GeomType
|
||||
from mujoco.mjx._src.types import Model
|
||||
# pylint: enable=g-importing-member
|
||||
import numpy as np
|
||||
|
||||
|
||||
# pair-wise collision functions
|
||||
@@ -81,6 +80,12 @@ def _add_candidate(
|
||||
if t1 > t2:
|
||||
t1, t2, g1, g2 = t2, t1, g2, g1
|
||||
|
||||
# MuJoCo does not collide planes with other planes or hfields
|
||||
if t1 == GeomType.PLANE and t2 == GeomType.PLANE:
|
||||
return
|
||||
if t1 == GeomType.PLANE and t2 == GeomType.HFIELD:
|
||||
return
|
||||
|
||||
def mesh_key(i):
|
||||
convex_data = [[None] * m.ngeom] * 3
|
||||
if isinstance(m, Model):
|
||||
@@ -284,13 +289,11 @@ def _collide_geoms(
|
||||
solimp=params.solimp,
|
||||
geom1=geom1,
|
||||
geom2=geom2,
|
||||
dim=np.array([]),
|
||||
efc_address=np.array([]),
|
||||
)
|
||||
return con
|
||||
|
||||
|
||||
def _max_contact_points(m: Model) -> int:
|
||||
def _max_contact_points(m: Union[Model, mujoco.MjModel]) -> int:
|
||||
"""Returns the maximum number of contact points when set as a numeric."""
|
||||
for i in range(m.nnumeric):
|
||||
name = m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0]
|
||||
@@ -334,7 +337,7 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet:
|
||||
return candidate_set
|
||||
|
||||
|
||||
def ncon(m: Model) -> int:
|
||||
def ncon(m: Union[Model, mujoco.MjModel]) -> int:
|
||||
"""Returns the number of contacts computed in MJX given a model."""
|
||||
if m.opt.disableflags & DisableBit.CONTACT:
|
||||
return 0
|
||||
@@ -354,9 +357,8 @@ def ncon(m: Model) -> int:
|
||||
|
||||
def collision(m: Model, d: Data) -> Data:
|
||||
"""Collides geometries."""
|
||||
ncon_ = ncon(m)
|
||||
if ncon_ == 0:
|
||||
return d.replace(contact=Contact.zero(), ncon=0)
|
||||
if ncon(m) == 0:
|
||||
return d.replace(contact=Contact.zero())
|
||||
|
||||
candidate_set = collision_candidates(m)
|
||||
|
||||
@@ -376,13 +378,4 @@ def collision(m: Model, d: Data) -> Data:
|
||||
_, idx = jax.lax.top_k(-contact.dist, k=max_contact_points)
|
||||
contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact)
|
||||
|
||||
if ncon_ != contact.dist.shape[0]:
|
||||
raise RuntimeError('Number of contacts does not match ncon.')
|
||||
|
||||
# TODO(robotics-simulation): move this logic to device_put
|
||||
ns = d.ne + d.nf + d.nl
|
||||
contact = contact.replace(efc_address=np.arange(ns, ns + ncon_ * 4, 4))
|
||||
# TODO(robotics-simulation): add support for other friction dimensions
|
||||
contact = contact.replace(dim=3 * np.ones(ncon_, dtype=np.int32))
|
||||
|
||||
return d.replace(contact=contact, ncon=ncon_)
|
||||
return d.replace(contact=contact)
|
||||
|
||||
@@ -253,7 +253,6 @@ class CapsuleCollisionTest(parameterized.TestCase):
|
||||
self.assertGreater(c.dist[1], 0)
|
||||
# extract the contact point with penetration
|
||||
c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
|
||||
c = c.replace(dim=c.dim[np.array([0])])
|
||||
for field in dataclasses.fields(Contact):
|
||||
_assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4)
|
||||
|
||||
@@ -281,7 +280,6 @@ class ConvexTest(absltest.TestCase):
|
||||
np.testing.assert_array_less(-dx.contact.dist[2:], 0)
|
||||
# extract the contact points with penetration
|
||||
c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
|
||||
c = c.replace(dim=c.dim[np.array([0, 1])])
|
||||
for field in dataclasses.fields(Contact):
|
||||
_assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2)
|
||||
|
||||
@@ -339,7 +337,6 @@ class ConvexTest(absltest.TestCase):
|
||||
np.testing.assert_array_less(-dx.contact.dist[1:], 0)
|
||||
# extract the contact point with penetration
|
||||
c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
|
||||
c = c.replace(dim=c.dim[np.array([0])])
|
||||
for field in dataclasses.fields(Contact):
|
||||
_assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2)
|
||||
|
||||
@@ -517,8 +514,8 @@ class TopKContactTest(absltest.TestCase):
|
||||
dx_all = collision_jit_fn(mx_all, dx)
|
||||
dx_top_k = collision_jit_fn(mx_top_k, dx)
|
||||
|
||||
self.assertEqual(dx_all.ncon, 3)
|
||||
self.assertEqual(dx_top_k.ncon, 2)
|
||||
self.assertEqual(dx_all.contact.dist.shape, (3,))
|
||||
self.assertEqual(dx_top_k.contact.dist.shape, (2,))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
@@ -14,11 +14,12 @@
|
||||
# ==============================================================================
|
||||
"""Core non-smooth constraint functions."""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import jax
|
||||
from jax import numpy as jp
|
||||
import mujoco
|
||||
from mujoco.mjx._src import collision_driver
|
||||
from mujoco.mjx._src import math
|
||||
from mujoco.mjx._src import support
|
||||
# pylint: disable=g-importing-member
|
||||
@@ -276,7 +277,7 @@ def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
|
||||
def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
|
||||
"""Calculates constraint rows for contacts."""
|
||||
|
||||
if (m.opt.disableflags & DisableBit.CONTACT) or d.ncon == 0:
|
||||
if collision_driver.ncon(m) == 0:
|
||||
return None
|
||||
|
||||
@jax.vmap
|
||||
@@ -313,7 +314,9 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
|
||||
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
|
||||
|
||||
|
||||
def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
|
||||
def count_constraints(
|
||||
m: Union[Model, mujoco.MjModel]
|
||||
) -> Tuple[int, int, int, int]:
|
||||
"""Returns equality, friction, limit, and contact constraint counts."""
|
||||
if m.opt.disableflags & DisableBit.CONSTRAINT:
|
||||
return 0, 0, 0, 0
|
||||
@@ -333,10 +336,7 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
|
||||
else:
|
||||
nl = int(m.jnt_limited.sum())
|
||||
|
||||
if m.opt.disableflags & DisableBit.CONTACT:
|
||||
nc = 0
|
||||
else:
|
||||
nc = d.ncon * 4
|
||||
nc = collision_driver.ncon(m) * 4
|
||||
|
||||
return ne, nf, nl, nc
|
||||
|
||||
@@ -344,10 +344,6 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
|
||||
def make_constraint(m: Model, d: Data) -> Data:
|
||||
"""Creates constraint jacobians and other supporting data."""
|
||||
|
||||
ns = sum(count_constraints(m, d)[:-1])
|
||||
# TODO(robotics-simulation): make device_put set nefc/efc_address instead
|
||||
d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + d.ncon * 4, 4)})
|
||||
|
||||
if m.opt.disableflags & DisableBit.CONSTRAINT:
|
||||
efcs = ()
|
||||
else:
|
||||
@@ -364,7 +360,7 @@ def make_constraint(m: Model, d: Data) -> Data:
|
||||
if not efcs:
|
||||
z = jp.empty(0)
|
||||
d = d.replace(efc_J=jp.empty((0, m.nv)))
|
||||
d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z, nefc=0)
|
||||
d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z)
|
||||
return d
|
||||
|
||||
efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs)
|
||||
@@ -378,6 +374,6 @@ def make_constraint(m: Model, d: Data) -> Data:
|
||||
|
||||
aref, r = fn(efc)
|
||||
d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref)
|
||||
d = d.replace(efc_frictionloss=efc.frictionloss, nefc=r.shape[0])
|
||||
d = d.replace(efc_frictionloss=efc.frictionloss)
|
||||
|
||||
return d
|
||||
|
||||
@@ -136,13 +136,8 @@ class ConstraintTest(parameterized.TestCase):
|
||||
m = test_util.load_test_file('ant.xml')
|
||||
d = mujoco.MjData(m)
|
||||
|
||||
m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONSTRAINT
|
||||
mx, dx = mjx.device_put(m), mjx.device_put(d)
|
||||
dx = constraint.make_constraint(mx, dx)
|
||||
self.assertGreater(dx.efc_J.shape[0], 1)
|
||||
|
||||
m.opt.disableflags = m.opt.disableflags | DisableBit.CONSTRAINT
|
||||
mx = mjx.device_put(m)
|
||||
mx, dx = mjx.device_put(m), mjx.device_put(d)
|
||||
dx = constraint.make_constraint(mx, dx)
|
||||
self.assertEqual(dx.efc_J.shape[0], 0)
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ import copy
|
||||
import dataclasses
|
||||
|
||||
import typing
|
||||
from typing import Dict, Optional, Sequence, TypeVar
|
||||
from typing import Any, Dict, Optional, Sequence, TypeVar
|
||||
import jax
|
||||
import numpy as np
|
||||
|
||||
@@ -113,6 +113,10 @@ class PyTreeNode:
|
||||
# stub for pytype
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def fields(cls) -> tuple[dataclasses.Field[Any], ...]:
|
||||
return dataclasses.fields(cls)
|
||||
|
||||
def tree_replace(
|
||||
self, params: Dict[str, Optional[jax.typing.ArrayLike]]
|
||||
) -> 'PyTreeNode':
|
||||
|
||||
@@ -25,6 +25,7 @@ import mujoco
|
||||
from mujoco.mjx._src import collision_driver
|
||||
from mujoco.mjx._src import mesh
|
||||
from mujoco.mjx._src import types
|
||||
import numpy as np
|
||||
|
||||
_MJ_TYPE_ATTR = {
|
||||
mujoco.mjtBias: (mujoco.MjModel.actuator_biastype,),
|
||||
@@ -68,7 +69,6 @@ _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.Model, 'actuator_trnid'): lambda x: x[:, 0],
|
||||
(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))
|
||||
@@ -265,9 +265,16 @@ def device_get_into(result, value):
|
||||
|
||||
else:
|
||||
if isinstance(result, mujoco.MjData):
|
||||
ncon = value.contact.dist.shape[0]
|
||||
nefc = value.efc_J.shape[0]
|
||||
mujoco._functions._realloc_con_efc( # pylint: disable=protected-access
|
||||
result, ncon=value.ncon, nefc=value.nefc
|
||||
result, ncon=ncon, nefc=nefc
|
||||
)
|
||||
result.ncon = ncon
|
||||
result.nefc = nefc
|
||||
efc_start = nefc - ncon * 4
|
||||
result.contact.efc_address[:] = np.arange(efc_start, nefc, 4)
|
||||
result.contact.dim[:] = 3
|
||||
|
||||
for f in dataclasses.fields(value): # type: ignore
|
||||
if (type(value), f.name) in _DERIVED:
|
||||
|
||||
+313
-70
@@ -14,94 +14,337 @@
|
||||
# ==============================================================================
|
||||
"""Functions to initialize, load, or save data."""
|
||||
|
||||
import copy
|
||||
from typing import List, Union
|
||||
|
||||
import jax
|
||||
from jax import numpy as jp
|
||||
import mujoco
|
||||
from mujoco.mjx._src import collision_driver
|
||||
from mujoco.mjx._src import constraint
|
||||
# pylint: disable=g-importing-member
|
||||
from mujoco.mjx._src.types import Contact
|
||||
from mujoco.mjx._src.types import Data
|
||||
from mujoco.mjx._src.types import Model
|
||||
# pylint: enable=g-importing-member
|
||||
from mujoco.mjx._src import mesh
|
||||
from mujoco.mjx._src import types
|
||||
import numpy as np
|
||||
|
||||
|
||||
def make_data(m: Model) -> Data:
|
||||
def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
|
||||
"""Puts mujoco.MjOption onto a device, resulting in mjx.Option."""
|
||||
if o.integrator not in set(types.IntegratorType):
|
||||
raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}')
|
||||
|
||||
if o.cone not in set(types.ConeType):
|
||||
raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}')
|
||||
|
||||
if o.solver not in set(types.SolverType):
|
||||
raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}')
|
||||
|
||||
for i in range(mujoco.mjtEnableBit.mjNENABLE):
|
||||
if o.enableflags & 2**i:
|
||||
raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}')
|
||||
|
||||
static_fields = {
|
||||
f.name: copy.copy(getattr(o, f.name))
|
||||
for f in types.Option.fields()
|
||||
if f.type in (int, bytes, np.ndarray)
|
||||
}
|
||||
static_fields['integrator'] = types.IntegratorType(o.integrator)
|
||||
static_fields['cone'] = types.ConeType(o.cone)
|
||||
static_fields['solver'] = types.SolverType(o.solver)
|
||||
static_fields['disableflags'] = types.DisableBit(o.disableflags)
|
||||
|
||||
device_fields = {
|
||||
f.name: copy.copy(getattr(o, f.name))
|
||||
for f in types.Option.fields()
|
||||
if f.type is jax.Array
|
||||
}
|
||||
device_fields = jax.device_put(device_fields, device=device)
|
||||
|
||||
has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
|
||||
|
||||
return types.Option(
|
||||
has_fluid_params=has_fluid_params,
|
||||
**static_fields,
|
||||
**device_fields,
|
||||
)
|
||||
|
||||
|
||||
def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic:
|
||||
"""Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic."""
|
||||
return types.Statistic(
|
||||
meaninertia=jax.device_put(s.meaninertia, device=device)
|
||||
)
|
||||
|
||||
|
||||
def put_model(m: mujoco.MjModel, device=None) -> types.Model:
|
||||
"""Puts mujoco.MjModel onto a device, resulting in mjx.Model."""
|
||||
if m.ntendon:
|
||||
raise NotImplementedError('tendons are not supported')
|
||||
|
||||
if (m.geom_condim != 3).any() or (m.pair_dim != 3).any():
|
||||
raise NotImplementedError('only condim=3 is supported')
|
||||
|
||||
# check collision geom types
|
||||
for g1, g2, *_ in collision_driver.collision_candidates(m):
|
||||
if collision_driver.get_collision_fn((g1, g2)) is None:
|
||||
g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
|
||||
raise NotImplementedError(f'({g1}, {g2}) has no collision function')
|
||||
|
||||
for enum_field, enum_type, mj_type in (
|
||||
(m.actuator_biastype, types.BiasType, mujoco.mjtBias),
|
||||
(m.actuator_dyntype, types.DynType, mujoco.mjtDyn),
|
||||
(m.actuator_gaintype, types.GainType, mujoco.mjtGain),
|
||||
(m.actuator_trntype, types.TrnType, mujoco.mjtTrn),
|
||||
(m.eq_type, types.EqType, mujoco.mjtEq),
|
||||
):
|
||||
missing = set(enum_field) - set(enum_type)
|
||||
if missing:
|
||||
raise NotImplementedError(
|
||||
f'{[mj_type(m) for m in missing]} not supported'
|
||||
)
|
||||
|
||||
opt = _put_option(m.opt, device=device)
|
||||
stat = _put_statistic(m.stat, device=device)
|
||||
|
||||
static_fields = {
|
||||
f.name: getattr(m, f.name)
|
||||
for f in types.Model.fields()
|
||||
if f.type in (int, bytes, np.ndarray)
|
||||
}
|
||||
|
||||
device_fields = {
|
||||
f.name: copy.copy(getattr(m, f.name)) # copy because device_put is async
|
||||
for f in types.Model.fields()
|
||||
if f.type is jax.Array
|
||||
}
|
||||
device_fields.update(mesh.get(m))
|
||||
device_fields = jax.device_put(device_fields, device=device)
|
||||
|
||||
return types.Model(
|
||||
opt=opt,
|
||||
stat=stat,
|
||||
**static_fields,
|
||||
**device_fields,
|
||||
)
|
||||
|
||||
|
||||
def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
|
||||
"""Allocate and initialize Data."""
|
||||
|
||||
ncon = collision_driver.ncon(m)
|
||||
ne, nf, nl, nc = constraint.count_constraints(m)
|
||||
nefc = ne + nf + nl + nc
|
||||
|
||||
zero_nv = jp.zeros(m.nv, dtype=jp.float32)
|
||||
zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32)
|
||||
zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32)
|
||||
zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32)
|
||||
zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32)
|
||||
zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=jp.float32)
|
||||
zero_nefc = jp.zeros(nefc, dtype=jp.float32)
|
||||
zero_na = jp.zeros(m.na, dtype=jp.float32)
|
||||
zero_nu = jp.zeros(m.nu, dtype=jp.float32)
|
||||
zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=jp.float32)
|
||||
zero_nm = jp.zeros(m.nM, dtype=jp.float32)
|
||||
|
||||
# create first d to get num contacts and nc
|
||||
d = Data(
|
||||
d = types.Data(
|
||||
solver_niter=jp.array(0, dtype=jp.int32),
|
||||
ne=0,
|
||||
nf=0,
|
||||
nl=0,
|
||||
nefc=0,
|
||||
ncon=0,
|
||||
time=jp.zeros((), dtype=jp.float32),
|
||||
time=jp.array(0.0),
|
||||
qpos=m.qpos0,
|
||||
qvel=jp.zeros(m.nv, dtype=jp.float32),
|
||||
act=jp.zeros(m.na, dtype=jp.float32),
|
||||
qacc_warmstart=jp.zeros(m.nv, dtype=jp.float32),
|
||||
ctrl=jp.zeros(m.nu, dtype=jp.float32),
|
||||
qfrc_applied=jp.zeros(m.nv, dtype=jp.float32),
|
||||
xfrc_applied=jp.zeros((m.nbody, 6), dtype=jp.float32),
|
||||
qvel=zero_nv,
|
||||
act=zero_na,
|
||||
qacc_warmstart=zero_nv,
|
||||
ctrl=zero_nu,
|
||||
qfrc_applied=zero_nv,
|
||||
xfrc_applied=zero_nbody_6,
|
||||
eq_active=jp.zeros(m.neq, dtype=jp.int32),
|
||||
qacc=jp.zeros(m.nv, dtype=jp.float32),
|
||||
act_dot=jp.zeros(m.na, dtype=jp.float32),
|
||||
xpos=jp.zeros((m.nbody, 3), dtype=jp.float32),
|
||||
qacc=zero_nv,
|
||||
act_dot=zero_na,
|
||||
xpos=zero_nbody_3,
|
||||
xquat=jp.zeros((m.nbody, 4), dtype=jp.float32),
|
||||
xmat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32),
|
||||
xipos=jp.zeros((m.nbody, 3), dtype=jp.float32),
|
||||
ximat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32),
|
||||
xanchor=jp.zeros((m.njnt, 3), dtype=jp.float32),
|
||||
xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32),
|
||||
xmat=zero_nbody_3_3,
|
||||
xipos=zero_nbody_3,
|
||||
ximat=zero_nbody_3_3,
|
||||
xanchor=zero_njnt_3,
|
||||
xaxis=zero_njnt_3,
|
||||
geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32),
|
||||
geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32),
|
||||
site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32),
|
||||
site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32),
|
||||
subtree_com=jp.zeros((m.nbody, 3), dtype=jp.float32),
|
||||
cdof=jp.zeros((m.nv, 6), dtype=jp.float32),
|
||||
cinert=jp.zeros((m.nbody, 10), dtype=jp.float32),
|
||||
actuator_length=jp.zeros(m.nu, dtype=jp.float32),
|
||||
subtree_com=zero_nbody_3,
|
||||
cdof=zero_nv_6,
|
||||
cinert=zero_nbody_10,
|
||||
actuator_length=zero_nu,
|
||||
actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32),
|
||||
crb=jp.zeros((m.nbody, 10), dtype=jp.float32),
|
||||
qM=jp.zeros(m.nM, dtype=jp.float32),
|
||||
qLD=jp.zeros(m.nM, dtype=jp.float32),
|
||||
qLDiagInv=jp.zeros(m.nv, dtype=jp.float32),
|
||||
qLDiagSqrtInv=jp.zeros(m.nv, dtype=jp.float32),
|
||||
contact=Contact.zero(),
|
||||
efc_J=jp.zeros((), dtype=jp.float32),
|
||||
efc_frictionloss=jp.zeros((), dtype=jp.float32),
|
||||
efc_D=jp.zeros((), dtype=jp.float32),
|
||||
actuator_velocity=jp.zeros(m.nu, dtype=jp.float32),
|
||||
cvel=jp.zeros((m.nbody, 6), dtype=jp.float32),
|
||||
cdof_dot=jp.zeros((m.nv, 6), dtype=jp.float32),
|
||||
qfrc_bias=jp.zeros(m.nv, dtype=jp.float32),
|
||||
qfrc_passive=jp.zeros(m.nv, dtype=jp.float32),
|
||||
efc_aref=jp.zeros((), dtype=jp.float32),
|
||||
actuator_force=jp.zeros(m.nu, dtype=jp.float32),
|
||||
qfrc_actuator=jp.zeros(m.nv, dtype=jp.float32),
|
||||
qfrc_smooth=jp.zeros(m.nv, dtype=jp.float32),
|
||||
qacc_smooth=jp.zeros(m.nv, dtype=jp.float32),
|
||||
qfrc_constraint=jp.zeros(m.nv, dtype=jp.float32),
|
||||
qfrc_inverse=jp.zeros(m.nv, dtype=jp.float32),
|
||||
efc_force=jp.zeros((), dtype=jp.float32),
|
||||
)
|
||||
|
||||
# get contact data with correct shapes
|
||||
ncon = collision_driver.ncon(m)
|
||||
d = d.replace(contact=Contact.zero((ncon,)), ncon=ncon)
|
||||
d = d.tree_replace({'contact.dim': 3 * np.ones(ncon)})
|
||||
|
||||
ne, nf, nl, nc = constraint.count_constraints(m, d)
|
||||
d = d.replace(ne=ne, nf=nf, nl=nl, nefc=ne + nf + nl + nc)
|
||||
ns = ne + nf + nl
|
||||
d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + ncon * 4, 4)})
|
||||
d = d.replace(
|
||||
efc_J=jp.zeros((d.nefc, m.nv), dtype=jp.float32),
|
||||
efc_frictionloss=jp.zeros(d.nefc, dtype=jp.float32),
|
||||
efc_D=jp.zeros(d.nefc, dtype=jp.float32),
|
||||
efc_aref=jp.zeros(d.nefc, dtype=jp.float32),
|
||||
efc_force=jp.zeros(d.nefc, dtype=jp.float32),
|
||||
crb=zero_nbody_10,
|
||||
qM=zero_nm,
|
||||
qLD=zero_nm,
|
||||
qLDiagInv=zero_nv,
|
||||
qLDiagSqrtInv=zero_nv,
|
||||
contact=types.Contact.zero(ncon),
|
||||
efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32),
|
||||
efc_frictionloss=zero_nefc,
|
||||
efc_D=zero_nefc,
|
||||
actuator_velocity=zero_nu,
|
||||
cvel=zero_nbody_6,
|
||||
cdof_dot=zero_nv_6,
|
||||
qfrc_bias=zero_nv,
|
||||
qfrc_passive=zero_nv,
|
||||
efc_aref=zero_nefc,
|
||||
actuator_force=zero_nu,
|
||||
qfrc_actuator=zero_nv,
|
||||
qfrc_smooth=zero_nv,
|
||||
qacc_smooth=zero_nv,
|
||||
qfrc_constraint=zero_nv,
|
||||
qfrc_inverse=zero_nv,
|
||||
efc_force=zero_nefc,
|
||||
)
|
||||
|
||||
return d
|
||||
|
||||
|
||||
def _get_contact(
|
||||
c: mujoco._structs._MjContactList,
|
||||
cx: types.Contact,
|
||||
efc_start: int,
|
||||
):
|
||||
"""Converts mjx.Contact to mujoco._structs._MjContactList."""
|
||||
con_id = np.nonzero(cx.dist <= 0)[0]
|
||||
for field in types.Contact.fields():
|
||||
value = getattr(cx, field.name)[con_id]
|
||||
if field.name == 'frame':
|
||||
value = value.reshape((-1, 9))
|
||||
getattr(c, field.name)[:] = value
|
||||
|
||||
ncon = con_id.shape[0]
|
||||
c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id]
|
||||
|
||||
|
||||
def get_data(
|
||||
m: mujoco.MjModel, d: types.Data
|
||||
) -> Union[mujoco.MjData, List[mujoco.MjData]]:
|
||||
"""Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData]."""
|
||||
dx = jax.device_get(d)
|
||||
batched = len(d.qpos.shape) > 1
|
||||
batch_size = d.qpos.shape[0] if batched else 1
|
||||
ne, nf, nl, nc = constraint.count_constraints(m)
|
||||
efc_type = np.array([
|
||||
mujoco.mjtConstraint.mjCNSTR_EQUALITY,
|
||||
mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF,
|
||||
mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT,
|
||||
mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL,
|
||||
]).repeat([ne, nf, nl, nc])
|
||||
|
||||
ds = []
|
||||
for i in range(batch_size):
|
||||
dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d
|
||||
ncon = (dx_i.contact.dist <= 0).sum()
|
||||
efc_active = (dx_i.efc_J != 0).any(axis=1)
|
||||
efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
|
||||
nefc, nc = efc_active.sum(), (efc_active & efc_con).sum()
|
||||
d_i = mujoco.MjData(m)
|
||||
d_i.nnzJ = nefc * m.nv
|
||||
mujoco._functions._realloc_con_efc(d_i, ncon=ncon, nefc=nefc) # pylint: disable=protected-access
|
||||
d_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc)
|
||||
d_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv)
|
||||
d_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
|
||||
|
||||
for field in types.Data.fields():
|
||||
if field.name == 'contact':
|
||||
_get_contact(d_i.contact, dx_i.contact, nefc - nc)
|
||||
continue
|
||||
|
||||
value = getattr(dx_i, field.name)
|
||||
|
||||
if field.name in ('xmat', 'ximat', 'geom_xmat'):
|
||||
value = value.reshape((-1, 9))
|
||||
|
||||
if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
|
||||
value = value[efc_active]
|
||||
|
||||
if field.name == 'efc_J':
|
||||
value = value[efc_active].reshape(-1)
|
||||
|
||||
if value.shape:
|
||||
getattr(d_i, field.name)[:] = value
|
||||
else:
|
||||
setattr(d_i, field.name, value)
|
||||
|
||||
d_i.efc_type[:] = efc_type[efc_active]
|
||||
ds.append(d_i)
|
||||
|
||||
return ds if batched else ds[0]
|
||||
|
||||
|
||||
def _put_contact(
|
||||
c: mujoco._structs._MjContactList, ncon: int, device=None
|
||||
) -> types.Contact:
|
||||
"""Puts mujoco.structs._MjContactList onto a device, resulting in mjx.Contact."""
|
||||
fields = {
|
||||
f.name: copy.copy(getattr(c, f.name)) for f in types.Contact.fields()
|
||||
}
|
||||
fields['frame'] = fields['frame'].reshape((-1, 3, 3))
|
||||
pad_size = ncon - c.dist.shape[0]
|
||||
pad_fn = lambda x: np.concatenate(
|
||||
(x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype))
|
||||
)
|
||||
fields = jax.tree_map(pad_fn, fields)
|
||||
fields['dist'][-pad_size:] = np.inf
|
||||
fields = jax.device_put(fields, device=device)
|
||||
|
||||
return types.Contact(**fields)
|
||||
|
||||
|
||||
def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
|
||||
"""Puts mujoco.MjData onto a device, resulting in mjx.Data."""
|
||||
ncon = collision_driver.ncon(m)
|
||||
ne, nf, nl, nc = constraint.count_constraints(m)
|
||||
nefc = ne + nf + nl + nc
|
||||
|
||||
for d_val, val, name in (
|
||||
(d.ncon, ncon, 'ncon'),
|
||||
(d.ne, ne, 'ne'),
|
||||
(d.nf, nf, 'nf'),
|
||||
(d.nl, nl, 'nl'),
|
||||
(d.nefc, nefc, 'nefc'),
|
||||
):
|
||||
if d_val > val:
|
||||
raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}')
|
||||
|
||||
fields = {
|
||||
f.name: copy.copy(getattr(d, f.name)) # copy because device_put is async
|
||||
for f in types.Data.fields()
|
||||
if f.type is jax.Array
|
||||
}
|
||||
|
||||
for fname in ('xmat', 'ximat', 'geom_xmat'):
|
||||
fields[fname] = fields[fname].reshape((-1, 3, 3))
|
||||
|
||||
# pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
|
||||
# efc_J is also optionally column-sparse (typically for large nv). MJX is
|
||||
# neither: it contains zeros for inactive constraints, and efc_J is always
|
||||
# (nefc, nv). this may change in the future.
|
||||
if mujoco.mj_isSparse(m):
|
||||
nr = d.efc_J_rownnz.shape[0]
|
||||
efc_j = np.zeros((nr, m.nv))
|
||||
for i in range(nr):
|
||||
rowadr = d.efc_J_rowadr[i]
|
||||
for j in range(d.efc_J_rownnz[i]):
|
||||
efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j]
|
||||
fields['efc_J'] = efc_j
|
||||
else:
|
||||
fields['efc_J'] = fields['efc_J'].reshape((-1, m.nv))
|
||||
|
||||
for fname in ('efc_J', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
|
||||
value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc)
|
||||
for i in range(4):
|
||||
value_beg = sum([ne, nf, nl][:i])
|
||||
d_beg = sum([d.ne, d.nf, d.nl][:i])
|
||||
size = [d.ne, d.nf, d.nl, d.nefc - d.nl - d.nf - d.ne][i]
|
||||
value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size]
|
||||
fields[fname] = value
|
||||
|
||||
fields = jax.device_put(fields, device=device)
|
||||
fields['contact'] = _put_contact(d.contact, ncon, device=device)
|
||||
|
||||
return types.Data(**fields)
|
||||
|
||||
+375
-11
@@ -17,26 +17,390 @@
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import jax
|
||||
from jax import numpy as jp
|
||||
import mujoco
|
||||
from mujoco import mjx
|
||||
from mujoco.mjx._src import test_util
|
||||
import numpy as np
|
||||
|
||||
|
||||
_MULTIPLE_CONVEX_OBJECTS = """
|
||||
<mujoco>
|
||||
<option timestep="0.001"/>
|
||||
<default>
|
||||
<geom solref=".006 1"/>
|
||||
</default>
|
||||
<asset>
|
||||
<mesh name="box" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1" scale="1 1 .1"/>
|
||||
<mesh name="boxoid" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 .5 -1 -1 2" scale=".3 .2 .1"/>
|
||||
<mesh name="pentaprism" vertex="1 0 0 0.309 0.951 0 -0.809 0.588 0 -0.809 -0.588 0 0.309 -0.951 0 1 0 1 0.309 0.951 1 -0.809 0.588 1 -0.809 -0.588 1 0.309 -0.951 1" scale=".2 .2 .1"/>
|
||||
</asset>
|
||||
<worldbody>
|
||||
<geom type="plane" pos="0 0 -.5" size="3 3 .01"/>
|
||||
<geom type="mesh" mesh="box" pos="0 0 -.15" euler="3 7 30"/>
|
||||
<body pos="-.3 -.3 .3">
|
||||
<freejoint/>
|
||||
<geom type="mesh" mesh="boxoid" rgba=".8 0 0 1" euler="3 5 -130"/>
|
||||
</body>
|
||||
<body pos=".3 .3 0.3">
|
||||
<freejoint/>
|
||||
<geom type="box" euler="3 5 -80" size=".3 .2 .1" rgba="0 .8 0 1"/>
|
||||
</body>
|
||||
<body pos=".3 .3 .6">
|
||||
<freejoint/>
|
||||
<geom type="mesh" mesh="pentaprism" rgba="0 0 .8 1"/>
|
||||
</body>
|
||||
<body pos=".6 -.3 .3">
|
||||
<joint name="joint" axis="1 0 0" type="hinge" range="-45 45"/>
|
||||
<geom type="capsule" size=".2 .05" rgba=".6 0 .6 1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<motor joint="joint"/>
|
||||
</actuator>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
_MULTIPLE_CONSTRAINTS = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<geom type="plane" size="3 3 .01"/>
|
||||
<body name="cap1" pos="-.3 -.3 .2">
|
||||
<freejoint/>
|
||||
<geom type="capsule" size=".2 .05"/>
|
||||
<body name="cap2" pos=".6 -.3 .3">
|
||||
<joint axis="0 1 0" type="hinge" range="-45 45"/>
|
||||
<joint axis="1 0 0" type="hinge" range="-0.001 0.001"/>
|
||||
<geom type="capsule" size=".2 .05"/>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
<equality>
|
||||
<connect body1="cap2" anchor="0 0 1"/>
|
||||
</equality>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
|
||||
class IoTest(parameterized.TestCase):
|
||||
|
||||
@parameterized.parameters(test_util.TEST_FILES)
|
||||
def test_make_data(self, fname):
|
||||
"""Test that data created by make_data matches data returned by step."""
|
||||
def test_put_model(self):
|
||||
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
|
||||
mx = mjx.put_model(m)
|
||||
self.assertEqual(mx.nq, m.nq)
|
||||
self.assertEqual(mx.nv, m.nv)
|
||||
self.assertEqual(mx.nu, m.nu)
|
||||
self.assertEqual(mx.na, m.na)
|
||||
self.assertEqual(mx.nbody, m.nbody)
|
||||
self.assertEqual(mx.njnt, m.njnt)
|
||||
self.assertEqual(mx.ngeom, m.ngeom)
|
||||
self.assertEqual(mx.nmesh, m.nmesh)
|
||||
self.assertEqual(mx.npair, m.npair)
|
||||
self.assertEqual(mx.nexclude, m.nexclude)
|
||||
self.assertEqual(mx.neq, m.neq)
|
||||
self.assertEqual(mx.nnumeric, m.nnumeric)
|
||||
self.assertEqual(mx.nM, m.nM)
|
||||
self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
|
||||
|
||||
m = test_util.load_test_file(fname)
|
||||
mx = mjx.device_put(m)
|
||||
dx = mjx.make_data(mx)
|
||||
dx_step = mjx.step(mx, dx)
|
||||
np.testing.assert_allclose(mx.body_parentid, m.body_parentid)
|
||||
np.testing.assert_allclose(mx.geom_type, m.geom_type)
|
||||
np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid)
|
||||
np.testing.assert_almost_equal(mx.geom_solref, m.geom_solref)
|
||||
np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos)
|
||||
self.assertLen(mx.geom_convex_face, 6)
|
||||
self.assertLen(mx.geom_convex_vert, 6)
|
||||
self.assertLen(mx.geom_convex_edge, 6)
|
||||
self.assertLen(mx.geom_convex_facenormal, 6)
|
||||
|
||||
_, dx_treedef = jax.tree_util.tree_flatten(dx)
|
||||
_, dx_step_treedef = jax.tree_util.tree_flatten(dx_step)
|
||||
np.testing.assert_allclose(mx.jnt_type, m.jnt_type)
|
||||
np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr)
|
||||
np.testing.assert_allclose(mx.jnt_bodyid, m.jnt_bodyid)
|
||||
np.testing.assert_allclose(mx.jnt_limited, m.jnt_limited)
|
||||
np.testing.assert_almost_equal(mx.jnt_axis, m.jnt_axis)
|
||||
|
||||
self.assertEqual(dx_treedef, dx_step_treedef)
|
||||
np.testing.assert_allclose(mx.actuator_trntype, m.actuator_trntype)
|
||||
np.testing.assert_allclose(mx.actuator_dyntype, m.actuator_dyntype)
|
||||
np.testing.assert_allclose(mx.actuator_gaintype, m.actuator_gaintype)
|
||||
np.testing.assert_allclose(mx.actuator_biastype, m.actuator_biastype)
|
||||
np.testing.assert_allclose(mx.actuator_trnid, m.actuator_trnid)
|
||||
|
||||
def test_fluid_params(self):
|
||||
"""Test that has_fluid_params is set when fluid params are present."""
|
||||
m = mjx.put_model(
|
||||
mujoco.MjModel.from_xml_string(
|
||||
'<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
|
||||
)
|
||||
)
|
||||
self.assertTrue(m.opt.has_fluid_params)
|
||||
|
||||
def test_put_model_implicit_not_implemented(self):
|
||||
"""Test that MJX guards against models with unimplemented features."""
|
||||
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.put_model(
|
||||
mujoco.MjModel.from_xml_string(
|
||||
'<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
|
||||
)
|
||||
)
|
||||
|
||||
def test_put_model_cone_not_implemented(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.put_model(
|
||||
mujoco.MjModel.from_xml_string(
|
||||
'<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
|
||||
)
|
||||
)
|
||||
|
||||
def test_put_model_pgs_not_implemented(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.put_model(
|
||||
mujoco.MjModel.from_xml_string(
|
||||
'<mujoco><option solver="PGS"/><worldbody/></mujoco>'
|
||||
)
|
||||
)
|
||||
|
||||
def test_put_model_site_actuator_not_implemented(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.put_model(mujoco.MjModel.from_xml_string("""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<site name="ref"/>
|
||||
<body>
|
||||
<site name="end"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<general site="end" refsite="ref"/>
|
||||
</actuator>
|
||||
</mujoco>"""))
|
||||
|
||||
def test_put_model_tendon_not_implemented(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.put_model(mujoco.MjModel.from_xml_string("""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body>
|
||||
<joint name="left_hip" type="hinge"/>
|
||||
<geom size="0.05"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<tendon>
|
||||
<fixed>
|
||||
<joint coef="1" joint="left_hip"/>
|
||||
</fixed>
|
||||
</tendon>
|
||||
</mujoco>"""))
|
||||
|
||||
def test_put_model_condim_not_implemented(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.put_model(mujoco.MjModel.from_xml_string("""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body>
|
||||
<freejoint/>
|
||||
<geom size="0.05" condim="1"/>
|
||||
</body>
|
||||
<body>
|
||||
<freejoint/>
|
||||
<geom size="0.05" condim="1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>"""))
|
||||
|
||||
def test_put_model_cylinder_not_implemented(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.put_model(mujoco.MjModel.from_xml_string("""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body>
|
||||
<freejoint/>
|
||||
<geom type="cylinder" size="0.05 0.05"/>
|
||||
</body>
|
||||
<body>
|
||||
<freejoint/>
|
||||
<geom size="0.05"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>"""))
|
||||
|
||||
def test_make_data(self):
|
||||
"""Test that make_data returns the correct shapes."""
|
||||
|
||||
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
|
||||
d = mjx.make_data(m)
|
||||
|
||||
nq = 22
|
||||
nbody = 5
|
||||
ncon = 46
|
||||
nv = 19
|
||||
nefc = 185
|
||||
nm = 64
|
||||
|
||||
self.assertEqual(d.qpos.shape, (nq,))
|
||||
self.assertEqual(d.qvel.shape, (nv,))
|
||||
self.assertEqual(d.act.shape, (0,))
|
||||
self.assertEqual(d.qacc_warmstart.shape, (nv,))
|
||||
self.assertEqual(d.ctrl.shape, (1,))
|
||||
self.assertEqual(d.qfrc_applied.shape, (nv,))
|
||||
self.assertEqual(d.xfrc_applied.shape, (nbody, 6))
|
||||
self.assertEqual(d.eq_active.shape, (0,))
|
||||
self.assertEqual(d.qacc.shape, (nv,))
|
||||
self.assertEqual(d.act_dot.shape, (0,))
|
||||
self.assertEqual(d.xpos.shape, (nbody, 3))
|
||||
self.assertEqual(d.xquat.shape, (nbody, 4))
|
||||
self.assertEqual(d.xmat.shape, (nbody, 3, 3))
|
||||
self.assertEqual(d.xipos.shape, (nbody, 3))
|
||||
self.assertEqual(d.ximat.shape, (nbody, 3, 3))
|
||||
self.assertEqual(d.xanchor.shape, (4, 3))
|
||||
self.assertEqual(d.xaxis.shape, (4, 3))
|
||||
self.assertEqual(d.geom_xpos.shape, (6, 3))
|
||||
self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
|
||||
self.assertEqual(d.subtree_com.shape, (nbody, 3))
|
||||
self.assertEqual(d.cdof.shape, (nv, 6))
|
||||
self.assertEqual(d.cinert.shape, (nbody, 10))
|
||||
self.assertEqual(d.crb.shape, (nbody, 10))
|
||||
self.assertEqual(d.actuator_length.shape, (1,))
|
||||
self.assertEqual(d.actuator_moment.shape, (1, nv))
|
||||
self.assertEqual(d.qM.shape, (nm,))
|
||||
self.assertEqual(d.qLD.shape, (nm,))
|
||||
self.assertEqual(d.qLDiagInv.shape, (nv,))
|
||||
self.assertEqual(d.qLDiagSqrtInv.shape, (nv,))
|
||||
self.assertEqual(d.contact.dist.shape, (ncon,))
|
||||
self.assertEqual(d.contact.pos.shape, (ncon, 3))
|
||||
self.assertEqual(d.contact.frame.shape, (ncon, 3, 3))
|
||||
self.assertEqual(d.contact.solref.shape, (ncon, 2))
|
||||
self.assertEqual(d.contact.solimp.shape, (ncon, 5))
|
||||
self.assertEqual(d.contact.geom1.shape, (ncon,))
|
||||
self.assertEqual(d.contact.geom2.shape, (ncon,))
|
||||
self.assertEqual(d.efc_J.shape, (nefc, nv))
|
||||
self.assertEqual(d.efc_frictionloss.shape, (nefc,))
|
||||
self.assertEqual(d.efc_D.shape, (nefc,))
|
||||
self.assertEqual(d.actuator_velocity.shape, (1,))
|
||||
self.assertEqual(d.cvel.shape, (nbody, 6))
|
||||
self.assertEqual(d.cdof_dot.shape, (nv, 6))
|
||||
self.assertEqual(d.qfrc_bias.shape, (nv,))
|
||||
self.assertEqual(d.qfrc_passive.shape, (nv,))
|
||||
self.assertEqual(d.efc_aref.shape, (nefc,))
|
||||
self.assertEqual(d.actuator_force.shape, (1,))
|
||||
self.assertEqual(d.qfrc_actuator.shape, (nv,))
|
||||
self.assertEqual(d.qfrc_smooth.shape, (nv,))
|
||||
self.assertEqual(d.qacc_smooth.shape, (nv,))
|
||||
self.assertEqual(d.qfrc_constraint.shape, (nv,))
|
||||
self.assertEqual(d.qfrc_inverse.shape, (nv,))
|
||||
self.assertEqual(d.efc_force.shape, (nefc,))
|
||||
|
||||
def test_put_data(self):
|
||||
"""Test that put_data puts the correct data for dense and sparse."""
|
||||
|
||||
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_step(m, d, 2)
|
||||
dx = mjx.put_data(m, d)
|
||||
|
||||
# check a few fields
|
||||
np.testing.assert_allclose(dx.qpos, d.qpos)
|
||||
np.testing.assert_allclose(dx.xpos, d.xpos)
|
||||
np.testing.assert_allclose(dx.cvel, d.cvel)
|
||||
np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
|
||||
np.testing.assert_allclose(dx.qM, d.qM)
|
||||
|
||||
# 4 contacts, 2 for each capsule against the plane
|
||||
self.assertEqual(dx.contact.dist.shape, (4,))
|
||||
self.assertEqual(d.ncon, 1) # however only 1 contact in this step
|
||||
np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0])
|
||||
self.assertTrue(np.isinf(dx.contact.dist[1:]).all())
|
||||
self.assertEqual(dx.contact.frame.shape, (4, 3, 3))
|
||||
np.testing.assert_allclose(
|
||||
dx.contact.frame[0].reshape(9), d.contact.frame[0]
|
||||
)
|
||||
np.testing.assert_allclose(dx.contact.frame[1:], 0)
|
||||
|
||||
# xmat, ximat, geom_xmat are all shape transformed
|
||||
self.assertEqual(dx.xmat.shape, (3, 3, 3))
|
||||
self.assertEqual(dx.ximat.shape, (3, 3, 3))
|
||||
self.assertEqual(dx.geom_xmat.shape, (3, 3, 3))
|
||||
np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat)
|
||||
np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat)
|
||||
np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat)
|
||||
|
||||
# efc_ are also shape transformed and padded
|
||||
self.assertEqual(dx.efc_J.shape, (21, 8)) # nefc, nv
|
||||
d_efc_j = d.efc_J.reshape((-1, 8))
|
||||
np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3]) # connect eq
|
||||
np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3]) # one active limit
|
||||
np.testing.assert_allclose(dx.efc_J[4], 0) # one inactive limit
|
||||
np.testing.assert_allclose(dx.efc_J[5:9], d_efc_j[4:8]) # contact
|
||||
np.testing.assert_allclose(dx.efc_J[9:], 0) # no contact
|
||||
|
||||
# check another efc_ too
|
||||
self.assertEqual(dx.efc_aref.shape, (21,)) # nefc
|
||||
np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3])
|
||||
np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3])
|
||||
np.testing.assert_allclose(dx.efc_aref[4], 0)
|
||||
np.testing.assert_allclose(dx.efc_aref[5:9], d.efc_aref[4:8])
|
||||
np.testing.assert_allclose(dx.efc_aref[9:], 0)
|
||||
|
||||
# check sparse transform is correct
|
||||
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_step(m, d, 2)
|
||||
dx_from_sparse = mjx.put_data(m, d)
|
||||
np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8)
|
||||
|
||||
def test_get_data(self):
|
||||
"""Test that get_data makes correct MjData."""
|
||||
|
||||
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_step(m, d, 2)
|
||||
dx = mjx.put_data(m, d)
|
||||
d_2: mujoco.MjData = mjx.get_data(m, dx)
|
||||
|
||||
# check a few fields
|
||||
np.testing.assert_allclose(d_2.qpos, d.qpos)
|
||||
np.testing.assert_allclose(d_2.xpos, d.xpos)
|
||||
np.testing.assert_allclose(d_2.cvel, d.cvel)
|
||||
np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot)
|
||||
np.testing.assert_allclose(d_2.qM, d.qM)
|
||||
|
||||
# only 1 contact active
|
||||
self.assertEqual(d_2.contact.dist.shape, (1,))
|
||||
self.assertEqual(d_2.ncon, 1)
|
||||
np.testing.assert_allclose(d_2.contact.dist, d.contact.dist)
|
||||
self.assertEqual(d_2.contact.frame.shape, (1, 9))
|
||||
np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
|
||||
|
||||
# xmat, ximat, geom_xmat are all shape transformed
|
||||
self.assertEqual(d_2.xmat.shape, (3, 9))
|
||||
self.assertEqual(d_2.ximat.shape, (3, 9))
|
||||
self.assertEqual(d_2.geom_xmat.shape, (3, 9))
|
||||
np.testing.assert_allclose(d_2.xmat, d.xmat)
|
||||
np.testing.assert_allclose(d_2.ximat, d.ximat)
|
||||
np.testing.assert_allclose(d_2.geom_xmat, d.geom_xmat)
|
||||
|
||||
# efc_* are also shape transformed and filtered
|
||||
self.assertEqual(d_2.efc_J.shape, (64,)) # nefc * nv
|
||||
np.testing.assert_allclose(d_2.efc_J, d.efc_J)
|
||||
self.assertEqual(d_2.efc_aref.shape, (8,)) # nefc
|
||||
np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
|
||||
|
||||
# efc_address is created on demand
|
||||
np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
|
||||
|
||||
def test_get_data_batched(self):
|
||||
"""Test that get_data makes correct List[MjData] for batched Data."""
|
||||
|
||||
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_step(m, d, 2)
|
||||
dx = mjx.put_data(m, d)
|
||||
# second data in batch has contact dist > 0, disables contact
|
||||
dx_b = jax.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
|
||||
ds = mjx.get_data(m, dx_b)
|
||||
self.assertLen(ds, 2)
|
||||
np.testing.assert_allclose(ds[0].qpos, d.qpos)
|
||||
np.testing.assert_allclose(ds[1].qpos, d.qpos + 0.05, atol=1e-8)
|
||||
self.assertEqual(ds[0].ncon, 1)
|
||||
self.assertEqual(ds[1].ncon, 0)
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
|
||||
@@ -221,7 +221,7 @@ def flat(
|
||||
'u': i,
|
||||
'a': m.actuator_actadr[i],
|
||||
'j': (
|
||||
m.actuator_trnid[i]
|
||||
m.actuator_trnid[i, 0]
|
||||
if m.actuator_trntype[i] == TrnType.JOINT
|
||||
else np.array(-1)
|
||||
),
|
||||
|
||||
@@ -210,15 +210,16 @@ class ScanTest(absltest.TestCase):
|
||||
m, fn, 'ujqva', 'ujqva', *args, group_by='u'
|
||||
)
|
||||
|
||||
actuator_trnid = m.actuator_trnid[:, 0]
|
||||
np.testing.assert_array_equal(gear, m.actuator_gear)
|
||||
np.testing.assert_array_equal(jnt_typ, m.jnt_type[m.actuator_trnid])
|
||||
np.testing.assert_array_equal(jnt_typ, m.jnt_type[actuator_trnid])
|
||||
np.testing.assert_array_equal(act, jp.array([1.4, 1.1]))
|
||||
expected_vadr = np.concatenate(
|
||||
[np.nonzero(m.dof_jntid == trnid)[0] for trnid in m.actuator_trnid]
|
||||
[np.nonzero(m.dof_jntid == trnid)[0] for trnid in actuator_trnid]
|
||||
)
|
||||
np.testing.assert_array_equal(vadr, expected_vadr)
|
||||
expected_qadr = np.concatenate(
|
||||
[np.nonzero(scan._q_jointid(m) == i)[0] for i in m.actuator_trnid]
|
||||
[np.nonzero(scan._q_jointid(m) == i)[0] for i in actuator_trnid]
|
||||
)
|
||||
np.testing.assert_array_equal(qadr, expected_qadr)
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ from typing import Optional
|
||||
import jax
|
||||
from jax import numpy as jp
|
||||
import mujoco
|
||||
from mujoco.mjx._src import constraint
|
||||
from mujoco.mjx._src import math
|
||||
from mujoco.mjx._src import smooth
|
||||
# pylint: disable=g-importing-member
|
||||
@@ -69,12 +70,12 @@ class _Context(PyTreeNode):
|
||||
# TODO(robotics-team): determine nv at which sparse mul is faster
|
||||
M = smooth.dense_m(m, d) if m.nv < 100 else None # pylint: disable=invalid-name
|
||||
ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc
|
||||
nv_0 = jp.zeros((m.nv,))
|
||||
nv_0 = jp.zeros(m.nv)
|
||||
ctx = _Context(
|
||||
qacc=d.qacc,
|
||||
qfrc_constraint=d.qfrc_constraint,
|
||||
Jaref=jaref,
|
||||
efc_force=jp.zeros(d.nefc),
|
||||
efc_force=d.efc_force,
|
||||
M=M,
|
||||
Ma=ma,
|
||||
grad=nv_0,
|
||||
@@ -111,7 +112,7 @@ class _LSPoint(PyTreeNode):
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
d: Data,
|
||||
m: Model,
|
||||
ctx: _Context,
|
||||
alpha: jax.Array,
|
||||
jv: jax.Array,
|
||||
@@ -122,7 +123,8 @@ class _LSPoint(PyTreeNode):
|
||||
# roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
|
||||
|
||||
# TODO(robotics-team): change this to support friction constraints
|
||||
active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
|
||||
ne, nf, *_ = constraint.count_constraints(m)
|
||||
active = ((ctx.Jaref + alpha * jv) < 0).at[:ne + nf].set(True)
|
||||
quad = jax.vmap(jp.multiply)(quad, active) # only active
|
||||
quad_total = quad_gauss + jp.sum(quad, axis=0)
|
||||
|
||||
@@ -177,12 +179,11 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
|
||||
Returns:
|
||||
context with new constraint force and costs
|
||||
"""
|
||||
del m
|
||||
|
||||
# TODO(robotics-team): add friction constraints
|
||||
|
||||
# only count active constraints
|
||||
active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
|
||||
ne, nf, *_ = constraint.count_constraints(m)
|
||||
active = (ctx.Jaref < 0).at[:ne + nf].set(True)
|
||||
|
||||
efc_force = d.efc_D * -ctx.Jaref * active
|
||||
qfrc_constraint = d.efc_J.T @ efc_force
|
||||
@@ -221,7 +222,8 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
|
||||
if m.opt.solver == SolverType.CG:
|
||||
mgrad = smooth.solve_m(m, d, grad)
|
||||
elif m.opt.solver == SolverType.NEWTON:
|
||||
active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
|
||||
ne, nf, *_ = constraint.count_constraints(m)
|
||||
active = (ctx.Jaref < 0).at[:ne + nf].set(True)
|
||||
h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
|
||||
h = smooth.dense_m(m, d) + h
|
||||
h_ = jax.scipy.linalg.cho_factor(h)
|
||||
@@ -265,7 +267,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
|
||||
quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
|
||||
quad = (quad * d.efc_D).T
|
||||
|
||||
point_fn = lambda alpha: _LSPoint.create(d, ctx, alpha, jv, quad, quad_gauss)
|
||||
point_fn = lambda a: _LSPoint.create(m, ctx, a, jv, quad, quad_gauss)
|
||||
|
||||
def cond(ctx: _LSContext) -> jax.Array:
|
||||
done = ctx.ls_iter >= m.opt.ls_iterations
|
||||
|
||||
@@ -21,6 +21,7 @@ import jax
|
||||
import jax.numpy as jp
|
||||
import mujoco
|
||||
# pylint: disable=g-importing-member
|
||||
from mujoco.mjx._src import dataclasses
|
||||
from mujoco.mjx._src.dataclasses import PyTreeNode
|
||||
# pylint: enable=g-importing-member
|
||||
import numpy as np
|
||||
@@ -213,7 +214,6 @@ class Option(PyTreeNode):
|
||||
integrator: integration mode
|
||||
cone: type of friction cone
|
||||
solver: solver algorithm
|
||||
integrator: integration mode
|
||||
iterations: number of main solver iterations
|
||||
ls_iterations: maximum number of CG/Newton linesearch iterations
|
||||
disableflags: bit flags for disabling standard features
|
||||
@@ -329,8 +329,8 @@ class Model(PyTreeNode):
|
||||
geom_margin: include in solver if dist<margin-gap (ngeom,)
|
||||
geom_gap: include in solver if dist<margin-gap (ngeom,)
|
||||
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)
|
||||
site_pos: local position offset rel. to body (nsite, 3)
|
||||
site_quat: local orientation offset rel. to body (nsite, 4)
|
||||
geom_convex_face: vertex face data, MJX only (ngeom,)
|
||||
geom_convex_vert: vertex data, MJX only (ngeom,)
|
||||
geom_convex_edge: unique edge data, MJX only (ngeom,)
|
||||
@@ -508,10 +508,8 @@ class Contact(PyTreeNode):
|
||||
solref: constraint solver reference, normal direction (mjNREF,)
|
||||
solreffriction: constraint solver reference, friction directions (mjNREF,)
|
||||
solimp: constraint solver impedance (mjNIMP,)
|
||||
dim: contact space dimensionality: 1, 3, 4 or 6
|
||||
geom1: id of geom 1
|
||||
geom2: id of geom 2
|
||||
efc_address: address in efc; -1: not included
|
||||
"""
|
||||
dist: jax.Array
|
||||
pos: jax.Array
|
||||
@@ -521,29 +519,25 @@ class Contact(PyTreeNode):
|
||||
solref: jax.Array
|
||||
solreffriction: jax.Array
|
||||
solimp: jax.Array
|
||||
# unsupported: mu, H
|
||||
dim: np.ndarray
|
||||
# unsupported: mu, H, dim
|
||||
geom1: jax.Array
|
||||
geom2: jax.Array
|
||||
efc_address: np.ndarray
|
||||
# unsupported: exclude
|
||||
# unsupported: efc_address, exclude
|
||||
|
||||
@classmethod
|
||||
def zero(cls, shape=(0,)) -> 'Contact':
|
||||
def zero(cls, ncon: int = 0) -> 'Contact':
|
||||
"""Returns a contact filled with zeros."""
|
||||
return Contact(
|
||||
dist=jp.zeros(shape),
|
||||
pos=jp.zeros(shape + (3,)),
|
||||
frame=jp.zeros(shape + (3, 3)),
|
||||
includemargin=jp.zeros(shape),
|
||||
friction=jp.zeros(shape + (5,)),
|
||||
solref=jp.zeros(shape + (mujoco.mjNREF,)),
|
||||
solreffriction=jp.zeros(shape + (mujoco.mjNREF,)),
|
||||
solimp=jp.zeros(shape + (mujoco.mjNIMP,)),
|
||||
dim=np.zeros(shape, dtype=np.int32),
|
||||
geom1=jp.zeros(shape, dtype=jp.int32),
|
||||
geom2=jp.zeros(shape, dtype=jp.int32),
|
||||
efc_address=np.zeros(shape, dtype=np.int32),
|
||||
dist=jp.zeros(ncon),
|
||||
pos=jp.zeros((ncon, 3,)),
|
||||
frame=jp.zeros((ncon, 3, 3)),
|
||||
includemargin=jp.zeros(ncon),
|
||||
friction=jp.zeros((ncon, 5)),
|
||||
solref=jp.zeros((ncon, mujoco.mjNREF)),
|
||||
solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
|
||||
solimp=jp.zeros((ncon, mujoco.mjNIMP,)),
|
||||
geom1=jp.zeros(ncon, dtype=jp.int32),
|
||||
geom2=jp.zeros(ncon, dtype=jp.int32),
|
||||
)
|
||||
|
||||
|
||||
@@ -552,11 +546,6 @@ class Data(PyTreeNode):
|
||||
|
||||
Attributes:
|
||||
solver_niter: number of solver iterations, per island (mjNISLAND,)
|
||||
ne: number of equality constraints
|
||||
nf: number of friction constraints
|
||||
nl: number of limit constraints
|
||||
nefc: number of constraints
|
||||
ncon: nubmer of contacts
|
||||
time: simulation time
|
||||
qpos: position (nq,)
|
||||
qvel: velocity (nv,)
|
||||
@@ -610,12 +599,6 @@ class Data(PyTreeNode):
|
||||
"""
|
||||
# solver statistics:
|
||||
solver_niter: jax.Array
|
||||
# sizes (variable in MJ, constant in MJX)
|
||||
ne: int
|
||||
nf: int
|
||||
nl: int
|
||||
nefc: int
|
||||
ncon: int
|
||||
# global properties:
|
||||
time: jax.Array
|
||||
# state:
|
||||
|
||||
@@ -82,7 +82,6 @@ class CollisionDriverIntegrationTest(parameterized.TestCase):
|
||||
mjx_contact = jax.tree_map(
|
||||
lambda x: x.take(np.array(idx), axis=0), dx.contact
|
||||
)
|
||||
mjx_contact = mjx_contact.replace(dim=mjx_contact.dim[idx])
|
||||
for field in dataclasses.fields(Contact):
|
||||
_assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user