Taylor Howell
d0faa3624e
PiperOrigin-RevId: 775658246 Change-Id: Ie6a71b971d4b165e0235689b7b4ee3246486e10d
111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
# Copyright 2025 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.
|
|
# ==============================================================================
|
|
"""Inverse dynamics functions."""
|
|
|
|
from jax import numpy as jp
|
|
from mujoco.mjx._src import derivative
|
|
from mujoco.mjx._src import forward
|
|
from mujoco.mjx._src import sensor
|
|
from mujoco.mjx._src import smooth
|
|
from mujoco.mjx._src import solver
|
|
from mujoco.mjx._src import support
|
|
# pylint: disable=g-importing-member
|
|
from mujoco.mjx._src.types import Data
|
|
from mujoco.mjx._src.types import DisableBit
|
|
from mujoco.mjx._src.types import EnableBit
|
|
from mujoco.mjx._src.types import IntegratorType
|
|
from mujoco.mjx._src.types import Model
|
|
|
|
|
|
def discrete_acc(m: Model, d: Data) -> Data:
|
|
"""Convert discrete-time qacc to continuous-time qacc."""
|
|
|
|
if m.opt.integrator == IntegratorType.RK4:
|
|
raise RuntimeError(
|
|
'discrete inverse dynamics is not supported by RK4 integrator'
|
|
)
|
|
elif m.opt.integrator == IntegratorType.EULER:
|
|
dsbl_eulerdamp = m.opt.disableflags & DisableBit.EULERDAMP
|
|
no_dof_damping = (m.dof_damping == 0).all()
|
|
if dsbl_eulerdamp or no_dof_damping:
|
|
return d
|
|
|
|
# set qfrc = (M + h*diag(B)) * qacc
|
|
qfrc = support.mul_m(m, d, d.qacc)
|
|
qfrc += m.opt.timestep * m.dof_damping * d.qacc
|
|
elif m.opt.integrator == IntegratorType.IMPLICITFAST:
|
|
qm = support.full_m(m, d)
|
|
|
|
# compute analytical derivative qDeriv; skip rne derivative
|
|
qderiv = derivative.deriv_smooth_vel(m, d)
|
|
if qderiv is not None:
|
|
# M = M - dt*qDeriv
|
|
qm -= m.opt.timestep * qderiv
|
|
|
|
# set qfrc = (M - dt*qDeriv) * qacc
|
|
qfrc = qm @ d.qacc
|
|
else:
|
|
raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.')
|
|
|
|
# solve for qacc: qfrc = M * qacc
|
|
qacc = smooth.solve_m(m, d, qfrc)
|
|
|
|
return d.replace(qacc=qacc)
|
|
|
|
|
|
def inv_constraint(m: Model, d: Data) -> Data:
|
|
"""Inverse constraint solver."""
|
|
|
|
# no constraints
|
|
if d._impl.efc_J.size == 0: # pytype: disable=attribute-error
|
|
return d.replace(qfrc_constraint=jp.zeros(m.nv))
|
|
|
|
# update
|
|
ctx = solver.Context.create(m, d, grad=False)
|
|
|
|
return d.tree_replace({
|
|
'qfrc_constraint': ctx.qfrc_constraint,
|
|
'_impl.efc_force': ctx.efc_force,
|
|
})
|
|
|
|
|
|
def inverse(m: Model, d: Data) -> Data:
|
|
"""Inverse dynamics."""
|
|
d = forward.fwd_position(m, d)
|
|
d = sensor.sensor_pos(m, d)
|
|
d = forward.fwd_velocity(m, d)
|
|
d = sensor.sensor_vel(m, d)
|
|
|
|
qacc = d.qacc
|
|
if m.opt.enableflags & EnableBit.INVDISCRETE:
|
|
d = discrete_acc(m, d)
|
|
|
|
d = inv_constraint(m, d)
|
|
d = smooth.rne(m, d)
|
|
d = smooth.tendon_bias(m, d)
|
|
d = sensor.sensor_acc(m, d)
|
|
|
|
qfrc_inverse = (
|
|
d.qfrc_bias
|
|
+ support.mul_m(m, d, d.qacc)
|
|
- d.qfrc_passive
|
|
- d.qfrc_constraint
|
|
)
|
|
|
|
if m.opt.enableflags & EnableBit.INVDISCRETE:
|
|
return d.replace(qfrc_inverse=qfrc_inverse, qacc=qacc)
|
|
else:
|
|
return d.replace(qfrc_inverse=qfrc_inverse)
|