Erik Frey
71333938ec
PiperOrigin-RevId: 627136131 Change-Id: Ia62eb20cb80395b569c7fc7f4d4ef87e12eaaf40
423 lines
14 KiB
Python
423 lines
14 KiB
Python
# Copyright 2023 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Core non-smooth constraint functions."""
|
|
|
|
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
|
|
from mujoco.mjx._src.dataclasses import PyTreeNode
|
|
from mujoco.mjx._src.types import ConstraintType
|
|
from mujoco.mjx._src.types import Contact
|
|
from mujoco.mjx._src.types import Data
|
|
from mujoco.mjx._src.types import DisableBit
|
|
from mujoco.mjx._src.types import EqType
|
|
from mujoco.mjx._src.types import JointType
|
|
from mujoco.mjx._src.types import Model
|
|
# pylint: enable=g-importing-member
|
|
import numpy as np
|
|
|
|
|
|
_CONDIM_EFC_COUNT = {1: 1, 3: 4, 4: 6, 6: 10}
|
|
|
|
|
|
class _Efc(PyTreeNode):
|
|
"""Support data for creating constraint matrices."""
|
|
J: jax.Array
|
|
pos: jax.Array
|
|
pos_norm: jax.Array
|
|
invweight: jax.Array
|
|
solref: jax.Array
|
|
solimp: jax.Array
|
|
frictionloss: jax.Array
|
|
|
|
|
|
def _kbi(
|
|
m: Model,
|
|
solref: jax.Array,
|
|
solimp: jax.Array,
|
|
pos: jax.Array,
|
|
) -> Tuple[jax.Array, jax.Array, jax.Array]:
|
|
"""Calculates stiffness, damping, and impedance of a constraint."""
|
|
timeconst, dampratio = solref
|
|
|
|
if not m.opt.disableflags & DisableBit.REFSAFE:
|
|
timeconst = jp.maximum(timeconst, 2 * m.opt.timestep) * (timeconst > 0)
|
|
|
|
dmin, dmax, width, mid, power = solimp
|
|
|
|
dmin = jp.clip(dmin, mujoco.mjMINIMP, mujoco.mjMAXIMP)
|
|
dmax = jp.clip(dmax, mujoco.mjMINIMP, mujoco.mjMAXIMP)
|
|
width = jp.maximum(0, width)
|
|
mid = jp.clip(mid, mujoco.mjMINIMP, mujoco.mjMAXIMP)
|
|
power = jp.maximum(1, power)
|
|
|
|
# See https://mujoco.readthedocs.io/en/latest/modeling.html#solver-parameters
|
|
k = 1 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio)
|
|
b = 2 / (dmax * timeconst)
|
|
# TODO(robotics-simulation): check various solparam settings in model gen test
|
|
k = jp.where(dampratio <= 0, -dampratio / (dmax * dmax), k)
|
|
b = jp.where(timeconst <= 0, -timeconst / dmax, b)
|
|
|
|
imp_x = jp.abs(pos) / width
|
|
imp_a = (1.0 / jp.power(mid, power - 1)) * jp.power(imp_x, power)
|
|
imp_b = 1 - (1.0 / jp.power(1 - mid, power - 1)) * jp.power(1 - imp_x, power)
|
|
imp_y = jp.where(imp_x < mid, imp_a, imp_b)
|
|
imp = dmin + imp_y * (dmax - dmin)
|
|
imp = jp.clip(imp, dmin, dmax)
|
|
imp = jp.where(imp_x > 1.0, dmax, imp)
|
|
|
|
return k, b, imp # corresponds to K, B, I of efc_KBIP
|
|
|
|
|
|
def _instantiate_equality_connect(m: Model, d: Data) -> Optional[_Efc]:
|
|
"""Calculates constraint rows for connect equality constraints."""
|
|
|
|
ids = np.nonzero(m.eq_type == EqType.CONNECT)[0]
|
|
|
|
if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
|
|
return None
|
|
|
|
id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
|
|
|
|
@jax.vmap
|
|
def fn(data, id1, id2):
|
|
anchor1, anchor2 = data[0:3], data[3:6]
|
|
# find global points
|
|
pos1 = d.xmat[id1] @ anchor1 + d.xpos[id1]
|
|
pos2 = d.xmat[id2] @ anchor2 + d.xpos[id2]
|
|
|
|
# compute position error
|
|
cpos = pos1 - pos2
|
|
|
|
# compute Jacobian difference (opposite of contact: 0 - 1)
|
|
jacp1, _ = support.jac(m, d, pos1, id1)
|
|
jacp2, _ = support.jac(m, d, pos2, id2)
|
|
j = (jacp1 - jacp2).T
|
|
|
|
return j, cpos, jp.repeat(math.norm(cpos), 3)
|
|
|
|
# concatenate to drop connect grouping dimension
|
|
j, pos, pos_norm = jax.tree_util.tree_map(jp.concatenate, fn(data, id1, id2))
|
|
invweight = m.body_invweight0[id1, 0] + m.body_invweight0[id2, 0]
|
|
invweight = jp.repeat(invweight, 3)
|
|
solref = jp.tile(m.eq_solref[ids], (3, 1))
|
|
solimp = jp.tile(m.eq_solimp[ids], (3, 1))
|
|
frictionloss = jp.zeros_like(pos_norm)
|
|
|
|
return _Efc(j, pos, pos_norm, invweight, solref, solimp, frictionloss)
|
|
|
|
|
|
def _instantiate_equality_weld(m: Model, d: Data) -> Optional[_Efc]:
|
|
"""Calculates constraint rows for weld equality constraints."""
|
|
|
|
ids = np.nonzero(m.eq_type == EqType.WELD)[0]
|
|
|
|
if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
|
|
return None
|
|
|
|
id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
|
|
|
|
@jax.vmap
|
|
def fn(data, id1, id2):
|
|
anchor1, anchor2 = data[0:3], data[3:6]
|
|
relpose, torquescale = data[6:10], data[10]
|
|
|
|
# find global points
|
|
pos1 = d.xmat[id1] @ anchor2 + d.xpos[id1]
|
|
pos2 = d.xmat[id2] @ anchor1 + d.xpos[id2]
|
|
|
|
# compute position error
|
|
cpos = pos1 - pos2
|
|
|
|
# compute Jacobian difference (opposite of contact: 0 - 1)
|
|
jacp1, jacr1 = support.jac(m, d, pos1, id1)
|
|
jacp2, jacr2 = support.jac(m, d, pos2, id2)
|
|
jacdifp = jacp1 - jacp2
|
|
jacdifr = (jacr1 - jacr2) * torquescale
|
|
|
|
# compute orientation error: neg(q1) * q0 * relpose (axis components only)
|
|
quat = math.quat_mul(d.xquat[id1], relpose)
|
|
quat1 = math.quat_inv(d.xquat[id2])
|
|
crot = math.quat_mul(quat1, quat)[1:] # copy axis components
|
|
|
|
# correct rotation Jacobian: 0.5 * neg(q1) * (jac0-jac1) * q0 * relpose
|
|
jac_fn = lambda j: math.quat_mul(math.quat_mul_axis(quat1, j), quat)[1:]
|
|
jacdifr = 0.5 * jax.vmap(jac_fn)(jacdifr)
|
|
|
|
j = jp.concatenate((jacdifp.T, jacdifr.T))
|
|
pos = jp.concatenate((cpos, crot * torquescale))
|
|
|
|
return j, pos, jp.repeat(math.norm(pos), 6)
|
|
|
|
# concatenate to drop weld grouping dimension
|
|
j, pos, pos_norm = jax.tree_util.tree_map(jp.concatenate, fn(data, id1, id2))
|
|
invweight = m.body_invweight0[id1] + m.body_invweight0[id2]
|
|
invweight = jp.repeat(invweight, 3)
|
|
solref = jp.tile(m.eq_solref[ids], (6, 1))
|
|
solimp = jp.tile(m.eq_solimp[ids], (6, 1))
|
|
frictionloss = jp.zeros_like(pos_norm)
|
|
|
|
return _Efc(j, pos, pos_norm, invweight, solref, solimp, frictionloss)
|
|
|
|
|
|
def _instantiate_equality_joint(m: Model, d: Data) -> Optional[_Efc]:
|
|
"""Calculates constraint rows for joint equality constraints."""
|
|
|
|
ids = np.nonzero(m.eq_type == EqType.JOINT)[0]
|
|
|
|
if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
|
|
return None
|
|
|
|
id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
|
|
dofadr1, dofadr2 = m.jnt_dofadr[id1], m.jnt_dofadr[id2]
|
|
qposadr1, qposadr2 = m.jnt_qposadr[id1], m.jnt_qposadr[id2]
|
|
|
|
@jax.vmap
|
|
def fn(data, id2, dofadr1, dofadr2, qposadr1, qposadr2):
|
|
pos1, pos2 = d.qpos[qposadr1], d.qpos[qposadr2]
|
|
ref1, ref2 = m.qpos0[qposadr1], m.qpos0[qposadr2]
|
|
pos2, ref2 = pos2 * (id2 > -1), ref2 * (id2 > -1)
|
|
|
|
dif = pos2 - ref2
|
|
dif_power = jp.power(dif, jp.arange(0, 5))
|
|
|
|
deriv = jp.dot(data[1:5], dif_power[:4] * jp.arange(1, 5))
|
|
j = jp.zeros((m.nv)).at[dofadr1].set(1.0).at[dofadr2].set(-deriv)
|
|
pos = pos1 - ref1 - jp.dot(data[:5], dif_power)
|
|
return j, pos
|
|
|
|
j, pos = fn(data, id2, dofadr1, dofadr2, qposadr1, qposadr2)
|
|
invweight = m.dof_invweight0[dofadr1] + m.dof_invweight0[dofadr2] * (id2 > -1)
|
|
solref, solimp = m.eq_solref[ids], m.eq_solimp[ids]
|
|
frictionloss = jp.zeros_like(pos)
|
|
|
|
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
|
|
|
|
|
|
def _instantiate_friction(m: Model, d: Data) -> Optional[_Efc]:
|
|
# TODO(robotics-team): implement _instantiate_friction
|
|
del m, d
|
|
return None
|
|
|
|
|
|
def _instantiate_limit_ball(m: Model, d: Data) -> Optional[_Efc]:
|
|
"""Calculates constraint rows for ball joint limits."""
|
|
|
|
ids = np.nonzero((m.jnt_type == JointType.BALL) & m.jnt_limited)[0]
|
|
|
|
if (m.opt.disableflags & DisableBit.LIMIT) or ids.size == 0:
|
|
return None
|
|
|
|
jnt_range = m.jnt_range[ids]
|
|
jnt_margin = m.jnt_margin[ids]
|
|
qposadr = np.array([np.arange(q, q + 4) for q in m.jnt_qposadr[ids]])
|
|
dofadr = np.array([np.arange(d, d + 3) for d in m.jnt_dofadr[ids]])
|
|
|
|
@jax.vmap
|
|
def fn(jnt_range, jnt_margin, qposadr, dofadr):
|
|
axis, angle = math.quat_to_axis_angle(d.qpos[qposadr])
|
|
j = jp.zeros(m.nv).at[dofadr].set(-axis)
|
|
pos = jp.amax(jnt_range) - angle - jnt_margin
|
|
active = pos < 0
|
|
return j * active, pos * active
|
|
|
|
j, pos = fn(jnt_range, jnt_margin, qposadr, dofadr)
|
|
invweight = m.dof_invweight0[m.jnt_dofadr[ids]]
|
|
solref, solimp = m.jnt_solref[ids], m.jnt_solimp[ids]
|
|
frictionloss = jp.zeros_like(pos)
|
|
|
|
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
|
|
|
|
|
|
def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
|
|
"""Calculates constraint rows for slide and hinge joint limits."""
|
|
|
|
slide_hinge = np.isin(m.jnt_type, (JointType.SLIDE, JointType.HINGE))
|
|
ids = np.nonzero(slide_hinge & m.jnt_limited)[0]
|
|
|
|
if (m.opt.disableflags & DisableBit.LIMIT) or ids.size == 0:
|
|
return None
|
|
|
|
jnt_range = m.jnt_range[ids]
|
|
jnt_margin = m.jnt_margin[ids]
|
|
qposadr = m.jnt_qposadr[ids]
|
|
dofadr = m.jnt_dofadr[ids]
|
|
|
|
@jax.vmap
|
|
def fn(jnt_range, jnt_margin, qposadr, dofadr):
|
|
dist_min = d.qpos[qposadr] - jnt_range[0]
|
|
dist_max = jnt_range[1] - d.qpos[qposadr]
|
|
j = jp.zeros(m.nv).at[dofadr].set((dist_min < dist_max) * 2 - 1)
|
|
pos = jp.minimum(dist_min, dist_max) - jnt_margin
|
|
active = pos < 0
|
|
return j * active, pos * active
|
|
|
|
j, pos = fn(jnt_range, jnt_margin, qposadr, dofadr)
|
|
invweight = m.dof_invweight0[dofadr]
|
|
solref, solimp = m.jnt_solref[ids], m.jnt_solimp[ids]
|
|
frictionloss = jp.zeros_like(pos)
|
|
|
|
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
|
|
|
|
|
|
def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
|
|
"""Calculates constraint rows for contacts."""
|
|
|
|
if d.ncon == 0:
|
|
return None
|
|
|
|
def contact_efc(c: Contact, condim: int):
|
|
|
|
@jax.vmap
|
|
def fn(c: Contact):
|
|
dist = c.dist - c.includemargin
|
|
active = dist < 0
|
|
body1, body2 = jp.array(m.geom_bodyid)[c.geom]
|
|
jac1p, jac1r = support.jac(m, d, c.pos, body1)
|
|
jac2p, jac2r = support.jac(m, d, c.pos, body2)
|
|
diff = c.frame @ (jac2p - jac1p).T
|
|
if condim > 3: # only calculate rotational diff if needed
|
|
diff = jp.concatenate((diff, c.frame @ (jac2r - jac1r).T), axis=0)
|
|
tran = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0]
|
|
|
|
if condim == 1:
|
|
return diff[0] * active, tran, dist * active, c.solref, c.solimp
|
|
|
|
# a pair of opposing pyramid edges per friction dimension
|
|
# repeat friction directions with positive and negative sign
|
|
fri = jp.repeat(c.friction[: condim - 1], 2, axis=0).at[1::2].mul(-1)
|
|
# repeat condims of jacdiff to match +/- friction directions
|
|
j = diff[0] + jp.repeat(diff[1:condim], 2, axis=0) * fri[:, None]
|
|
# pyramidal has common invweight across all edges
|
|
diag_approx = tran + fri[0] * fri[0] * tran
|
|
inv_w = diag_approx * 2 * fri[0] * fri[0] / m.opt.impratio
|
|
repeat_fn = lambda x: jp.repeat(x[None], (condim - 1) * 2, axis=0)
|
|
inv_w, pos, solref, solimp = jax.tree_util.tree_map(
|
|
repeat_fn, (inv_w, dist, c.solref, c.solimp)
|
|
)
|
|
return j * active, inv_w, pos * active, solref, solimp
|
|
|
|
return fn(c)
|
|
|
|
# group efc calculations by condim
|
|
dims, begs = np.unique(d.contact.dim, return_index=True)
|
|
efcs = []
|
|
for i in range(len(dims)):
|
|
dim, beg = dims[i], begs[i]
|
|
end = begs[i + 1] if i < len(dims) - 1 else None
|
|
c = jax.tree_util.tree_map(lambda x, b=beg, e=end: x[b:e], d.contact)
|
|
efc = contact_efc(c, dim)
|
|
if dim > 1:
|
|
# remove efc grouping dimension
|
|
efc = jax.tree_util.tree_map(jp.concatenate, efc)
|
|
efcs.append(efc)
|
|
|
|
efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs)
|
|
j, invweight, pos, solref, solimp = efc
|
|
frictionloss = jp.zeros_like(pos)
|
|
|
|
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
|
|
|
|
|
|
def counts(efc_type: np.ndarray) -> Tuple[int, int, int, int]:
|
|
"""Returns equality, friction, limit, and contact constraint counts."""
|
|
ne = (efc_type == ConstraintType.EQUALITY).sum()
|
|
nf = 0 # no support for friction loss yet
|
|
nl = (efc_type == ConstraintType.LIMIT_JOINT).sum()
|
|
nc_f = (efc_type == ConstraintType.CONTACT_FRICTIONLESS).sum()
|
|
nc_p = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum()
|
|
nc = nc_f + nc_p
|
|
|
|
return ne, nf, nl, nc
|
|
|
|
|
|
def make_efc_type(
|
|
m: Union[Model, mujoco.MjModel], dim: Optional[np.ndarray] = None
|
|
) -> np.ndarray:
|
|
"""Returns efc_type that outlines the type of each constraint row."""
|
|
if m.opt.disableflags & DisableBit.CONSTRAINT:
|
|
return np.empty(0, dtype=int)
|
|
|
|
dim = collision_driver.make_condim(m) if dim is None else dim
|
|
efc_types = []
|
|
|
|
if not m.opt.disableflags & DisableBit.EQUALITY:
|
|
num_rows = (m.eq_type == EqType.CONNECT).sum() * 3
|
|
num_rows += (m.eq_type == EqType.WELD).sum() * 6
|
|
num_rows += (m.eq_type == EqType.JOINT).sum()
|
|
efc_types.extend([ConstraintType.EQUALITY] * num_rows)
|
|
|
|
if not m.opt.disableflags & DisableBit.LIMIT:
|
|
efc_types.extend([ConstraintType.LIMIT_JOINT] * m.jnt_limited.sum())
|
|
|
|
if not m.opt.disableflags & DisableBit.CONTACT:
|
|
num_rows = sum(_CONDIM_EFC_COUNT[d] for d in dim)
|
|
efc_types.extend([ConstraintType.CONTACT_PYRAMIDAL] * num_rows)
|
|
|
|
return np.array(efc_types)
|
|
|
|
|
|
def make_efc_address(efc_type: np.ndarray, dim: np.ndarray) -> np.ndarray:
|
|
"""Returns efc_address that maps contacts to constraint row address."""
|
|
nc = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum()
|
|
nc_start = efc_type.size - nc
|
|
offsets = np.cumsum([0] + [_CONDIM_EFC_COUNT[d] for d in dim])[:-1]
|
|
|
|
return nc_start + offsets
|
|
|
|
|
|
def make_constraint(m: Model, d: Data) -> Data:
|
|
"""Creates constraint jacobians and other supporting data."""
|
|
|
|
if m.opt.disableflags & DisableBit.CONSTRAINT:
|
|
efcs = ()
|
|
else:
|
|
efcs = tuple(efc for efc in (
|
|
_instantiate_equality_connect(m, d),
|
|
_instantiate_equality_weld(m, d),
|
|
_instantiate_equality_joint(m, d),
|
|
_instantiate_friction(m, d),
|
|
_instantiate_limit_ball(m, d),
|
|
_instantiate_limit_slide_hinge(m, d),
|
|
_instantiate_contact(m, d),
|
|
) if efc is not None)
|
|
|
|
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)
|
|
return d
|
|
|
|
efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs)
|
|
|
|
@jax.vmap
|
|
def fn(efc):
|
|
k, b, imp = _kbi(m, efc.solref, efc.solimp, efc.pos_norm)
|
|
r = jp.maximum(efc.invweight * (1 - imp) / imp, mujoco.mjMINVAL)
|
|
aref = -b * (efc.J @ d.qvel) - k * imp * efc.pos
|
|
return aref, r
|
|
|
|
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)
|
|
|
|
return d
|