# Copyright 2025 The Newton Developers
#
# 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.
# ==============================================================================
from typing import Any, Tuple
import warp as wp
from mujoco_warp._src import math
from mujoco_warp._src import ray
from mujoco_warp._src import smooth
from mujoco_warp._src import support
from mujoco_warp._src.collision_sdf import get_sdf_params
from mujoco_warp._src.collision_sdf import sdf
from mujoco_warp._src.types import MJ_MAXVAL
from mujoco_warp._src.types import MJ_MINVAL
from mujoco_warp._src.types import ConeType
from mujoco_warp._src.types import ConstraintType
from mujoco_warp._src.types import ContactType
from mujoco_warp._src.types import Data
from mujoco_warp._src.types import DataType
from mujoco_warp._src.types import DisableBit
from mujoco_warp._src.types import JointType
from mujoco_warp._src.types import Model
from mujoco_warp._src.types import ObjType
from mujoco_warp._src.types import SensorType
from mujoco_warp._src.types import TrnType
from mujoco_warp._src.types import vec5
from mujoco_warp._src.types import vec6
from mujoco_warp._src.types import vec8
from mujoco_warp._src.types import vec8i
from mujoco_warp._src.util_misc import inside_geom
from mujoco_warp._src.warp_util import cache_kernel
from mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
@wp.func
def _write_scalar(
# Model:
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
# In:
sensorid: int,
sensor: Any,
# Out:
out: wp.array(dtype=float),
):
adr = sensor_adr[sensorid]
cutoff = sensor_cutoff[sensorid]
if cutoff > 0.0 and not (sensor_type[sensorid] == int(SensorType.GEOMFROMTO.value)):
datatype = sensor_datatype[sensorid]
if datatype == DataType.REAL:
out[adr] = wp.clamp(sensor, -cutoff, cutoff)
return
elif datatype == DataType.POSITIVE:
out[adr] = wp.min(sensor, cutoff)
return
out[adr] = sensor
@wp.func
def _write_vector(
# Model:
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
# In:
sensorid: int,
sensordim: int,
sensor: Any,
# Out:
out: wp.array(dtype=float),
):
adr = sensor_adr[sensorid]
cutoff = sensor_cutoff[sensorid]
if cutoff > 0.0 and not (sensor_type[sensorid] == int(SensorType.GEOMFROMTO.value)):
datatype = sensor_datatype[sensorid]
if datatype == DataType.REAL:
for i in range(sensordim):
out[adr + i] = wp.clamp(sensor[i], -cutoff, cutoff)
return
elif datatype == DataType.POSITIVE:
for i in range(sensordim):
out[adr + i] = wp.min(sensor[i], cutoff)
return
for i in range(sensordim):
out[adr + i] = sensor[i]
@wp.func
def _magnetometer(
# Model:
opt_magnetic: wp.array(dtype=wp.vec3),
# Data in:
site_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
worldid: int,
objid: int,
) -> wp.vec3:
magnetic = opt_magnetic[worldid % opt_magnetic.shape[0]]
return wp.transpose(site_xmat_in[worldid, objid]) @ magnetic
@wp.func
def _cam_projection(
# Model:
cam_fovy: wp.array2d(dtype=float),
cam_resolution: wp.array(dtype=wp.vec2i),
cam_sensorsize: wp.array(dtype=wp.vec2),
cam_intrinsic: wp.array2d(dtype=wp.vec4),
# Data in:
site_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
worldid: int,
objid: int,
refid: int,
) -> wp.vec2:
sensorsize = cam_sensorsize[refid]
intrinsic = cam_intrinsic[worldid % cam_intrinsic.shape[0], refid]
fovy = cam_fovy[worldid % cam_fovy.shape[0], refid]
res = cam_resolution[refid]
target_xpos = site_xpos_in[worldid, objid]
xpos = cam_xpos_in[worldid, refid]
xmat = cam_xmat_in[worldid, refid]
translation = wp.mat44f(1.0, 0.0, 0.0, -xpos[0], 0.0, 1.0, 0.0, -xpos[1], 0.0, 0.0, 1.0, -xpos[2], 0.0, 0.0, 0.0, 1.0)
rotation = wp.mat44f(
xmat[0, 0], xmat[1, 0], xmat[2, 0], 0.0,
xmat[0, 1], xmat[1, 1], xmat[2, 1], 0.0,
xmat[0, 2], xmat[1, 2], xmat[2, 2], 0.0,
0.0, 0.0, 0.0, 1.0,
) # fmt: skip
# focal transformation matrix (3 x 4)
if sensorsize[0] != 0.0 and sensorsize[1] != 0.0:
fx = intrinsic[0] / (sensorsize[0] + MJ_MINVAL) * float(res[0])
fy = intrinsic[1] / (sensorsize[1] + MJ_MINVAL) * float(res[1])
focal = wp.mat44f(-fx, 0.0, 0.0, 0.0, 0.0, fy, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0)
else:
f = 0.5 / wp.tan(fovy * wp.static(wp.pi / 360.0)) * float(res[1])
focal = wp.mat44f(-f, 0.0, 0.0, 0.0, 0.0, f, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0)
# image matrix (3 x 3)
image = wp.mat44f(
1.0, 0.0, 0.5 * float(res[0]), 0.0, 0.0, 1.0, 0.5 * float(res[1]), 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0
)
# projection matrix (3 x 4): product of all 4 matrices
# TODO(team): compute proj directly
proj = image @ focal @ rotation @ translation
# projection matrix multiples homogeneous [x, y, z, 1] vectors
pos_hom = wp.vec4(target_xpos[0], target_xpos[1], target_xpos[2], 1.0)
# project world coordinates into pixel space, see:
# https://en.wikipedia.org/wiki/3D_projection#Mathematical_formula
pixel_coord_hom = proj @ pos_hom
# avoid dividing by tiny numbers
denom = pixel_coord_hom[2]
if wp.abs(denom) < MJ_MINVAL:
denom = wp.clamp(denom, -MJ_MINVAL, MJ_MINVAL)
# compute projection
return wp.vec2f(pixel_coord_hom[0], pixel_coord_hom[1]) / denom
@wp.kernel
def _sensor_rangefinder_init(
# Model:
sensor_objid: wp.array(dtype=int),
sensor_rangefinder_adr: wp.array(dtype=int),
# Data in:
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
# Out:
pnt_out: wp.array2d(dtype=wp.vec3),
vec_out: wp.array2d(dtype=wp.vec3),
):
worldid, rfid = wp.tid()
sensorid = sensor_rangefinder_adr[rfid]
objid = sensor_objid[sensorid]
site_xpos = site_xpos_in[worldid, objid]
site_xmat = site_xmat_in[worldid, objid]
pnt_out[worldid, rfid] = site_xpos
vec_out[worldid, rfid] = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2])
@wp.func
def _joint_pos(jnt_qposadr: wp.array(dtype=int), qpos_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float:
return qpos_in[worldid, jnt_qposadr[objid]]
@wp.func
def _tendon_pos(ten_length_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float:
return ten_length_in[worldid, objid]
@wp.func
def _actuator_pos(actuator_length_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float:
return actuator_length_in[worldid, objid]
@wp.func
def _ball_quat(jnt_qposadr: wp.array(dtype=int), qpos_in: wp.array2d(dtype=float), worldid: int, objid: int) -> wp.quat:
adr = jnt_qposadr[objid]
quat = wp.quat(
qpos_in[worldid, adr + 0],
qpos_in[worldid, adr + 1],
qpos_in[worldid, adr + 2],
qpos_in[worldid, adr + 3],
)
return wp.normalize(quat)
@wp.kernel
def _limit_pos(
# Model:
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_objid: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
sensor_limitpos_adr: wp.array(dtype=int),
# Data in:
ne_in: wp.array(dtype=int),
nf_in: wp.array(dtype=int),
nl_in: wp.array(dtype=int),
efc_type_in: wp.array2d(dtype=int),
efc_id_in: wp.array2d(dtype=int),
efc_pos_in: wp.array2d(dtype=float),
efc_margin_in: wp.array2d(dtype=float),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, efcid, limitposid = wp.tid()
ne = ne_in[worldid]
nf = nf_in[worldid]
nl = nl_in[worldid]
# skip if not limit
if efcid < ne + nf or efcid >= ne + nf + nl:
return
sensorid = sensor_limitpos_adr[limitposid]
if efc_id_in[worldid, efcid] == sensor_objid[sensorid]:
efc_type = efc_type_in[worldid, efcid]
if efc_type == ConstraintType.LIMIT_JOINT or efc_type == ConstraintType.LIMIT_TENDON:
val = efc_pos_in[worldid, efcid] - efc_margin_in[worldid, efcid]
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, sensordata_out[worldid])
@wp.func
def _frame_pos(
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xmat_in: wp.array2d(dtype=wp.mat33),
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
worldid: int,
objid: int,
objtype: int,
refid: int,
reftype: int,
) -> wp.vec3:
if objtype == ObjType.BODY:
xpos = xipos_in[worldid, objid]
elif objtype == ObjType.XBODY:
xpos = xpos_in[worldid, objid]
elif objtype == ObjType.GEOM:
xpos = geom_xpos_in[worldid, objid]
elif objtype == ObjType.SITE:
xpos = site_xpos_in[worldid, objid]
elif objtype == ObjType.CAMERA:
xpos = cam_xpos_in[worldid, objid]
else: # UNKNOWN
xpos = wp.vec3(0.0)
if refid == -1:
return xpos
if reftype == ObjType.BODY:
xpos_ref = xipos_in[worldid, refid]
xmat_ref = ximat_in[worldid, refid]
elif objtype == ObjType.XBODY:
xpos_ref = xpos_in[worldid, refid]
xmat_ref = xmat_in[worldid, refid]
elif reftype == ObjType.GEOM:
xpos_ref = geom_xpos_in[worldid, refid]
xmat_ref = geom_xmat_in[worldid, refid]
elif reftype == ObjType.SITE:
xpos_ref = site_xpos_in[worldid, refid]
xmat_ref = site_xmat_in[worldid, refid]
elif reftype == ObjType.CAMERA:
xpos_ref = cam_xpos_in[worldid, refid]
xmat_ref = cam_xmat_in[worldid, refid]
else: # UNKNOWN
xpos_ref = wp.vec3(0.0)
xmat_ref = wp.identity(3, wp.float32)
return wp.transpose(xmat_ref) @ (xpos - xpos_ref)
@wp.func
def _frame_axis(
# Data in:
xmat_in: wp.array2d(dtype=wp.mat33),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
site_xmat_in: wp.array2d(dtype=wp.mat33),
cam_xmat_in: wp.array2d(dtype=wp.mat33),
# In:
worldid: int,
objid: int,
objtype: int,
refid: int,
reftype: int,
frame_axis: int,
) -> wp.vec3:
if objtype == ObjType.BODY:
xmat = ximat_in[worldid, objid]
axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
elif objtype == ObjType.XBODY:
xmat = xmat_in[worldid, objid]
axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
elif objtype == ObjType.GEOM:
xmat = geom_xmat_in[worldid, objid]
axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
elif objtype == ObjType.SITE:
xmat = site_xmat_in[worldid, objid]
axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
elif objtype == ObjType.CAMERA:
xmat = cam_xmat_in[worldid, objid]
axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
else: # UNKNOWN
axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis])
if refid == -1:
return axis
if reftype == ObjType.BODY:
xmat_ref = ximat_in[worldid, refid]
elif reftype == ObjType.XBODY:
xmat_ref = xmat_in[worldid, refid]
elif reftype == ObjType.GEOM:
xmat_ref = geom_xmat_in[worldid, refid]
elif reftype == ObjType.SITE:
xmat_ref = site_xmat_in[worldid, refid]
elif reftype == ObjType.CAMERA:
xmat_ref = cam_xmat_in[worldid, refid]
else: # UNKNOWN
xmat_ref = wp.identity(3, dtype=wp.float32)
return wp.transpose(xmat_ref) @ axis
@wp.func
def _frame_quat(
# Model:
body_iquat: wp.array2d(dtype=wp.quat),
geom_bodyid: wp.array(dtype=int),
geom_quat: wp.array2d(dtype=wp.quat),
site_bodyid: wp.array(dtype=int),
site_quat: wp.array2d(dtype=wp.quat),
cam_bodyid: wp.array(dtype=int),
cam_quat: wp.array2d(dtype=wp.quat),
# Data in:
xquat_in: wp.array2d(dtype=wp.quat),
# In:
worldid: int,
objid: int,
objtype: int,
refid: int,
reftype: int,
) -> wp.quat:
body_iquat_id = worldid % body_iquat.shape[0]
geom_quat_id = worldid % geom_quat.shape[0]
site_quat_id = worldid % site_quat.shape[0]
cam_quat_id = worldid % cam_quat.shape[0]
if objtype == ObjType.BODY:
quat = math.mul_quat(xquat_in[worldid, objid], body_iquat[body_iquat_id, objid])
elif objtype == ObjType.XBODY:
quat = xquat_in[worldid, objid]
elif objtype == ObjType.GEOM:
quat = math.mul_quat(xquat_in[worldid, geom_bodyid[objid]], geom_quat[geom_quat_id, objid])
elif objtype == ObjType.SITE:
quat = math.mul_quat(xquat_in[worldid, site_bodyid[objid]], site_quat[site_quat_id, objid])
elif objtype == ObjType.CAMERA:
quat = math.mul_quat(xquat_in[worldid, cam_bodyid[objid]], cam_quat[cam_quat_id, objid])
else: # UNKNOWN
quat = wp.quat(1.0, 0.0, 0.0, 0.0)
if refid == -1:
return quat
if reftype == ObjType.BODY:
refquat = math.mul_quat(xquat_in[worldid, refid], body_iquat[body_iquat_id, refid])
elif reftype == ObjType.XBODY:
refquat = xquat_in[worldid, refid]
elif reftype == ObjType.GEOM:
refquat = math.mul_quat(xquat_in[worldid, geom_bodyid[refid]], geom_quat[geom_quat_id, refid])
elif reftype == ObjType.SITE:
refquat = math.mul_quat(xquat_in[worldid, site_bodyid[refid]], site_quat[site_quat_id, refid])
elif reftype == ObjType.CAMERA:
refquat = math.mul_quat(xquat_in[worldid, cam_bodyid[refid]], cam_quat[cam_quat_id, refid])
else: # UNKNOWN
refquat = wp.quat(1.0, 0.0, 0.0, 0.0)
return math.mul_quat(math.quat_inv(refquat), quat)
@wp.func
def _subtree_com(subtree_com_in: wp.array2d(dtype=wp.vec3), worldid: int, objid: int) -> wp.vec3:
return subtree_com_in[worldid, objid]
@wp.func
def _clock(time_in: wp.array(dtype=float), worldid: int) -> float:
return time_in[worldid]
@wp.kernel
def _sensor_pos(
# Model:
ngeom: int,
opt_magnetic: wp.array(dtype=wp.vec3),
body_geomnum: wp.array(dtype=int),
body_geomadr: wp.array(dtype=int),
body_iquat: wp.array2d(dtype=wp.quat),
jnt_qposadr: wp.array(dtype=int),
geom_type: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
geom_quat: wp.array2d(dtype=wp.quat),
site_type: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
site_size: wp.array(dtype=wp.vec3),
site_quat: wp.array2d(dtype=wp.quat),
cam_bodyid: wp.array(dtype=int),
cam_quat: wp.array2d(dtype=wp.quat),
cam_fovy: wp.array2d(dtype=float),
cam_resolution: wp.array(dtype=wp.vec2i),
cam_sensorsize: wp.array(dtype=wp.vec2),
cam_intrinsic: wp.array2d(dtype=wp.vec4),
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_objtype: wp.array(dtype=int),
sensor_objid: wp.array(dtype=int),
sensor_reftype: wp.array(dtype=int),
sensor_refid: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
nxn_pairid: wp.array(dtype=wp.vec2i),
sensor_pos_adr: wp.array(dtype=int),
rangefinder_sensor_adr: wp.array(dtype=int),
# Data in:
time_in: wp.array(dtype=float),
energy_in: wp.array(dtype=wp.vec2),
qpos_in: wp.array2d(dtype=float),
xpos_in: wp.array2d(dtype=wp.vec3),
xquat_in: wp.array2d(dtype=wp.quat),
xmat_in: wp.array2d(dtype=wp.mat33),
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
ten_length_in: wp.array2d(dtype=float),
actuator_length_in: wp.array2d(dtype=float),
# In:
rangefinder_dist_in: wp.array2d(dtype=float),
sensor_collision_in: wp.array4d(dtype=float),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, posid = wp.tid()
sensorid = sensor_pos_adr[posid]
sensortype = sensor_type[sensorid]
objid = sensor_objid[sensorid]
out = sensordata_out[worldid]
if sensortype == SensorType.MAGNETOMETER:
vec3 = _magnetometer(opt_magnetic, site_xmat_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.CAMPROJECTION:
refid = sensor_refid[sensorid]
vec2 = _cam_projection(
cam_fovy, cam_resolution, cam_sensorsize, cam_intrinsic, site_xpos_in, cam_xpos_in, cam_xmat_in, worldid, objid, refid
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 2, vec2, out)
elif sensortype == SensorType.RANGEFINDER:
val = rangefinder_dist_in[worldid, rangefinder_sensor_adr[sensorid]]
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.JOINTPOS:
val = _joint_pos(jnt_qposadr, qpos_in, worldid, objid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.TENDONPOS:
val = _tendon_pos(ten_length_in, worldid, objid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.ACTUATORPOS:
val = _actuator_pos(actuator_length_in, worldid, objid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.BALLQUAT:
quat = _ball_quat(jnt_qposadr, qpos_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, out)
elif sensortype == SensorType.FRAMEPOS:
objtype = sensor_objtype[sensorid]
refid = sensor_refid[sensorid]
reftype = sensor_reftype[sensorid]
vec3 = _frame_pos(
xpos_in,
xmat_in,
xipos_in,
ximat_in,
geom_xpos_in,
geom_xmat_in,
site_xpos_in,
site_xmat_in,
cam_xpos_in,
cam_xmat_in,
worldid,
objid,
objtype,
refid,
reftype,
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.FRAMEXAXIS or sensortype == SensorType.FRAMEYAXIS or sensortype == SensorType.FRAMEZAXIS:
objtype = sensor_objtype[sensorid]
refid = sensor_refid[sensorid]
reftype = sensor_reftype[sensorid]
if sensortype == SensorType.FRAMEXAXIS:
axis = 0
elif sensortype == SensorType.FRAMEYAXIS:
axis = 1
elif sensortype == SensorType.FRAMEZAXIS:
axis = 2
vec3 = _frame_axis(
xmat_in, ximat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.FRAMEQUAT:
objtype = sensor_objtype[sensorid]
refid = sensor_refid[sensorid]
reftype = sensor_reftype[sensorid]
quat = _frame_quat(
body_iquat,
geom_bodyid,
geom_quat,
site_bodyid,
site_quat,
cam_bodyid,
cam_quat,
xquat_in,
worldid,
objid,
objtype,
refid,
reftype,
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, out)
elif sensortype == SensorType.SUBTREECOM:
vec3 = _subtree_com(subtree_com_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.GEOMDIST or sensortype == SensorType.GEOMNORMAL or sensortype == SensorType.GEOMFROMTO:
objtype = sensor_objtype[sensorid]
objid = sensor_objid[sensorid]
reftype = sensor_reftype[sensorid]
refid = sensor_refid[sensorid]
# initialize
dist = float(sensor_cutoff[sensorid])
pnts = vec6(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
flip = bool(False)
# check for flip direction
if objtype == int(ObjType.BODY.value):
n1 = body_geomnum[objid]
id1 = body_geomadr[objid]
else:
n1 = 1
id1 = objid
if reftype == int(ObjType.BODY.value):
n2 = body_geomnum[refid]
id2 = body_geomadr[refid]
else:
n2 = 1
id2 = refid
for geom1 in range(n1):
geomid1 = id1 + geom1
for geom2 in range(n2):
geomid2 = id2 + geom2
if geomid1 <= geomid2:
pairid = math.upper_tri_index(ngeom, geomid1, geomid2)
else:
pairid = math.upper_tri_index(ngeom, geomid2, geomid1)
collisionid = nxn_pairid[pairid][1]
for i in range(8):
dist_new = sensor_collision_in[worldid, collisionid, i, 0]
if dist_new < dist:
dist = dist_new
if sensortype == SensorType.GEOMNORMAL or sensortype == SensorType.GEOMFROMTO:
pnts = vec6(
sensor_collision_in[worldid, collisionid, i, 1],
sensor_collision_in[worldid, collisionid, i, 2],
sensor_collision_in[worldid, collisionid, i, 3],
sensor_collision_in[worldid, collisionid, i, 4],
sensor_collision_in[worldid, collisionid, i, 5],
sensor_collision_in[worldid, collisionid, i, 6],
)
if geom_type[geomid1] > geom_type[geomid2]:
flip = True
elif geom_type[geomid1] == geom_type[geomid2]:
flip = geomid1 > geomid2
else:
flip = False
if sensortype == int(SensorType.GEOMDIST.value):
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, dist, out)
elif sensortype == int(SensorType.GEOMNORMAL.value):
if dist <= sensor_cutoff[sensorid]:
normal = wp.normalize(wp.vec3(pnts[3] - pnts[0], pnts[4] - pnts[1], pnts[5] - pnts[2]))
if flip:
normal *= -1.0
else:
normal = wp.vec3(0.0, 0.0, 0.0)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, normal, out)
elif sensortype == int(SensorType.GEOMFROMTO.value):
if dist <= sensor_cutoff[sensorid]:
if flip:
fromto = vec6(pnts[3], pnts[4], pnts[5], pnts[0], pnts[1], pnts[2])
else:
fromto = pnts
else:
fromto = vec6(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 6, fromto, out)
elif sensortype == SensorType.INSIDESITE:
objtype = sensor_objtype[sensorid]
if objtype == ObjType.XBODY:
xpos = xpos_in[worldid, objid]
elif objtype == ObjType.BODY:
xpos = xipos_in[worldid, objid]
elif objtype == ObjType.GEOM:
xpos = geom_xpos_in[worldid, objid]
elif objtype == ObjType.SITE:
xpos = site_xpos_in[worldid, objid]
elif objtype == ObjType.CAMERA:
xpos = cam_xpos_in[worldid, objid]
else:
return # should not occur
refid = sensor_refid[sensorid]
val_bool = inside_geom(site_xpos_in[worldid, refid], site_xmat_in[worldid, refid], site_size[refid], site_type[refid], xpos)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, float(val_bool), out)
elif sensortype == SensorType.E_POTENTIAL:
val = energy_in[worldid][0]
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.E_KINETIC:
val = energy_in[worldid][1]
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.CLOCK:
val = _clock(time_in, worldid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
@wp.kernel
def _sensor_collision(
# Model:
ngeom: int,
nxn_pairid: wp.array(dtype=wp.vec2i),
# Data in:
contact_dist_in: wp.array(dtype=float),
contact_pos_in: wp.array(dtype=wp.vec3),
contact_frame_in: wp.array(dtype=wp.mat33),
contact_geom_in: wp.array(dtype=wp.vec2i),
contact_worldid_in: wp.array(dtype=int),
contact_type_in: wp.array(dtype=int),
contact_geomcollisionid_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
# Out:
sensor_collision_out: wp.array4d(dtype=float),
):
conid = wp.tid()
if conid >= nacon_in[0]:
return
if not contact_type_in[conid] & ContactType.SENSOR:
return
geom = contact_geom_in[conid]
if geom[0] <= geom[1]:
pairid = math.upper_tri_index(ngeom, geom[0], geom[1])
else:
pairid = math.upper_tri_index(ngeom, geom[1], geom[0])
worldid = contact_worldid_in[conid]
collisionid = nxn_pairid[pairid][1]
geomcollisionid = contact_geomcollisionid_in[conid]
dist = contact_dist_in[conid]
pos = contact_pos_in[conid]
frame = contact_frame_in[conid]
normal = wp.vec3(frame[0, 0], frame[0, 1], frame[0, 2])
pnt1 = pos - 0.5 * dist * normal
pnt2 = pos + 0.5 * dist * normal
sensor_collision_out[worldid, collisionid, geomcollisionid, 0] = dist
sensor_collision_out[worldid, collisionid, geomcollisionid, 1] = pnt1[0]
sensor_collision_out[worldid, collisionid, geomcollisionid, 2] = pnt1[1]
sensor_collision_out[worldid, collisionid, geomcollisionid, 3] = pnt1[2]
sensor_collision_out[worldid, collisionid, geomcollisionid, 4] = pnt2[0]
sensor_collision_out[worldid, collisionid, geomcollisionid, 5] = pnt2[1]
sensor_collision_out[worldid, collisionid, geomcollisionid, 6] = pnt2[2]
[docs]
@event_scope
def sensor_pos(m: Model, d: Data):
"""Compute position-dependent sensor values."""
if m.opt.disableflags & DisableBit.SENSOR:
return
# rangefinder
rangefinder_dist = wp.empty((d.nworld, m.nrangefinder), dtype=float)
if m.sensor_rangefinder_adr.size > 0:
rangefinder_pnt = wp.empty((d.nworld, m.nrangefinder), dtype=wp.vec3)
rangefinder_vec = wp.empty((d.nworld, m.nrangefinder), dtype=wp.vec3)
rangefinder_geomid = wp.empty((d.nworld, m.nrangefinder), dtype=int)
rangefinder_normal = wp.empty((d.nworld, m.nrangefinder), dtype=wp.vec3)
# get position and direction
wp.launch(
_sensor_rangefinder_init,
dim=(d.nworld, m.sensor_rangefinder_adr.size),
inputs=[m.sensor_objid, m.sensor_rangefinder_adr, d.site_xpos, d.site_xmat],
outputs=[rangefinder_pnt, rangefinder_vec],
)
# get distances
ray.rays(
m,
d,
rangefinder_pnt,
rangefinder_vec,
vec6(MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL, MJ_MAXVAL),
True,
m.sensor_rangefinder_bodyid,
rangefinder_dist,
rangefinder_geomid,
rangefinder_normal,
)
if m.sensor_e_potential:
energy_pos(m, d)
if m.sensor_e_kinetic:
energy_vel(m, d)
# collision sensors (distance, normal, fromto)
sensor_collision = wp.full((d.nworld, m.nsensorcollision, 8, 7), 1.0e32, dtype=float)
if m.nsensorcollision:
wp.launch(
_sensor_collision,
dim=d.naconmax,
inputs=[
m.ngeom,
m.nxn_pairid,
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.geom,
d.contact.worldid,
d.contact.type,
d.contact.geomcollisionid,
d.nacon,
],
outputs=[sensor_collision],
)
wp.launch(
_sensor_pos,
dim=(d.nworld, m.sensor_pos_adr.size),
inputs=[
m.ngeom,
m.opt.magnetic,
m.body_geomnum,
m.body_geomadr,
m.body_iquat,
m.jnt_qposadr,
m.geom_type,
m.geom_bodyid,
m.geom_quat,
m.site_type,
m.site_bodyid,
m.site_size,
m.site_quat,
m.cam_bodyid,
m.cam_quat,
m.cam_fovy,
m.cam_resolution,
m.cam_sensorsize,
m.cam_intrinsic,
m.sensor_type,
m.sensor_datatype,
m.sensor_objtype,
m.sensor_objid,
m.sensor_reftype,
m.sensor_refid,
m.sensor_adr,
m.sensor_cutoff,
m.nxn_pairid,
m.sensor_pos_adr,
m.rangefinder_sensor_adr,
d.time,
d.energy,
d.qpos,
d.xpos,
d.xquat,
d.xmat,
d.xipos,
d.ximat,
d.geom_xpos,
d.geom_xmat,
d.site_xpos,
d.site_xmat,
d.cam_xpos,
d.cam_xmat,
d.subtree_com,
d.ten_length,
d.actuator_length,
rangefinder_dist,
sensor_collision,
],
outputs=[d.sensordata],
)
# jointlimitpos and tendonlimitpos
wp.launch(
_limit_pos,
dim=(d.nworld, d.njmax, m.sensor_limitpos_adr.size),
inputs=[
m.sensor_type,
m.sensor_datatype,
m.sensor_objid,
m.sensor_adr,
m.sensor_cutoff,
m.sensor_limitpos_adr,
d.ne,
d.nf,
d.nl,
d.efc.type,
d.efc.id,
d.efc.pos,
d.efc.margin,
],
outputs=[
d.sensordata,
],
)
@wp.func
def _velocimeter(
# Model:
body_rootid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
# Data in:
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
) -> wp.vec3:
bodyid = site_bodyid[objid]
pos = site_xpos_in[worldid, objid]
rot = site_xmat_in[worldid, objid]
cvel = cvel_in[worldid, bodyid]
ang = wp.spatial_top(cvel)
lin = wp.spatial_bottom(cvel)
subtree_com = subtree_com_in[worldid, body_rootid[bodyid]]
dif = pos - subtree_com
return wp.transpose(rot) @ (lin - wp.cross(dif, ang))
@wp.func
def _gyro(
# Model:
site_bodyid: wp.array(dtype=int),
# Data in:
site_xmat_in: wp.array2d(dtype=wp.mat33),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
) -> wp.vec3:
bodyid = site_bodyid[objid]
rot = site_xmat_in[worldid, objid]
cvel = cvel_in[worldid, bodyid]
ang = wp.spatial_top(cvel)
return wp.transpose(rot) @ ang
@wp.func
def _joint_vel(jnt_dofadr: wp.array(dtype=int), qvel_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float:
return qvel_in[worldid, jnt_dofadr[objid]]
@wp.func
def _tendon_vel(ten_velocity_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float:
return ten_velocity_in[worldid, objid]
@wp.func
def _actuator_vel(actuator_velocity_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float:
return actuator_velocity_in[worldid, objid]
@wp.func
def _ball_ang_vel(jnt_dofadr: wp.array(dtype=int), qvel_in: wp.array2d(dtype=float), worldid: int, objid: int) -> wp.vec3:
adr = jnt_dofadr[objid]
return wp.vec3(qvel_in[worldid, adr + 0], qvel_in[worldid, adr + 1], qvel_in[worldid, adr + 2])
@wp.kernel
def _limit_vel(
# Model:
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_objid: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
sensor_limitvel_adr: wp.array(dtype=int),
# Data in:
ne_in: wp.array(dtype=int),
nf_in: wp.array(dtype=int),
nl_in: wp.array(dtype=int),
efc_type_in: wp.array2d(dtype=int),
efc_id_in: wp.array2d(dtype=int),
efc_vel_in: wp.array2d(dtype=float),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, efcid, limitvelid = wp.tid()
ne = ne_in[worldid]
nf = nf_in[worldid]
nl = nl_in[worldid]
# skip if not limit
if efcid < ne + nf or efcid >= ne + nf + nl:
return
sensorid = sensor_limitvel_adr[limitvelid]
if efc_id_in[worldid, efcid] == sensor_objid[sensorid]:
efc_type = efc_type_in[worldid, efcid]
if efc_type == ConstraintType.LIMIT_JOINT or efc_type == ConstraintType.LIMIT_TENDON:
_write_scalar(
sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, efc_vel_in[worldid, efcid], sensordata_out[worldid]
)
@wp.func
def _cvel_offset(
# Model:
body_rootid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
cam_bodyid: wp.array(dtype=int),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xipos_in: wp.array2d(dtype=wp.vec3),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
site_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objtype: int,
objid: int,
) -> Tuple[wp.spatial_vector, wp.vec3]:
if objtype == ObjType.BODY:
pos = xipos_in[worldid, objid]
bodyid = objid
elif objtype == ObjType.XBODY:
pos = xpos_in[worldid, objid]
bodyid = objid
elif objtype == ObjType.GEOM:
pos = geom_xpos_in[worldid, objid]
bodyid = geom_bodyid[objid]
elif objtype == ObjType.SITE:
pos = site_xpos_in[worldid, objid]
bodyid = site_bodyid[objid]
elif objtype == ObjType.CAMERA:
pos = cam_xpos_in[worldid, objid]
bodyid = cam_bodyid[objid]
else: # UNKNOWN
pos = wp.vec3(0.0)
bodyid = 0
return cvel_in[worldid, bodyid], pos - subtree_com_in[worldid, body_rootid[bodyid]]
@wp.func
def _frame_linvel(
# Model:
body_rootid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
cam_bodyid: wp.array(dtype=int),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xmat_in: wp.array2d(dtype=wp.mat33),
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
objtype: int,
refid: int,
reftype: int,
) -> wp.vec3:
if objtype == ObjType.BODY:
xpos = xipos_in[worldid, objid]
elif objtype == ObjType.XBODY:
xpos = xpos_in[worldid, objid]
elif objtype == ObjType.GEOM:
xpos = geom_xpos_in[worldid, objid]
elif objtype == ObjType.SITE:
xpos = site_xpos_in[worldid, objid]
elif objtype == ObjType.CAMERA:
xpos = cam_xpos_in[worldid, objid]
else: # UNKNOWN
xpos = wp.vec3(0.0)
if reftype == ObjType.BODY:
xposref = xipos_in[worldid, refid]
xmatref = ximat_in[worldid, refid]
elif reftype == ObjType.XBODY:
xposref = xpos_in[worldid, refid]
xmatref = xmat_in[worldid, refid]
elif reftype == ObjType.GEOM:
xposref = geom_xpos_in[worldid, refid]
xmatref = geom_xmat_in[worldid, refid]
elif reftype == ObjType.SITE:
xposref = site_xpos_in[worldid, refid]
xmatref = site_xmat_in[worldid, refid]
elif reftype == ObjType.CAMERA:
xposref = cam_xpos_in[worldid, refid]
xmatref = cam_xmat_in[worldid, refid]
else: # UNKNOWN
xposref = wp.vec3(0.0)
xmatref = wp.identity(3, dtype=float)
cvel, offset = _cvel_offset(
body_rootid,
geom_bodyid,
site_bodyid,
cam_bodyid,
xpos_in,
xipos_in,
geom_xpos_in,
site_xpos_in,
cam_xpos_in,
subtree_com_in,
cvel_in,
worldid,
objtype,
objid,
)
cvelref, offsetref = _cvel_offset(
body_rootid,
geom_bodyid,
site_bodyid,
cam_bodyid,
xpos_in,
xipos_in,
geom_xpos_in,
site_xpos_in,
cam_xpos_in,
subtree_com_in,
cvel_in,
worldid,
reftype,
refid,
)
clinvel = wp.spatial_bottom(cvel)
cangvel = wp.spatial_top(cvel)
cangvelref = wp.spatial_top(cvelref)
xlinvel = clinvel - wp.cross(offset, cangvel)
if refid > -1:
clinvelref = wp.spatial_bottom(cvelref)
xlinvelref = clinvelref - wp.cross(offsetref, cangvelref)
rvec = xpos - xposref
rel_vel = xlinvel - xlinvelref + wp.cross(rvec, cangvelref)
return wp.transpose(xmatref) @ rel_vel
else:
return xlinvel
@wp.func
def _frame_angvel(
# Model:
body_rootid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
cam_bodyid: wp.array(dtype=int),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xmat_in: wp.array2d(dtype=wp.mat33),
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
objtype: int,
refid: int,
reftype: int,
) -> wp.vec3:
cvel, _ = _cvel_offset(
body_rootid,
geom_bodyid,
site_bodyid,
cam_bodyid,
xpos_in,
xipos_in,
geom_xpos_in,
site_xpos_in,
cam_xpos_in,
subtree_com_in,
cvel_in,
worldid,
objtype,
objid,
)
cangvel = wp.spatial_top(cvel)
if refid > -1:
if reftype == ObjType.BODY:
xmatref = ximat_in[worldid, refid]
elif reftype == ObjType.XBODY:
xmatref = xmat_in[worldid, refid]
elif reftype == ObjType.GEOM:
xmatref = geom_xmat_in[worldid, refid]
elif reftype == ObjType.SITE:
xmatref = site_xmat_in[worldid, refid]
elif reftype == ObjType.CAMERA:
xmatref = cam_xmat_in[worldid, refid]
else: # UNKNOWN
xmatref = wp.identity(3, dtype=float)
cvelref, _ = _cvel_offset(
body_rootid,
geom_bodyid,
site_bodyid,
cam_bodyid,
xpos_in,
xipos_in,
geom_xpos_in,
site_xpos_in,
cam_xpos_in,
subtree_com_in,
cvel_in,
worldid,
reftype,
refid,
)
cangvelref = wp.spatial_top(cvelref)
return wp.transpose(xmatref) @ (cangvel - cangvelref)
else:
return cangvel
@wp.func
def _subtree_linvel(subtree_linvel_in: wp.array2d(dtype=wp.vec3), worldid: int, objid: int) -> wp.vec3:
return subtree_linvel_in[worldid, objid]
@wp.func
def _subtree_angmom(subtree_angmom_in: wp.array2d(dtype=wp.vec3), worldid: int, objid: int) -> wp.vec3:
return subtree_angmom_in[worldid, objid]
@wp.kernel
def _sensor_vel(
# Model:
body_rootid: wp.array(dtype=int),
jnt_dofadr: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
cam_bodyid: wp.array(dtype=int),
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_objtype: wp.array(dtype=int),
sensor_objid: wp.array(dtype=int),
sensor_reftype: wp.array(dtype=int),
sensor_refid: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
sensor_vel_adr: wp.array(dtype=int),
# Data in:
qvel_in: wp.array2d(dtype=float),
xpos_in: wp.array2d(dtype=wp.vec3),
xmat_in: wp.array2d(dtype=wp.mat33),
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
ten_velocity_in: wp.array2d(dtype=float),
actuator_velocity_in: wp.array2d(dtype=float),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
subtree_linvel_in: wp.array2d(dtype=wp.vec3),
subtree_angmom_in: wp.array2d(dtype=wp.vec3),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, velid = wp.tid()
sensorid = sensor_vel_adr[velid]
sensortype = sensor_type[sensorid]
objid = sensor_objid[sensorid]
out = sensordata_out[worldid]
if sensortype == SensorType.VELOCIMETER:
vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.GYRO:
vec3 = _gyro(site_bodyid, site_xmat_in, cvel_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.JOINTVEL:
val = _joint_vel(jnt_dofadr, qvel_in, worldid, objid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.TENDONVEL:
val = _tendon_vel(ten_velocity_in, worldid, objid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.ACTUATORVEL:
val = _actuator_vel(actuator_velocity_in, worldid, objid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.BALLANGVEL:
vec3 = _ball_ang_vel(jnt_dofadr, qvel_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.FRAMELINVEL:
objtype = sensor_objtype[sensorid]
refid = sensor_refid[sensorid]
reftype = sensor_reftype[sensorid]
frame_linvel = _frame_linvel(
body_rootid,
geom_bodyid,
site_bodyid,
cam_bodyid,
xpos_in,
xmat_in,
xipos_in,
ximat_in,
geom_xpos_in,
geom_xmat_in,
site_xpos_in,
site_xmat_in,
cam_xpos_in,
cam_xmat_in,
subtree_com_in,
cvel_in,
worldid,
objid,
objtype,
refid,
reftype,
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_linvel, out)
elif sensortype == SensorType.FRAMEANGVEL:
objtype = sensor_objtype[sensorid]
refid = sensor_refid[sensorid]
reftype = sensor_reftype[sensorid]
frame_angvel = _frame_angvel(
body_rootid,
geom_bodyid,
site_bodyid,
cam_bodyid,
xpos_in,
xmat_in,
xipos_in,
ximat_in,
geom_xpos_in,
geom_xmat_in,
site_xpos_in,
site_xmat_in,
cam_xpos_in,
cam_xmat_in,
subtree_com_in,
cvel_in,
worldid,
objid,
objtype,
refid,
reftype,
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_angvel, out)
elif sensortype == SensorType.SUBTREELINVEL:
vec3 = _subtree_linvel(subtree_linvel_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.SUBTREEANGMOM:
vec3 = _subtree_angmom(subtree_angmom_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
[docs]
@event_scope
def sensor_vel(m: Model, d: Data):
"""Compute velocity-dependent sensor values."""
if m.opt.disableflags & DisableBit.SENSOR:
return
if m.sensor_subtree_vel:
smooth.subtree_vel(m, d)
wp.launch(
_sensor_vel,
dim=(d.nworld, m.sensor_vel_adr.size),
inputs=[
m.body_rootid,
m.jnt_dofadr,
m.geom_bodyid,
m.site_bodyid,
m.cam_bodyid,
m.sensor_type,
m.sensor_datatype,
m.sensor_objtype,
m.sensor_objid,
m.sensor_reftype,
m.sensor_refid,
m.sensor_adr,
m.sensor_cutoff,
m.sensor_vel_adr,
d.qvel,
d.xpos,
d.xmat,
d.xipos,
d.ximat,
d.geom_xpos,
d.geom_xmat,
d.site_xpos,
d.site_xmat,
d.cam_xpos,
d.cam_xmat,
d.subtree_com,
d.ten_velocity,
d.actuator_velocity,
d.cvel,
d.subtree_linvel,
d.subtree_angmom,
],
outputs=[d.sensordata],
)
wp.launch(
_limit_vel,
dim=(d.nworld, d.njmax, m.sensor_limitvel_adr.size),
inputs=[
m.sensor_type,
m.sensor_datatype,
m.sensor_objid,
m.sensor_adr,
m.sensor_cutoff,
m.sensor_limitvel_adr,
d.ne,
d.nf,
d.nl,
d.efc.type,
d.efc.id,
d.efc.vel,
],
outputs=[
d.sensordata,
],
)
@wp.func
def _accelerometer(
# Model:
body_rootid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
# Data in:
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
cacc_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
) -> wp.vec3:
bodyid = site_bodyid[objid]
rot = site_xmat_in[worldid, objid]
rotT = wp.transpose(rot)
cvel = cvel_in[worldid, bodyid]
cvel_top = wp.spatial_top(cvel)
cvel_bottom = wp.spatial_bottom(cvel)
cacc = cacc_in[worldid, bodyid]
cacc_top = wp.spatial_top(cacc)
cacc_bottom = wp.spatial_bottom(cacc)
dif = site_xpos_in[worldid, objid] - subtree_com_in[worldid, body_rootid[bodyid]]
ang = rotT @ cvel_top
lin = rotT @ (cvel_bottom - wp.cross(dif, cvel_top))
acc = rotT @ (cacc_bottom - wp.cross(dif, cacc_top))
correction = wp.cross(ang, lin)
return acc + correction
@wp.func
def _force(
# Model:
site_bodyid: wp.array(dtype=int),
# Data in:
site_xmat_in: wp.array2d(dtype=wp.mat33),
cfrc_int_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
) -> wp.vec3:
bodyid = site_bodyid[objid]
cfrc_int = cfrc_int_in[worldid, bodyid]
site_xmat = site_xmat_in[worldid, objid]
return wp.transpose(site_xmat) @ wp.spatial_bottom(cfrc_int)
@wp.func
def _torque(
# Model:
body_rootid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
# Data in:
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cfrc_int_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
) -> wp.vec3:
bodyid = site_bodyid[objid]
cfrc_int = cfrc_int_in[worldid, bodyid]
site_xmat = site_xmat_in[worldid, objid]
dif = site_xpos_in[worldid, objid] - subtree_com_in[worldid, body_rootid[bodyid]]
return wp.transpose(site_xmat) @ (wp.spatial_top(cfrc_int) - wp.cross(dif, wp.spatial_bottom(cfrc_int)))
@wp.func
def _actuator_force(actuator_force_in: wp.array2d(dtype=float), worldid: int, objid: int) -> float:
return actuator_force_in[worldid, objid]
@wp.func
def _joint_actuator_force(
# Model:
jnt_dofadr: wp.array(dtype=int),
# Data in:
qfrc_actuator_in: wp.array2d(dtype=float),
# In:
worldid: int,
objid: int,
) -> float:
return qfrc_actuator_in[worldid, jnt_dofadr[objid]]
@wp.kernel
def _tendon_actuator_force(
# Model:
actuator_trntype: wp.array(dtype=int),
actuator_trnid: wp.array(dtype=wp.vec2i),
sensor_objid: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_tendonactfrc_adr: wp.array(dtype=int),
# Data in:
actuator_force_in: wp.array2d(dtype=float),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, tenactfrcid, actid = wp.tid()
sensorid = sensor_tendonactfrc_adr[tenactfrcid]
if actuator_trntype[actid] == TrnType.TENDON and actuator_trnid[actid][0] == sensor_objid[sensorid]:
adr = sensor_adr[sensorid]
sensordata_out[worldid, adr] += actuator_force_in[worldid, actid]
@wp.kernel
def _tendon_actuator_force_cutoff(
# Model:
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
sensor_tendonactfrc_adr: wp.array(dtype=int),
# Data in:
sensordata_in: wp.array2d(dtype=float),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, tenactfrcid = wp.tid()
sensorid = sensor_tendonactfrc_adr[tenactfrcid]
adr = sensor_adr[sensorid]
val = sensordata_in[worldid, adr]
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, sensordata_out[worldid])
@wp.kernel
def _limit_frc(
# Model:
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_objid: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
sensor_limitfrc_adr: wp.array(dtype=int),
# Data in:
ne_in: wp.array(dtype=int),
nf_in: wp.array(dtype=int),
nl_in: wp.array(dtype=int),
efc_type_in: wp.array2d(dtype=int),
efc_id_in: wp.array2d(dtype=int),
efc_force_in: wp.array2d(dtype=float),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, efcid, limitfrcid = wp.tid()
ne = ne_in[worldid]
nf = nf_in[worldid]
nl = nl_in[worldid]
# skip if not limit
if efcid < ne + nf or efcid >= ne + nf + nl:
return
sensorid = sensor_limitfrc_adr[limitfrcid]
if efc_id_in[worldid, efcid] == sensor_objid[sensorid]:
efc_type = efc_type_in[worldid, efcid]
if efc_type == ConstraintType.LIMIT_JOINT or efc_type == ConstraintType.LIMIT_TENDON:
_write_scalar(
sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, efc_force_in[worldid, efcid], sensordata_out[worldid]
)
@wp.func
def _framelinacc(
# Model:
body_rootid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
cam_bodyid: wp.array(dtype=int),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xipos_in: wp.array2d(dtype=wp.vec3),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
site_xpos_in: wp.array2d(dtype=wp.vec3),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
cacc_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
objtype: int,
) -> wp.vec3:
if objtype == ObjType.BODY:
bodyid = objid
pos = xipos_in[worldid, objid]
elif objtype == ObjType.XBODY:
bodyid = objid
pos = xpos_in[worldid, objid]
elif objtype == ObjType.GEOM:
bodyid = geom_bodyid[objid]
pos = geom_xpos_in[worldid, objid]
elif objtype == ObjType.SITE:
bodyid = site_bodyid[objid]
pos = site_xpos_in[worldid, objid]
elif objtype == ObjType.CAMERA:
bodyid = cam_bodyid[objid]
pos = cam_xpos_in[worldid, objid]
else: # UNKNOWN
bodyid = 0
pos = wp.vec3(0.0)
cacc = cacc_in[worldid, bodyid]
cvel = cvel_in[worldid, bodyid]
offset = pos - subtree_com_in[worldid, body_rootid[bodyid]]
ang = wp.spatial_top(cvel)
lin = wp.spatial_bottom(cvel) - wp.cross(offset, ang)
acc = wp.spatial_bottom(cacc) - wp.cross(offset, wp.spatial_top(cacc))
correction = wp.cross(ang, lin)
return acc + correction
@wp.func
def _frameangacc(
# Model:
geom_bodyid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
cam_bodyid: wp.array(dtype=int),
# Data in:
cacc_in: wp.array2d(dtype=wp.spatial_vector),
# In:
worldid: int,
objid: int,
objtype: int,
) -> wp.vec3:
if objtype == ObjType.BODY or objtype == ObjType.XBODY:
bodyid = objid
elif objtype == ObjType.GEOM:
bodyid = geom_bodyid[objid]
elif objtype == ObjType.SITE:
bodyid = site_bodyid[objid]
elif objtype == ObjType.CAMERA:
bodyid = cam_bodyid[objid]
else: # UNKNOWN
bodyid = 0
return wp.spatial_top(cacc_in[worldid, bodyid])
@wp.kernel
def _sensor_acc(
# Model:
opt_cone: int,
body_rootid: wp.array(dtype=int),
jnt_dofadr: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
cam_bodyid: wp.array(dtype=int),
sensor_type: wp.array(dtype=int),
sensor_datatype: wp.array(dtype=int),
sensor_objtype: wp.array(dtype=int),
sensor_objid: wp.array(dtype=int),
sensor_intprm: wp.array2d(dtype=int),
sensor_dim: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_cutoff: wp.array(dtype=float),
sensor_acc_adr: wp.array(dtype=int),
sensor_adr_to_contact_adr: wp.array(dtype=int),
# Data in:
xpos_in: wp.array2d(dtype=wp.vec3),
xipos_in: wp.array2d(dtype=wp.vec3),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
cam_xpos_in: wp.array2d(dtype=wp.vec3),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
actuator_force_in: wp.array2d(dtype=float),
qfrc_actuator_in: wp.array2d(dtype=float),
cacc_in: wp.array2d(dtype=wp.spatial_vector),
cfrc_int_in: wp.array2d(dtype=wp.spatial_vector),
contact_dist_in: wp.array(dtype=float),
contact_pos_in: wp.array(dtype=wp.vec3),
contact_frame_in: wp.array(dtype=wp.mat33),
contact_friction_in: wp.array(dtype=vec5),
contact_dim_in: wp.array(dtype=int),
contact_efc_address_in: wp.array2d(dtype=int),
efc_force_in: wp.array2d(dtype=float),
njmax_in: int,
nacon_in: wp.array(dtype=int),
# In:
sensor_contact_nmatch_in: wp.array2d(dtype=int),
sensor_contact_matchid_in: wp.array3d(dtype=int),
sensor_contact_direction_in: wp.array3d(dtype=float),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
worldid, accid = wp.tid()
sensorid = sensor_acc_adr[accid]
sensortype = sensor_type[sensorid]
objid = sensor_objid[sensorid]
out = sensordata_out[worldid]
if sensortype == SensorType.CONTACT:
dataspec = sensor_intprm[sensorid, 0]
dim = sensor_dim[sensorid]
objtype = sensor_objtype[sensorid]
reduce = sensor_intprm[sensorid, 1]
# found, force, torque, dist, pos, normal, tangent
# TODO(thowell): precompute slot size
found = False
force = False
torque = False
dist = False
pos = False
normal = False
tangent = False
size = int(0)
for i in range(7):
if dataspec & (1 << i):
if i == 0:
found = True
size += 1
elif i == 1:
force = True
size += 3
elif i == 2:
torque = True
size += 3
elif i == 3:
dist = True
size += 1
elif i == 4:
pos = True
size += 3
elif i == 5:
normal = True
size += 3
elif i == 6:
tangent = True
size += 3
num = dim // size # number of slots
adr = sensor_adr[sensorid]
contactsensorid = sensor_adr_to_contact_adr[sensorid]
nmatch = sensor_contact_nmatch_in[worldid, contactsensorid]
if reduce == 3: # netforce
# Single-pass computation: first compute centroid, then wrench about centroid
# Pass 1: compute force-weighted centroid of contact positions
net_pos = wp.vec3(0.0)
net_force = wp.vec3(0.0)
net_torque = wp.vec3(0.0)
total_force_magnitude = float(0.0)
for i in range(nmatch):
cid = sensor_contact_matchid_in[worldid, contactsensorid, i]
dir = sensor_contact_direction_in[worldid, contactsensorid, i]
contact_forcetorque = support.contact_force_fn(
opt_cone,
contact_frame_in,
contact_friction_in,
contact_dim_in,
contact_efc_address_in,
efc_force_in,
njmax_in,
nacon_in,
worldid,
cid,
False,
)
# Accumulate for centroid computation (unsigned force magnitude)
weight = wp.norm_l2(wp.spatial_top(contact_forcetorque))
contact_pos = contact_pos_in[cid]
net_pos += weight * contact_pos
total_force_magnitude += weight
# Apply direction and transform to global frame
contact_forcetorque *= dir
force_local = wp.spatial_top(contact_forcetorque)
torque_local = wp.spatial_bottom(contact_forcetorque)
frame = contact_frame_in[cid]
frameT = wp.transpose(frame)
force_global = frameT @ force_local
torque_global = frameT @ torque_local
# Accumulate force and torque (about origin for now)
net_force += force_global
net_torque += torque_global
# Accumulate moment contribution: will adjust after centroid is computed
net_torque += wp.cross(contact_pos, force_global)
# Finalize centroid
net_pos /= wp.max(total_force_magnitude, MJ_MINVAL)
# Adjust torque: subtract moment from centroid (since we accumulated about origin)
# torque_about_centroid = torque_about_origin - centroid x total_force
net_torque -= wp.cross(net_pos, net_force)
adr_slot = adr
if found:
out[adr_slot] = float(nmatch)
adr_slot += 1
if force:
out[adr_slot + 0] = net_force[0]
out[adr_slot + 1] = net_force[1]
out[adr_slot + 2] = net_force[2]
adr_slot += 3
if torque:
out[adr_slot + 0] = net_torque[0]
out[adr_slot + 1] = net_torque[1]
out[adr_slot + 2] = net_torque[2]
adr_slot += 3
if dist:
out[adr_slot] = 0.0
adr_slot += 1
if pos:
out[adr_slot + 0] = net_pos[0]
out[adr_slot + 1] = net_pos[1]
out[adr_slot + 2] = net_pos[2]
adr_slot += 3
if normal:
out[adr_slot + 0] = 1.0
out[adr_slot + 1] = 0.0
out[adr_slot + 2] = 0.0
adr_slot += 3
if tangent:
out[adr_slot + 0] = 0.0
out[adr_slot + 1] = 1.0
out[adr_slot + 2] = 0.0
else:
nslots = wp.min(nmatch, num)
for i in range(nslots):
# sorted contact id
cid = sensor_contact_matchid_in[worldid, contactsensorid, i]
# contact direction
dir = sensor_contact_direction_in[worldid, contactsensorid, i]
adr_slot = adr + i * size
if found:
out[adr_slot] = float(nmatch)
adr_slot += 1
if force or torque:
contact_forcetorque = support.contact_force_fn(
opt_cone,
contact_frame_in,
contact_friction_in,
contact_dim_in,
contact_efc_address_in,
efc_force_in,
njmax_in,
nacon_in,
worldid,
cid,
False,
)
if force:
out[adr_slot + 0] = contact_forcetorque[0]
out[adr_slot + 1] = contact_forcetorque[1]
out[adr_slot + 2] = dir * contact_forcetorque[2]
adr_slot += 3
if torque:
out[adr_slot + 0] = contact_forcetorque[3]
out[adr_slot + 1] = contact_forcetorque[4]
out[adr_slot + 2] = dir * contact_forcetorque[5]
adr_slot += 3
if dist:
out[adr_slot] = contact_dist_in[cid]
adr_slot += 1
if pos:
contact_pos = contact_pos_in[cid]
out[adr_slot + 0] = contact_pos[0]
out[adr_slot + 1] = contact_pos[1]
out[adr_slot + 2] = contact_pos[2]
adr_slot += 3
if normal:
contact_normal = contact_frame_in[cid][0]
out[adr_slot + 0] = dir * contact_normal[0]
out[adr_slot + 1] = dir * contact_normal[1]
out[adr_slot + 2] = dir * contact_normal[2]
adr_slot += 3
if tangent:
contact_tangent = contact_frame_in[cid][1]
out[adr_slot + 0] = dir * contact_tangent[0]
out[adr_slot + 1] = dir * contact_tangent[1]
out[adr_slot + 2] = dir * contact_tangent[2]
# zero remaining slots
for i in range(nmatch, num):
for j in range(size):
out[adr + i * size + j] = 0.0
elif sensortype == SensorType.ACCELEROMETER:
vec3 = _accelerometer(
body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.FORCE:
vec3 = _force(site_bodyid, site_xmat_in, cfrc_int_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.TORQUE:
vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cfrc_int_in, worldid, objid)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.ACTUATORFRC:
val = _actuator_force(actuator_force_in, worldid, objid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.JOINTACTFRC:
val = _joint_actuator_force(jnt_dofadr, qfrc_actuator_in, worldid, objid)
_write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out)
elif sensortype == SensorType.FRAMELINACC:
objtype = sensor_objtype[sensorid]
vec3 = _framelinacc(
body_rootid,
geom_bodyid,
site_bodyid,
cam_bodyid,
xpos_in,
xipos_in,
geom_xpos_in,
site_xpos_in,
cam_xpos_in,
subtree_com_in,
cvel_in,
cacc_in,
worldid,
objid,
objtype,
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
elif sensortype == SensorType.FRAMEANGACC:
objtype = sensor_objtype[sensorid]
vec3 = _frameangacc(
geom_bodyid,
site_bodyid,
cam_bodyid,
cacc_in,
worldid,
objid,
objtype,
)
_write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
@wp.kernel
def _sensor_touch(
# Model:
opt_cone: int,
geom_bodyid: wp.array(dtype=int),
site_type: wp.array(dtype=int),
site_bodyid: wp.array(dtype=int),
site_size: wp.array(dtype=wp.vec3),
sensor_objid: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
sensor_touch_adr: wp.array(dtype=int),
# Data in:
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
contact_pos_in: wp.array(dtype=wp.vec3),
contact_frame_in: wp.array(dtype=wp.mat33),
contact_dim_in: wp.array(dtype=int),
contact_geom_in: wp.array(dtype=wp.vec2i),
contact_efc_address_in: wp.array2d(dtype=int),
contact_worldid_in: wp.array(dtype=int),
efc_force_in: wp.array2d(dtype=float),
nacon_in: wp.array(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
conid, sensortouchadrid = wp.tid()
if conid >= nacon_in[0]:
return
sensorid = sensor_touch_adr[sensortouchadrid]
objid = sensor_objid[sensorid]
bodyid = site_bodyid[objid]
# find contact in sensor zone, add normal force
# contacting bodies
geom = contact_geom_in[conid]
conbody = wp.vec2i(geom_bodyid[geom[0]], geom_bodyid[geom[1]])
# select contacts involving sensorized body
worldid = contact_worldid_in[conid]
efc_address0 = contact_efc_address_in[conid, 0]
if efc_address0 >= 0 and (bodyid == conbody[0] or bodyid == conbody[1]):
# get contact normal force
normalforce = efc_force_in[worldid, efc_address0]
if opt_cone == ConeType.PYRAMIDAL:
dim = contact_dim_in[conid]
for i in range(1, 2 * (dim - 1)):
normalforce += efc_force_in[worldid, contact_efc_address_in[conid, i]]
if normalforce <= 0.0:
return
# convert contact normal force to global frame, normalize
frame = contact_frame_in[conid]
conray = wp.vec3(frame[0, 0], frame[0, 1], frame[0, 2]) * normalforce
conray, _ = math.normalize_with_norm(conray)
# flip ray direction if sensor is on body2
if bodyid == conbody[1]:
conray = -conray
# add if ray-zone intersection (always true when contact.pos inside zone)
dist, normal = ray.ray_geom(
site_xpos_in[worldid, objid],
site_xmat_in[worldid, objid],
site_size[objid],
contact_pos_in[conid],
conray,
site_type[objid],
)
if dist >= 0.0:
adr = sensor_adr[sensorid]
wp.atomic_add(sensordata_out[worldid], adr, normalforce)
@wp.func
def _transform_spatial(vec: wp.spatial_vector, dif: wp.vec3) -> wp.vec3:
return wp.spatial_bottom(vec) - wp.cross(dif, wp.spatial_top(vec))
@wp.kernel
def _sensor_tactile(
# Model:
body_rootid: wp.array(dtype=int),
body_weldid: wp.array(dtype=int),
oct_child: wp.array(dtype=vec8i),
oct_aabb: wp.array2d(dtype=wp.vec3),
oct_coeff: wp.array(dtype=vec8),
geom_type: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
mesh_vertadr: wp.array(dtype=int),
mesh_normaladr: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
mesh_normal: wp.array(dtype=wp.vec3),
mesh_quat: wp.array(dtype=wp.quat),
sensor_objid: wp.array(dtype=int),
sensor_refid: wp.array(dtype=int),
sensor_dim: wp.array(dtype=int),
sensor_adr: wp.array(dtype=int),
plugin: wp.array(dtype=int),
plugin_attr: wp.array(dtype=wp.vec3f),
geom_plugin_index: wp.array(dtype=int),
taxel_vertadr: wp.array(dtype=int),
taxel_sensorid: wp.array(dtype=int),
# Data in:
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
contact_geom_in: wp.array(dtype=wp.vec2i),
contact_worldid_in: wp.array(dtype=int),
nacon_in: wp.array(dtype=int),
# Data out:
sensordata_out: wp.array2d(dtype=float),
):
conid, taxelid = wp.tid()
if conid >= nacon_in[0]:
return
worldid = contact_worldid_in[conid]
# get sensor_id
sensor_id = taxel_sensorid[taxelid]
# get parent weld id
mesh_id = sensor_objid[sensor_id]
geom_id = sensor_refid[sensor_id]
parent_body = geom_bodyid[geom_id]
parent_weld = body_weldid[parent_body]
# contact geom
body1 = body_weldid[geom_bodyid[contact_geom_in[conid][0]]]
body2 = body_weldid[geom_bodyid[contact_geom_in[conid][1]]]
if body1 == parent_weld:
geom = contact_geom_in[conid][1]
elif body2 == parent_weld:
geom = contact_geom_in[conid][0]
else:
return
body = geom_bodyid[geom]
# vertex local position
vertid = taxel_vertadr[taxelid] - mesh_vertadr[mesh_id]
pos = mesh_vert[vertid + mesh_vertadr[mesh_id]]
# position in global frame
xpos = geom_xmat_in[worldid, geom_id] @ pos
xpos += geom_xpos_in[worldid, geom_id]
# position in other geom frame
tmp = xpos - geom_xpos_in[worldid, geom]
lpos = wp.transpose(geom_xmat_in[worldid, geom]) @ tmp
plugin_id = geom_plugin_index[geom]
contact_type = geom_type[geom]
plugin_attributes, plugin_index, volume_data, mesh_data = get_sdf_params(
oct_child,
oct_aabb,
oct_coeff,
plugin,
plugin_attr,
contact_type,
geom_size[worldid % geom_size.shape[0], geom],
plugin_id,
mesh_id,
)
depth = wp.min(sdf(contact_type, lpos, plugin_attributes, plugin_index, volume_data, mesh_data), 0.0)
if depth >= 0.0:
return
# get velocity in global
vel_sensor = _transform_spatial(cvel_in[worldid, parent_weld], xpos - subtree_com_in[worldid, body_rootid[parent_weld]])
vel_other = _transform_spatial(
cvel_in[worldid, body], geom_xpos_in[worldid, geom] - subtree_com_in[worldid, body_rootid[body]]
)
vel_rel = vel_sensor - vel_other
# get contact force/torque, rotate into node frame
offset = mesh_normaladr[mesh_id] + 3 * vertid
normal = math.rot_vec_quat(mesh_normal[offset], mesh_quat[mesh_id])
tang1 = math.rot_vec_quat(mesh_normal[offset + 1], mesh_quat[mesh_id])
tang2 = math.rot_vec_quat(mesh_normal[offset + 2], mesh_quat[mesh_id])
kMaxDepth = 0.05
pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
force = wp.mul(normal, pressure)
# one row of mat^T * force
forceT = wp.vec3()
forceT[0] = wp.dot(force, normal)
forceT[1] = wp.abs(wp.dot(vel_rel, tang1))
forceT[2] = wp.abs(wp.dot(vel_rel, tang2))
# add to sensor output
dim = sensor_dim[sensor_id] / 3
wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
wp.atomic_add(sensordata_out[worldid], sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
@wp.func
def _check_match(body_parentid: wp.array(dtype=int), body: int, geom: int, objtype: int, objid: int) -> bool:
"""Check if a contact body/geom matches a sensor spec (objtype, objid)."""
if objtype == ObjType.UNKNOWN:
return True
if objtype == ObjType.SITE:
return True # already passed site filter test
if objtype == ObjType.GEOM:
return objid == geom
if objtype == ObjType.BODY:
return objid == body
if objtype == ObjType.XBODY:
# traverse up the tree from body, return true if we land on id
while body > objid:
body = body_parentid[body]
return body == objid
return False
@wp.kernel
def _contact_match(
# Model:
opt_cone: int,
opt_contact_sensor_maxmatch: int,
body_parentid: wp.array(dtype=int),
geom_bodyid: wp.array(dtype=int),
site_type: wp.array(dtype=int),
site_size: wp.array(dtype=wp.vec3),
sensor_objtype: wp.array(dtype=int),
sensor_objid: wp.array(dtype=int),
sensor_reftype: wp.array(dtype=int),
sensor_refid: wp.array(dtype=int),
sensor_intprm: wp.array2d(dtype=int),
sensor_contact_adr: wp.array(dtype=int),
# Data in:
site_xpos_in: wp.array2d(dtype=wp.vec3),
site_xmat_in: wp.array2d(dtype=wp.mat33),
contact_dist_in: wp.array(dtype=float),
contact_pos_in: wp.array(dtype=wp.vec3),
contact_frame_in: wp.array(dtype=wp.mat33),
contact_friction_in: wp.array(dtype=vec5),
contact_dim_in: wp.array(dtype=int),
contact_geom_in: wp.array(dtype=wp.vec2i),
contact_efc_address_in: wp.array2d(dtype=int),
contact_worldid_in: wp.array(dtype=int),
contact_type_in: wp.array(dtype=int),
efc_force_in: wp.array2d(dtype=float),
njmax_in: int,
nacon_in: wp.array(dtype=int),
# Out:
sensor_contact_nmatch_out: wp.array2d(dtype=int),
sensor_contact_matchid_out: wp.array3d(dtype=int),
sensor_contact_criteria_out: wp.array3d(dtype=float),
sensor_contact_direction_out: wp.array3d(dtype=float),
):
contactsensorid, contactid = wp.tid()
sensorid = sensor_contact_adr[contactsensorid]
if contactid >= nacon_in[0]:
return
if not contact_type_in[contactid] & ContactType.CONSTRAINT:
return
# sensor information
objtype = sensor_objtype[sensorid]
objid = sensor_objid[sensorid]
reftype = sensor_reftype[sensorid]
refid = sensor_refid[sensorid]
reduce = sensor_intprm[sensorid, 1]
worldid = contact_worldid_in[contactid]
# site filter
if objtype == ObjType.SITE:
if not inside_geom(
site_xpos_in[worldid, objid], site_xmat_in[worldid, objid], site_size[objid], site_type[objid], contact_pos_in[contactid]
):
return
# unknown-unknown match
if objtype == ObjType.UNKNOWN and reftype == ObjType.UNKNOWN:
dir = 1.0
else:
# contact information
geom = contact_geom_in[contactid]
geom1 = geom[0]
geom2 = geom[1]
body1 = geom_bodyid[geom1]
body2 = geom_bodyid[geom2]
# check match of sensor objects with contact objects
match11 = _check_match(body_parentid, body1, geom1, objtype, objid)
match12 = _check_match(body_parentid, body2, geom2, objtype, objid)
match21 = _check_match(body_parentid, body1, geom1, reftype, refid)
match22 = _check_match(body_parentid, body2, geom2, reftype, refid)
# if a sensor object is specified, it must be involved in the contact
if not match11 and not match12:
return
if not match21 and not match22:
return
# determine direction
dir = 1.0
if objtype != ObjType.UNKNOWN and reftype != ObjType.UNKNOWN:
# both obj1 and obj2 specified: direction depends on order
order_regular = match11 and match22
order_reverse = match12 and match21
if not order_regular and not order_reverse:
return
if order_reverse and not order_regular:
dir = -1.0
elif objtype != ObjType.UNKNOWN:
if not match11:
dir = -1.0
elif reftype != ObjType.UNKNOWN:
if not match22:
dir = -1.0
contactmatchid = wp.atomic_add(sensor_contact_nmatch_out[worldid], contactsensorid, 1)
if contactmatchid >= opt_contact_sensor_maxmatch:
# TODO(team): alternative to wp.printf for reporting overflow?
wp.printf("contact match overflow: please increase Option.contact_sensor_maxmatch to %u\n", contactmatchid)
return
sensor_contact_matchid_out[worldid, contactsensorid, contactmatchid] = contactid
if reduce == 1: # mindist
sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = contact_dist_in[contactid]
elif reduce == 2: # maxforce
contact_force = support.contact_force_fn(
opt_cone,
contact_frame_in,
contact_friction_in,
contact_dim_in,
contact_efc_address_in,
efc_force_in,
njmax_in,
nacon_in,
worldid,
contactid,
False,
)
force_magnitude = (
contact_force[0] * contact_force[0] + contact_force[1] * contact_force[1] + contact_force[2] * contact_force[2]
)
sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = -force_magnitude
# contact direction
sensor_contact_direction_out[worldid, contactsensorid, contactmatchid] = dir
@cache_kernel
def _contact_sort(maxmatch: int):
@wp.kernel(module="unique", enable_backward=False)
def contact_sort(
# Model:
sensor_intprm: wp.array2d(dtype=int),
sensor_contact_adr: wp.array(dtype=int),
# In:
sensor_contact_nmatch_in: wp.array2d(dtype=int),
sensor_contact_matchid_in: wp.array3d(dtype=int),
sensor_contact_criteria_in: wp.array3d(dtype=float),
# Out:
sensor_contact_matchid_out: wp.array3d(dtype=int),
):
worldid, contactsensorid = wp.tid()
worldid, contactsensorid = wp.tid()
sensorid = sensor_contact_adr[contactsensorid]
reduce = sensor_intprm[sensorid, 1]
if reduce == 0 or reduce == 3: # none or netforce
return
nmatch = sensor_contact_nmatch_in[worldid, contactsensorid]
# skip sort
if nmatch <= 1:
return
criteria_tile = wp.tile_load(sensor_contact_criteria_in[worldid, contactsensorid], shape=maxmatch)
matchid_tile = wp.tile_load(sensor_contact_matchid_in[worldid, contactsensorid], shape=maxmatch)
wp.tile_sort(criteria_tile, matchid_tile)
wp.tile_store(sensor_contact_matchid_out[worldid, contactsensorid], matchid_tile)
return contact_sort
[docs]
@event_scope
def sensor_acc(m: Model, d: Data):
"""Compute acceleration-dependent sensor values."""
if m.opt.disableflags & DisableBit.SENSOR:
return
wp.launch(
_sensor_touch,
dim=(d.naconmax, m.sensor_touch_adr.size),
inputs=[
m.opt.cone,
m.geom_bodyid,
m.site_type,
m.site_bodyid,
m.site_size,
m.sensor_objid,
m.sensor_adr,
m.sensor_touch_adr,
d.site_xpos,
d.site_xmat,
d.contact.pos,
d.contact.frame,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
d.efc.force,
d.nacon,
],
outputs=[
d.sensordata,
],
)
wp.launch(
_sensor_tactile,
dim=(d.naconmax, m.nsensortaxel),
inputs=[
m.body_rootid,
m.body_weldid,
m.oct_child,
m.oct_aabb,
m.oct_coeff,
m.geom_type,
m.geom_bodyid,
m.geom_size,
m.mesh_vertadr,
m.mesh_normaladr,
m.mesh_vert,
m.mesh_normal,
m.mesh_quat,
m.sensor_objid,
m.sensor_refid,
m.sensor_dim,
m.sensor_adr,
m.plugin,
m.plugin_attr,
m.geom_plugin_index,
m.taxel_vertadr,
m.taxel_sensorid,
d.geom_xpos,
d.geom_xmat,
d.subtree_com,
d.cvel,
d.contact.geom,
d.contact.worldid,
d.nacon,
],
outputs=[
d.sensordata,
],
)
sensor_contact_nmatch = wp.empty((d.nworld, m.nsensorcontact), dtype=int)
sensor_contact_matchid = wp.empty((d.nworld, m.nsensorcontact, m.opt.contact_sensor_maxmatch), dtype=int)
sensor_contact_direction = wp.empty((d.nworld, m.nsensorcontact, m.opt.contact_sensor_maxmatch), dtype=float)
if m.nsensorcontact:
sensor_contact_criteria = wp.empty((d.nworld, m.nsensorcontact, m.opt.contact_sensor_maxmatch), dtype=float)
# TODO(team): fill_ operations in one kernel?
sensor_contact_nmatch.fill_(0)
sensor_contact_matchid.fill_(-1)
sensor_contact_criteria.fill_(1.0e32)
wp.launch(
_contact_match,
dim=(m.sensor_contact_adr.size, d.naconmax),
inputs=[
m.opt.cone,
m.opt.contact_sensor_maxmatch,
m.body_parentid,
m.geom_bodyid,
m.site_type,
m.site_size,
m.sensor_objtype,
m.sensor_objid,
m.sensor_reftype,
m.sensor_refid,
m.sensor_intprm,
m.sensor_contact_adr,
d.site_xpos,
d.site_xmat,
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.friction,
d.contact.dim,
d.contact.geom,
d.contact.efc_address,
d.contact.worldid,
d.contact.type,
d.efc.force,
d.njmax,
d.nacon,
],
outputs=[sensor_contact_nmatch, sensor_contact_matchid, sensor_contact_criteria, sensor_contact_direction],
)
# sorting
wp.launch_tiled(
_contact_sort(m.opt.contact_sensor_maxmatch),
dim=(d.nworld, m.sensor_contact_adr.size),
inputs=[m.sensor_intprm, m.sensor_contact_adr, sensor_contact_nmatch, sensor_contact_matchid, sensor_contact_criteria],
outputs=[sensor_contact_matchid],
block_dim=m.block_dim.contact_sort,
)
if m.sensor_rne_postconstraint:
smooth.rne_postconstraint(m, d)
wp.launch(
_sensor_acc,
dim=(d.nworld, m.sensor_acc_adr.size),
inputs=[
m.opt.cone,
m.body_rootid,
m.jnt_dofadr,
m.geom_bodyid,
m.site_bodyid,
m.cam_bodyid,
m.sensor_type,
m.sensor_datatype,
m.sensor_objtype,
m.sensor_objid,
m.sensor_intprm,
m.sensor_dim,
m.sensor_adr,
m.sensor_cutoff,
m.sensor_acc_adr,
m.sensor_adr_to_contact_adr,
d.xpos,
d.xipos,
d.geom_xpos,
d.site_xpos,
d.site_xmat,
d.cam_xpos,
d.subtree_com,
d.cvel,
d.actuator_force,
d.qfrc_actuator,
d.cacc,
d.cfrc_int,
d.contact.dist,
d.contact.pos,
d.contact.frame,
d.contact.friction,
d.contact.dim,
d.contact.efc_address,
d.efc.force,
d.njmax,
d.nacon,
sensor_contact_nmatch,
sensor_contact_matchid,
sensor_contact_direction,
],
outputs=[d.sensordata],
)
wp.launch(
_tendon_actuator_force,
dim=(d.nworld, m.sensor_tendonactfrc_adr.size, m.nu),
inputs=[
m.actuator_trntype,
m.actuator_trnid,
m.sensor_objid,
m.sensor_adr,
m.sensor_tendonactfrc_adr,
d.actuator_force,
],
outputs=[
d.sensordata,
],
)
wp.launch(
_tendon_actuator_force_cutoff,
dim=(d.nworld, m.sensor_tendonactfrc_adr.size),
inputs=[
m.sensor_type,
m.sensor_datatype,
m.sensor_adr,
m.sensor_cutoff,
m.sensor_tendonactfrc_adr,
d.sensordata,
],
outputs=[d.sensordata],
)
wp.launch(
_limit_frc,
dim=(d.nworld, d.njmax, m.sensor_limitfrc_adr.size),
inputs=[
m.sensor_type,
m.sensor_datatype,
m.sensor_objid,
m.sensor_adr,
m.sensor_cutoff,
m.sensor_limitfrc_adr,
d.ne,
d.nf,
d.nl,
d.efc.type,
d.efc.id,
d.efc.force,
],
outputs=[
d.sensordata,
],
)
@wp.kernel
def _energy_pos_zero(
# Data out:
energy_out: wp.array(dtype=wp.vec2),
):
worldid = wp.tid()
energy_out[worldid][0] = 0.0
@wp.kernel
def _energy_pos_gravity(
# Model:
opt_gravity: wp.array(dtype=wp.vec3),
body_mass: wp.array2d(dtype=float),
# Data in:
xipos_in: wp.array2d(dtype=wp.vec3),
# Data out:
energy_out: wp.array(dtype=wp.vec2),
):
worldid, bodyid = wp.tid()
gravity = opt_gravity[worldid % opt_gravity.shape[0]]
bodyid += 1 # skip world body
energy = wp.vec2(
body_mass[worldid % body_mass.shape[0], bodyid] * wp.dot(gravity, xipos_in[worldid, bodyid]),
0.0,
)
wp.atomic_sub(energy_out, worldid, energy)
@wp.kernel
def _energy_pos_passive_joint(
# Model:
qpos_spring: wp.array2d(dtype=float),
jnt_type: wp.array(dtype=int),
jnt_qposadr: wp.array(dtype=int),
jnt_stiffness: wp.array2d(dtype=float),
# Data in:
qpos_in: wp.array2d(dtype=float),
# Data out:
energy_out: wp.array(dtype=wp.vec2),
):
worldid, jntid = wp.tid()
jnt_stiffness_id = worldid % jnt_stiffness.shape[0]
stiffness = jnt_stiffness[jnt_stiffness_id, jntid]
if stiffness == 0.0:
return
padr = jnt_qposadr[jntid]
jnttype = jnt_type[jntid]
qpos_spring_id = worldid % qpos_spring.shape[0]
if jnttype == JointType.FREE:
dif0 = wp.vec3(
qpos_in[worldid, padr + 0] - qpos_spring[qpos_spring_id, padr + 0],
qpos_in[worldid, padr + 1] - qpos_spring[qpos_spring_id, padr + 1],
qpos_in[worldid, padr + 2] - qpos_spring[qpos_spring_id, padr + 2],
)
# convert quaternion difference into angular "velocity"
quat1 = wp.quat(
qpos_in[worldid, padr + 3],
qpos_in[worldid, padr + 4],
qpos_in[worldid, padr + 5],
qpos_in[worldid, padr + 6],
)
quat1 = wp.normalize(quat1)
quat_spring = wp.quat(
qpos_spring[qpos_spring_id, padr + 3],
qpos_spring[qpos_spring_id, padr + 4],
qpos_spring[qpos_spring_id, padr + 5],
qpos_spring[qpos_spring_id, padr + 6],
)
dif1 = math.quat_sub(quat1, quat_spring)
energy = wp.vec2(
0.5 * stiffness * (wp.dot(dif0, dif0) + wp.dot(dif1, dif1)),
0.0,
)
wp.atomic_add(energy_out, worldid, energy)
elif jnttype == JointType.BALL:
quat = wp.quat(
qpos_in[worldid, padr + 0],
qpos_in[worldid, padr + 1],
qpos_in[worldid, padr + 2],
qpos_in[worldid, padr + 3],
)
quat = wp.normalize(quat)
quat_spring = wp.quat(
qpos_spring[qpos_spring_id, padr + 0],
qpos_spring[qpos_spring_id, padr + 1],
qpos_spring[qpos_spring_id, padr + 2],
qpos_spring[qpos_spring_id, padr + 3],
)
dif = math.quat_sub(quat, quat_spring)
energy = wp.vec2(
0.5 * stiffness * wp.dot(dif, dif),
0.0,
)
wp.atomic_add(energy_out, worldid, energy)
elif jnttype == JointType.SLIDE or jnttype == JointType.HINGE:
dif_ = qpos_in[worldid, padr] - qpos_spring[qpos_spring_id, padr]
energy = wp.vec2(
0.5 * stiffness * dif_ * dif_,
0.0,
)
wp.atomic_add(energy_out, worldid, energy)
@wp.kernel
def _energy_pos_passive_tendon(
# Model:
tendon_stiffness: wp.array2d(dtype=float),
tendon_lengthspring: wp.array2d(dtype=wp.vec2),
# Data in:
ten_length_in: wp.array2d(dtype=float),
# Data out:
energy_out: wp.array(dtype=wp.vec2),
):
worldid, tenid = wp.tid()
tendon_stiffness_id = worldid % tendon_stiffness.shape[0]
stiffness = tendon_stiffness[tendon_stiffness_id, tenid]
if stiffness == 0.0:
return
length = ten_length_in[worldid, tenid]
# compute spring displacement
tendon_lengthspring_id = worldid % tendon_lengthspring.shape[0]
lengthspring = tendon_lengthspring[tendon_lengthspring_id, tenid]
lower = lengthspring[0]
upper = lengthspring[1]
if length > upper:
displacement = upper - length
elif length < lower:
displacement = lower - length
else:
displacement = 0.0
energy = wp.vec2(0.5 * stiffness * displacement * displacement, 0.0)
wp.atomic_add(energy_out, worldid, energy)
[docs]
def energy_pos(m: Model, d: Data):
"""Position-dependent energy (potential)."""
wp.launch(_energy_pos_zero, dim=d.nworld, outputs=[d.energy])
# init potential energy: -sum_i(body_i.mass * dot(gravity, body_i.pos))
if not m.opt.disableflags & DisableBit.GRAVITY:
wp.launch(
_energy_pos_gravity, dim=(d.nworld, m.nbody - 1), inputs=[m.opt.gravity, m.body_mass, d.xipos], outputs=[d.energy]
)
if not m.opt.disableflags & DisableBit.SPRING:
# add joint-level springs
wp.launch(
_energy_pos_passive_joint,
dim=(d.nworld, m.njnt),
inputs=[
m.qpos_spring,
m.jnt_type,
m.jnt_qposadr,
m.jnt_stiffness,
d.qpos,
],
outputs=[d.energy],
)
# add tendon-level springs
if m.ntendon:
wp.launch(
_energy_pos_passive_tendon,
dim=(d.nworld, m.ntendon),
inputs=[
m.tendon_stiffness,
m.tendon_lengthspring,
d.ten_length,
],
outputs=[d.energy],
)
# TODO(team): flex
@cache_kernel
def _energy_vel_kinetic(nv: int):
@wp.kernel(module="unique", enable_backward=False)
def energy_vel_kinetic(
# Data in:
qvel_in: wp.array2d(dtype=float),
# In:
Mqvel: wp.array2d(dtype=float),
# Data out:
energy_out: wp.array(dtype=wp.vec2),
):
worldid = wp.tid()
qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
Mqvel_tile = wp.tile_load(Mqvel[worldid], shape=wp.static(nv))
# qvel * (M @ qvel)
qvelMqvel_tile = wp.tile_map(wp.mul, qvel_tile, Mqvel_tile)
# sum(qvel * (M @ qvel))
quadratic_tile = wp.tile_reduce(wp.add, qvelMqvel_tile)
energy_out[worldid][1] = 0.5 * quadratic_tile[0]
return energy_vel_kinetic
[docs]
def energy_vel(m: Model, d: Data):
"""Velocity-dependent energy (kinetic)."""
# kinetic energy: 0.5 * qvel.T @ M @ qvel
# M @ qvel
mv = wp.zeros((d.nworld, m.nv), dtype=float)
support.mul_m(m, d, mv, d.qvel)
wp.launch_tiled(
_energy_vel_kinetic(m.nv),
dim=d.nworld,
inputs=[d.qvel, mv],
outputs=[d.energy],
block_dim=m.block_dim.energy_vel_kinetic,
)