def _warmstart(self):
"""Calculate qacc_warmstart. Initializing qacc_warmstart in the
simulator will result in a more rapid computation of qacc when
other fields (e.g. qpos, qvel, ctrl) are only slightly changed."""
# TODO: This is largely undocumented in MuJoCo; I'm following the lead
# of derivative.cpp. Does this have any implications for correctness,
# or just performance?
curr_state = np.concatenate((self.get_state(), self.sim.data.ctrl))
if np.array_equal(self.warmstart_for, curr_state):
# self.qacc_warmstart has already been computed for this x.
return
self.sim.forward() # run full computation
for i in range(self.NWARMUP):
# iterate several times to get a good qacc_warmstart
mjfunc.mj_forwardSkip(self.env.sim.model, self.env.sim.data,
SkipStages.VEL.value, 1)
self.warmstart[:] = self.sim.data.qacc_warmstart
self.warmstart_for = curr_state
def _finite_diff(sim, qacc_warmstart, vary, eps):
# Select the variable to perform finite differences over
variables = {
VaryValue.POS: sim.data.qpos,
VaryValue.VEL: sim.data.qvel,
VaryValue.CTRL: sim.data.ctrl,
}
variable = variables[vary]
# Select what parts of the forward dynamics computation can be skipped.
# Everything depends on position, so perform full computation in that case.
# When varying velocity, can skip position-dependent computations.
# When varying control, can skip acceleration/force-dependent computations.
skip_stages = {
VaryValue.POS: SkipStages.NONE,
VaryValue.VEL: SkipStages.POS,
VaryValue.CTRL: SkipStages.VEL,
}
skip_stage = skip_stages[vary]
# Make a copy of variable, to restore after each perturbation
original = np.copy(variable)
# Run forward simulation and copy the qacc output. We will take finite
# differences w.r.t. this value later.
sim.data.qacc_warmstart[:] = qacc_warmstart
sim.forward() # TODO: this computation is probably often redundant
center = np.copy(sim.data.qacc)
if vary in [VaryValue.VEL, VaryValue.CTRL]:
nin = len(variable)
nout = len(center)
jacobian = np.zeros((nout, nin), dtype=np.float64)
for j in range(nin):
# perturb
variable[j] += eps
# compute finite-difference
sim.data.qacc_warmstart[:] = qacc_warmstart
mjfunc.mj_forwardSkip(sim.model, sim.data, skip_stage.value, 1)
jacobian[:, j] = (sim.data.qacc - center) / eps
# unperturb
variable[j] = original[j]
elif vary == VaryValue.POS:
# When varying qpos, positions corresponding to ball and free joints
# are specified in quaternions. The derivative w.r.t. a quaternion
# is only defined in the tangent space to the configuration manifold.
# Accordingly, we do not vary the quaternion directly, but instead
# use mju_quatIntegrate to perturb with a certain angular velocity.
# Note that the Jacobian is always square with size sim.model.nv,
# which in the presence of quaternions is smaller than sim.model.nq.
nv = sim.model.nv
jacobian = np.zeros((nv, nv), dtype=np.float64)
# If qvel[j] belongs to a quaternion, then quatadr[j] gives the address
# of that quaternion and dofpos[j] gives the corresponding degree of
# freedom within the quaternion.
#
# If qvel[j] is not part of a quaternion, then quatadr[j] is -1
# and ids[j] gives the index in qpos to that joint. If there are no
# quaternions, ids will just be a sequence of consecutive numbers.
# But in the presence of a quaternion, it will skip a number.
# Common variables
joint_ids = sim.model.dof_jntid
jnt_types = sim.model.jnt_type[joint_ids]
jnt_qposadr = sim.model.jnt_qposadr[joint_ids]
dofadrs = sim.model.jnt_dofadr[joint_ids]
# Compute ids
js = np.arange(nv)
ids = jnt_qposadr + js - dofadrs
# Compute quatadr and dofpos
quatadr = np.tile(-1, nv)
dofpos = np.zeros(nv)
# A ball joint is always a quaternion.
ball_joints = (jnt_types == JointTypes.BALL.value)
quatadr[ball_joints] = jnt_qposadr[ball_joints]
dofpos[ball_joints] = js[ball_joints] - dofadrs[ball_joints]
# A free joint consists of a Cartesian (x,y,z) and a quaternion.
free_joints = (jnt_types == JointTypes.FREE.value) & (js > dofadrs + 3)
quatadr[free_joints] = jnt_qposadr[free_joints] + 3
dofpos[free_joints] = js[free_joints] - dofadrs[ball_joints] - 3
for j in js:
# perturb
qa = quatadr[j]
if qa >= 0: # quaternion perturbation
angvel = np.array([0, 0, 0])
angvel[dofpos[j]] = eps
# side-effect: perturbs variable[qa:qa+4]
mjfunc.mju_quatIntegrate(variable[qa], angvel, 1)
else: # simple perturbation
variable[ids[j]] += eps
# compute finite-difference
sim.data.qacc_warmstart[:] = qacc_warmstart
mjfunc.mj_forwardSkip(sim.model, sim.data, skip_stage.value, 1)
jacobian[:, j] = (sim.data.qacc - center) / eps
# unperturb
if qa >= 0: # quaternion perturbation
variable[qa:qa + 4] = original[qa:qa + 4]
else:
variable[ids[j]] = original[ids[j]]
else:
assert False
return jacobian
