Mustafa Hasekioglu
528 lines
16 KiB
Python
528 lines
16 KiB
Python
# Copyright 2026 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.
|
|
# ==============================================================================
|
|
|
|
"""DO NOT EDIT. This file is auto-generated."""
|
|
|
|
import dataclasses
|
|
import functools
|
|
|
|
import jax
|
|
import warp as wp
|
|
|
|
from mujoco.mjx._src import types
|
|
import mujoco.mjx.third_party.mujoco_warp as mjwarp
|
|
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
|
|
from mujoco.mjx.warp import ffi
|
|
|
|
_m = mjwarp.Model(
|
|
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
|
|
)
|
|
_d = mjwarp.Data(
|
|
**{f.name: None for f in dataclasses.fields(mjwarp.Data) if f.init}
|
|
)
|
|
_o = mjwarp.Option(
|
|
**{f.name: None for f in dataclasses.fields(mjwarp.Option) if f.init}
|
|
)
|
|
_s = mjwarp.Statistic(
|
|
**{f.name: None for f in dataclasses.fields(mjwarp.Statistic) if f.init}
|
|
)
|
|
_c = mjwarp.Contact(
|
|
**{f.name: None for f in dataclasses.fields(mjwarp.Contact) if f.init}
|
|
)
|
|
_e = mjwarp.Constraint(
|
|
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
|
|
)
|
|
_cb = mjwp_types.Callback(
|
|
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
|
|
)
|
|
|
|
|
|
@ffi.format_args_for_warp
|
|
def _collision_shim(
|
|
# Model
|
|
nworld: int,
|
|
block_dim: mjwp_types.BlockDim,
|
|
flex_conaffinity: wp.array[int],
|
|
flex_condim: wp.array[int],
|
|
flex_contype: wp.array[int],
|
|
flex_dim: wp.array[int],
|
|
flex_elem: wp.array[int],
|
|
flex_elemadr: wp.array[int],
|
|
flex_elemdataadr: wp.array[int],
|
|
flex_elemnum: wp.array[int],
|
|
flex_friction: wp.array[wp.vec3],
|
|
flex_gap: wp.array[float],
|
|
flex_margin: wp.array[float],
|
|
flex_priority: wp.array[int],
|
|
flex_radius: wp.array[float],
|
|
flex_shell: wp.array[int],
|
|
flex_shelldataadr: wp.array[int],
|
|
flex_shellnum: wp.array[int],
|
|
flex_solimp: wp.array[mjwp_types.vec5],
|
|
flex_solmix: wp.array[float],
|
|
flex_solref: wp.array[wp.vec2],
|
|
flex_vertadr: wp.array[int],
|
|
flex_vertflexid: wp.array[int],
|
|
geom_aabb: wp.array3d[wp.vec3],
|
|
geom_bodyid: wp.array[int],
|
|
geom_conaffinity: wp.array[int],
|
|
geom_condim: wp.array[int],
|
|
geom_contype: wp.array[int],
|
|
geom_dataid: wp.array2d[int],
|
|
geom_friction: wp.array2d[wp.vec3],
|
|
geom_gap: wp.array2d[float],
|
|
geom_margin: wp.array2d[float],
|
|
geom_pair_type_count: tuple[int, ...],
|
|
geom_plugin_index: wp.array[int],
|
|
geom_priority: wp.array[int],
|
|
geom_rbound: wp.array2d[float],
|
|
geom_size: wp.array2d[wp.vec3],
|
|
geom_solimp: wp.array2d[mjwp_types.vec5],
|
|
geom_solmix: wp.array2d[float],
|
|
geom_solref: wp.array2d[wp.vec2],
|
|
geom_type: wp.array[int],
|
|
has_sdf_geom: bool,
|
|
hfield_adr: wp.array[int],
|
|
hfield_data: wp.array[float],
|
|
hfield_ncol: wp.array[int],
|
|
hfield_nrow: wp.array[int],
|
|
hfield_size: wp.array[wp.vec4],
|
|
mesh_face: wp.array[wp.vec3i],
|
|
mesh_faceadr: wp.array[int],
|
|
mesh_graph: wp.array[int],
|
|
mesh_graphadr: wp.array[int],
|
|
mesh_octadr: wp.array[int],
|
|
mesh_polyadr: wp.array[int],
|
|
mesh_polymap: wp.array[int],
|
|
mesh_polymapadr: wp.array[int],
|
|
mesh_polymapnum: wp.array[int],
|
|
mesh_polynormal: wp.array[wp.vec3],
|
|
mesh_polynum: wp.array[int],
|
|
mesh_polyvert: wp.array[int],
|
|
mesh_polyvertadr: wp.array[int],
|
|
mesh_polyvertnum: wp.array[int],
|
|
mesh_vert: wp.array[wp.vec3],
|
|
mesh_vertadr: wp.array[int],
|
|
mesh_vertnum: wp.array[int],
|
|
nflex: int,
|
|
nflexelem: int,
|
|
nflexshelldata: int,
|
|
nflexvert: int,
|
|
ngeom: int,
|
|
nmaxmeshdeg: int,
|
|
nmaxpolygon: int,
|
|
nmeshface: int,
|
|
nxn_geom_pair_filtered: wp.array[wp.vec2i],
|
|
nxn_pairid: wp.array[wp.vec2i],
|
|
nxn_pairid_filtered: wp.array[wp.vec2i],
|
|
oct_aabb: wp.array2d[wp.vec3],
|
|
oct_child: wp.array[mjwp_types.vec8i],
|
|
oct_coeff: wp.array[mjwp_types.vec8],
|
|
pair_dim: wp.array[int],
|
|
pair_friction: wp.array2d[mjwp_types.vec5],
|
|
pair_gap: wp.array2d[float],
|
|
pair_margin: wp.array2d[float],
|
|
pair_solimp: wp.array2d[mjwp_types.vec5],
|
|
pair_solref: wp.array2d[wp.vec2],
|
|
pair_solreffriction: wp.array2d[wp.vec2],
|
|
plugin: wp.array[int],
|
|
plugin_attr: wp.array[mjwp_types.vec_pluginattr],
|
|
opt__broadphase: int,
|
|
opt__broadphase_filter: int,
|
|
opt__ccd_iterations: int,
|
|
opt__ccd_tolerance: wp.array[float],
|
|
opt__disableflags: int,
|
|
opt__enableflags: int,
|
|
opt__sdf_initpoints: int,
|
|
opt__sdf_iterations: int,
|
|
# Data
|
|
naccdmax: int,
|
|
naconmax: int,
|
|
body_awake: wp.array2d[int],
|
|
flexvert_xpos: wp.array2d[wp.vec3],
|
|
geom_xmat: wp.array2d[wp.mat33],
|
|
geom_xpos: wp.array2d[wp.vec3],
|
|
nacon: wp.array[int],
|
|
ncollision: wp.array[int],
|
|
contact__dim: wp.array[int],
|
|
contact__dist: wp.array[float],
|
|
contact__efc_address: wp.array2d[int],
|
|
contact__flex: wp.array[wp.vec2i],
|
|
contact__frame: wp.array[wp.mat33],
|
|
contact__friction: wp.array[mjwp_types.vec5],
|
|
contact__geom: wp.array[wp.vec2i],
|
|
contact__geomcollisionid: wp.array[int],
|
|
contact__includemargin: wp.array[float],
|
|
contact__pos: wp.array[wp.vec3],
|
|
contact__solimp: wp.array[mjwp_types.vec5],
|
|
contact__solref: wp.array[wp.vec2],
|
|
contact__solreffriction: wp.array[wp.vec2],
|
|
contact__type: wp.array[int],
|
|
contact__vert: wp.array[wp.vec2i],
|
|
contact__worldid: wp.array[int],
|
|
):
|
|
_m.stat = _s
|
|
_m.opt = _o
|
|
_m.callback = _cb
|
|
_d.efc = _e
|
|
_d.contact = _c
|
|
_m.block_dim = block_dim
|
|
_m.flex_conaffinity = flex_conaffinity
|
|
_m.flex_condim = flex_condim
|
|
_m.flex_contype = flex_contype
|
|
_m.flex_dim = flex_dim
|
|
_m.flex_elem = flex_elem
|
|
_m.flex_elemadr = flex_elemadr
|
|
_m.flex_elemdataadr = flex_elemdataadr
|
|
_m.flex_elemnum = flex_elemnum
|
|
_m.flex_friction = flex_friction
|
|
_m.flex_gap = flex_gap
|
|
_m.flex_margin = flex_margin
|
|
_m.flex_priority = flex_priority
|
|
_m.flex_radius = flex_radius
|
|
_m.flex_shell = flex_shell
|
|
_m.flex_shelldataadr = flex_shelldataadr
|
|
_m.flex_shellnum = flex_shellnum
|
|
_m.flex_solimp = flex_solimp
|
|
_m.flex_solmix = flex_solmix
|
|
_m.flex_solref = flex_solref
|
|
_m.flex_vertadr = flex_vertadr
|
|
_m.flex_vertflexid = flex_vertflexid
|
|
_m.geom_aabb = geom_aabb
|
|
_m.geom_bodyid = geom_bodyid
|
|
_m.geom_conaffinity = geom_conaffinity
|
|
_m.geom_condim = geom_condim
|
|
_m.geom_contype = geom_contype
|
|
_m.geom_dataid = geom_dataid
|
|
_m.geom_friction = geom_friction
|
|
_m.geom_gap = geom_gap
|
|
_m.geom_margin = geom_margin
|
|
_m.geom_pair_type_count = geom_pair_type_count
|
|
_m.geom_plugin_index = geom_plugin_index
|
|
_m.geom_priority = geom_priority
|
|
_m.geom_rbound = geom_rbound
|
|
_m.geom_size = geom_size
|
|
_m.geom_solimp = geom_solimp
|
|
_m.geom_solmix = geom_solmix
|
|
_m.geom_solref = geom_solref
|
|
_m.geom_type = geom_type
|
|
_m.has_sdf_geom = has_sdf_geom
|
|
_m.hfield_adr = hfield_adr
|
|
_m.hfield_data = hfield_data
|
|
_m.hfield_ncol = hfield_ncol
|
|
_m.hfield_nrow = hfield_nrow
|
|
_m.hfield_size = hfield_size
|
|
_m.mesh_face = mesh_face
|
|
_m.mesh_faceadr = mesh_faceadr
|
|
_m.mesh_graph = mesh_graph
|
|
_m.mesh_graphadr = mesh_graphadr
|
|
_m.mesh_octadr = mesh_octadr
|
|
_m.mesh_polyadr = mesh_polyadr
|
|
_m.mesh_polymap = mesh_polymap
|
|
_m.mesh_polymapadr = mesh_polymapadr
|
|
_m.mesh_polymapnum = mesh_polymapnum
|
|
_m.mesh_polynormal = mesh_polynormal
|
|
_m.mesh_polynum = mesh_polynum
|
|
_m.mesh_polyvert = mesh_polyvert
|
|
_m.mesh_polyvertadr = mesh_polyvertadr
|
|
_m.mesh_polyvertnum = mesh_polyvertnum
|
|
_m.mesh_vert = mesh_vert
|
|
_m.mesh_vertadr = mesh_vertadr
|
|
_m.mesh_vertnum = mesh_vertnum
|
|
_m.nflex = nflex
|
|
_m.nflexelem = nflexelem
|
|
_m.nflexshelldata = nflexshelldata
|
|
_m.nflexvert = nflexvert
|
|
_m.ngeom = ngeom
|
|
_m.nmaxmeshdeg = nmaxmeshdeg
|
|
_m.nmaxpolygon = nmaxpolygon
|
|
_m.nmeshface = nmeshface
|
|
_m.nxn_geom_pair_filtered = nxn_geom_pair_filtered
|
|
_m.nxn_pairid = nxn_pairid
|
|
_m.nxn_pairid_filtered = nxn_pairid_filtered
|
|
_m.oct_aabb = oct_aabb
|
|
_m.oct_child = oct_child
|
|
_m.oct_coeff = oct_coeff
|
|
_m.opt.broadphase = opt__broadphase
|
|
_m.opt.broadphase_filter = opt__broadphase_filter
|
|
_m.opt.ccd_iterations = opt__ccd_iterations
|
|
_m.opt.ccd_tolerance = opt__ccd_tolerance
|
|
_m.opt.disableflags = opt__disableflags
|
|
_m.opt.enableflags = opt__enableflags
|
|
_m.opt.sdf_initpoints = opt__sdf_initpoints
|
|
_m.opt.sdf_iterations = opt__sdf_iterations
|
|
_m.pair_dim = pair_dim
|
|
_m.pair_friction = pair_friction
|
|
_m.pair_gap = pair_gap
|
|
_m.pair_margin = pair_margin
|
|
_m.pair_solimp = pair_solimp
|
|
_m.pair_solref = pair_solref
|
|
_m.pair_solreffriction = pair_solreffriction
|
|
_m.plugin = plugin
|
|
_m.plugin_attr = plugin_attr
|
|
_d.body_awake = body_awake
|
|
_d.contact.dim = contact__dim
|
|
_d.contact.dist = contact__dist
|
|
_d.contact.efc_address = contact__efc_address
|
|
_d.contact.flex = contact__flex
|
|
_d.contact.frame = contact__frame
|
|
_d.contact.friction = contact__friction
|
|
_d.contact.geom = contact__geom
|
|
_d.contact.geomcollisionid = contact__geomcollisionid
|
|
_d.contact.includemargin = contact__includemargin
|
|
_d.contact.pos = contact__pos
|
|
_d.contact.solimp = contact__solimp
|
|
_d.contact.solref = contact__solref
|
|
_d.contact.solreffriction = contact__solreffriction
|
|
_d.contact.type = contact__type
|
|
_d.contact.vert = contact__vert
|
|
_d.contact.worldid = contact__worldid
|
|
_d.flexvert_xpos = flexvert_xpos
|
|
_d.geom_xmat = geom_xmat
|
|
_d.geom_xpos = geom_xpos
|
|
_d.naccdmax = naccdmax
|
|
_d.nacon = nacon
|
|
_d.naconmax = naconmax
|
|
_d.ncollision = ncollision
|
|
_d.nworld = nworld
|
|
mjwarp.collision(_m, _d)
|
|
|
|
|
|
def _collision_jax_impl(m: types.Model, d: types.Data):
|
|
output_dims = {
|
|
'nacon': d._impl.nacon.shape,
|
|
'ncollision': d._impl.ncollision.shape,
|
|
'contact__dim': d._impl.contact__dim.shape,
|
|
'contact__dist': d._impl.contact__dist.shape,
|
|
'contact__efc_address': d._impl.contact__efc_address.shape,
|
|
'contact__flex': d._impl.contact__flex.shape,
|
|
'contact__frame': d._impl.contact__frame.shape,
|
|
'contact__friction': d._impl.contact__friction.shape,
|
|
'contact__geom': d._impl.contact__geom.shape,
|
|
'contact__geomcollisionid': d._impl.contact__geomcollisionid.shape,
|
|
'contact__includemargin': d._impl.contact__includemargin.shape,
|
|
'contact__pos': d._impl.contact__pos.shape,
|
|
'contact__solimp': d._impl.contact__solimp.shape,
|
|
'contact__solref': d._impl.contact__solref.shape,
|
|
'contact__solreffriction': d._impl.contact__solreffriction.shape,
|
|
'contact__type': d._impl.contact__type.shape,
|
|
'contact__vert': d._impl.contact__vert.shape,
|
|
'contact__worldid': d._impl.contact__worldid.shape,
|
|
}
|
|
jf = ffi.jax_callable_variadic_tuple(
|
|
_collision_shim,
|
|
num_outputs=18,
|
|
output_dims=output_dims,
|
|
vmap_method=None,
|
|
in_out_argnames=set([
|
|
'nacon',
|
|
'ncollision',
|
|
'contact__dim',
|
|
'contact__dist',
|
|
'contact__efc_address',
|
|
'contact__flex',
|
|
'contact__frame',
|
|
'contact__friction',
|
|
'contact__geom',
|
|
'contact__geomcollisionid',
|
|
'contact__includemargin',
|
|
'contact__pos',
|
|
'contact__solimp',
|
|
'contact__solref',
|
|
'contact__solreffriction',
|
|
'contact__type',
|
|
'contact__vert',
|
|
'contact__worldid',
|
|
]),
|
|
stage_in_argnames=set([
|
|
'geom_aabb',
|
|
'geom_friction',
|
|
'geom_gap',
|
|
'geom_margin',
|
|
'geom_rbound',
|
|
'geom_size',
|
|
'geom_solimp',
|
|
'geom_solmix',
|
|
'geom_solref',
|
|
'geom_xmat',
|
|
'geom_xpos',
|
|
'hfield_data',
|
|
'pair_friction',
|
|
'pair_gap',
|
|
'pair_margin',
|
|
'pair_solimp',
|
|
'pair_solref',
|
|
'pair_solreffriction',
|
|
]),
|
|
stage_out_argnames=set([]),
|
|
graph_mode=m.opt._impl.graph_mode,
|
|
has_side_effect=False,
|
|
)
|
|
out = jf(
|
|
d.qpos.shape[0],
|
|
m._impl.block_dim,
|
|
m._impl.flex_conaffinity,
|
|
m._impl.flex_condim,
|
|
m._impl.flex_contype,
|
|
m._impl.flex_dim,
|
|
m._impl.flex_elem,
|
|
m._impl.flex_elemadr,
|
|
m._impl.flex_elemdataadr,
|
|
m._impl.flex_elemnum,
|
|
m._impl.flex_friction,
|
|
m._impl.flex_gap,
|
|
m._impl.flex_margin,
|
|
m._impl.flex_priority,
|
|
m._impl.flex_radius,
|
|
m._impl.flex_shell,
|
|
m._impl.flex_shelldataadr,
|
|
m._impl.flex_shellnum,
|
|
m._impl.flex_solimp,
|
|
m._impl.flex_solmix,
|
|
m._impl.flex_solref,
|
|
m.flex_vertadr,
|
|
m._impl.flex_vertflexid,
|
|
m.geom_aabb,
|
|
m.geom_bodyid,
|
|
m.geom_conaffinity,
|
|
m.geom_condim,
|
|
m.geom_contype,
|
|
jax.numpy.expand_dims(m.geom_dataid, 0),
|
|
m.geom_friction,
|
|
m.geom_gap,
|
|
m.geom_margin,
|
|
m._impl.geom_pair_type_count,
|
|
m._impl.geom_plugin_index,
|
|
m.geom_priority,
|
|
m.geom_rbound,
|
|
m.geom_size,
|
|
m.geom_solimp,
|
|
m.geom_solmix,
|
|
m.geom_solref,
|
|
m.geom_type,
|
|
m._impl.has_sdf_geom,
|
|
m.hfield_adr,
|
|
m.hfield_data,
|
|
m.hfield_ncol,
|
|
m.hfield_nrow,
|
|
m.hfield_size,
|
|
m.mesh_face,
|
|
m.mesh_faceadr,
|
|
m.mesh_graph,
|
|
m.mesh_graphadr,
|
|
m.mesh_octadr,
|
|
m._impl.mesh_polyadr,
|
|
m._impl.mesh_polymap,
|
|
m._impl.mesh_polymapadr,
|
|
m._impl.mesh_polymapnum,
|
|
m._impl.mesh_polynormal,
|
|
m._impl.mesh_polynum,
|
|
m._impl.mesh_polyvert,
|
|
m._impl.mesh_polyvertadr,
|
|
m._impl.mesh_polyvertnum,
|
|
m.mesh_vert,
|
|
m.mesh_vertadr,
|
|
m.mesh_vertnum,
|
|
m.nflex,
|
|
m._impl.nflexelem,
|
|
m._impl.nflexshelldata,
|
|
m._impl.nflexvert,
|
|
m.ngeom,
|
|
m._impl.nmaxmeshdeg,
|
|
m._impl.nmaxpolygon,
|
|
m.nmeshface,
|
|
m._impl.nxn_geom_pair_filtered,
|
|
m._impl.nxn_pairid,
|
|
m._impl.nxn_pairid_filtered,
|
|
m._impl.oct_aabb,
|
|
m._impl.oct_child,
|
|
m._impl.oct_coeff,
|
|
m.pair_dim,
|
|
m.pair_friction,
|
|
m.pair_gap,
|
|
m.pair_margin,
|
|
m.pair_solimp,
|
|
m.pair_solref,
|
|
m.pair_solreffriction,
|
|
m._impl.plugin,
|
|
m._impl.plugin_attr,
|
|
m.opt._impl.broadphase,
|
|
m.opt._impl.broadphase_filter,
|
|
m.opt._impl.ccd_iterations,
|
|
m.opt._impl.ccd_tolerance,
|
|
m.opt.disableflags,
|
|
m.opt.enableflags,
|
|
m.opt._impl.sdf_initpoints,
|
|
m.opt._impl.sdf_iterations,
|
|
d._impl.naccdmax,
|
|
d._impl.naconmax,
|
|
d._impl.body_awake,
|
|
d._impl.flexvert_xpos,
|
|
d.geom_xmat,
|
|
d.geom_xpos,
|
|
d._impl.nacon,
|
|
d._impl.ncollision,
|
|
d._impl.contact__dim,
|
|
d._impl.contact__dist,
|
|
d._impl.contact__efc_address,
|
|
d._impl.contact__flex,
|
|
d._impl.contact__frame,
|
|
d._impl.contact__friction,
|
|
d._impl.contact__geom,
|
|
d._impl.contact__geomcollisionid,
|
|
d._impl.contact__includemargin,
|
|
d._impl.contact__pos,
|
|
d._impl.contact__solimp,
|
|
d._impl.contact__solref,
|
|
d._impl.contact__solreffriction,
|
|
d._impl.contact__type,
|
|
d._impl.contact__vert,
|
|
d._impl.contact__worldid,
|
|
)
|
|
d = d.tree_replace({
|
|
'_impl.nacon': out[0],
|
|
'_impl.ncollision': out[1],
|
|
'_impl.contact__dim': out[2],
|
|
'_impl.contact__dist': out[3],
|
|
'_impl.contact__efc_address': out[4],
|
|
'_impl.contact__flex': out[5],
|
|
'_impl.contact__frame': out[6],
|
|
'_impl.contact__friction': out[7],
|
|
'_impl.contact__geom': out[8],
|
|
'_impl.contact__geomcollisionid': out[9],
|
|
'_impl.contact__includemargin': out[10],
|
|
'_impl.contact__pos': out[11],
|
|
'_impl.contact__solimp': out[12],
|
|
'_impl.contact__solref': out[13],
|
|
'_impl.contact__solreffriction': out[14],
|
|
'_impl.contact__type': out[15],
|
|
'_impl.contact__vert': out[16],
|
|
'_impl.contact__worldid': out[17],
|
|
})
|
|
return d
|
|
|
|
|
|
@jax.custom_batching.custom_vmap
|
|
@ffi.marshal_jax_warp_callable
|
|
def collision(m: types.Model, d: types.Data):
|
|
return _collision_jax_impl(m, d)
|
|
|
|
|
|
@collision.def_vmap
|
|
@ffi.marshal_custom_vmap
|
|
def collision_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
|
|
d = collision(m, d)
|
|
return d, is_batched[1]
|