Mustafa Hasekioglu
4608 lines
139 KiB
Python
4608 lines
139 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 _forward_shim(
|
|
# Model
|
|
nworld: int,
|
|
M_elemid: wp.array2d[int],
|
|
M_fullm_i: wp.array[int],
|
|
M_fullm_j: wp.array[int],
|
|
M_mulm_col: wp.array[int],
|
|
M_mulm_madr: wp.array[int],
|
|
M_mulm_rowadr: wp.array[int],
|
|
M_rowadr: wp.array[int],
|
|
M_rownnz: wp.array[int],
|
|
M_tiles: tuple[mjwp_types.TileSet, ...],
|
|
actuator_acc0: wp.array2d[float],
|
|
actuator_actadr: wp.array[int],
|
|
actuator_actearly: wp.array[bool],
|
|
actuator_actlimited: wp.array[bool],
|
|
actuator_actnum: wp.array[int],
|
|
actuator_actrange: wp.array2d[wp.vec2],
|
|
actuator_biasprm: wp.array2d[mjwp_types.vec10f],
|
|
actuator_biastype: wp.array[int],
|
|
actuator_cranklength: wp.array2d[float],
|
|
actuator_ctrllimited: wp.array[bool],
|
|
actuator_ctrlrange: wp.array2d[wp.vec2],
|
|
actuator_delay: wp.array[float],
|
|
actuator_dynprm: wp.array2d[mjwp_types.vec10f],
|
|
actuator_dyntype: wp.array[int],
|
|
actuator_forcelimited: wp.array[bool],
|
|
actuator_forcerange: wp.array2d[wp.vec2],
|
|
actuator_gainprm: wp.array2d[mjwp_types.vec10f],
|
|
actuator_gaintype: wp.array[int],
|
|
actuator_gear: wp.array2d[wp.spatial_vector],
|
|
actuator_history: wp.array[wp.vec2i],
|
|
actuator_historyadr: wp.array[int],
|
|
actuator_lengthrange: wp.array2d[wp.vec2],
|
|
actuator_trnid: wp.array[wp.vec2i],
|
|
actuator_trntype: wp.array[int],
|
|
actuator_trntype_body_adr: wp.array[int],
|
|
block_dim: mjwp_types.BlockDim,
|
|
body_branch_start: wp.array[int],
|
|
body_branches: wp.array[int],
|
|
body_dofadr: wp.array[int],
|
|
body_dofnum: wp.array[int],
|
|
body_fluid_ellipsoid: wp.array[bool],
|
|
body_geomadr: wp.array[int],
|
|
body_geomnum: wp.array[int],
|
|
body_gravcomp: wp.array2d[float],
|
|
body_inertia: wp.array2d[wp.vec3],
|
|
body_invweight0: wp.array2d[wp.vec2],
|
|
body_ipos: wp.array2d[wp.vec3],
|
|
body_iquat: wp.array2d[wp.quat],
|
|
body_isdofancestor: wp.array2d[int],
|
|
body_jntadr: wp.array[int],
|
|
body_jntnum: wp.array[int],
|
|
body_mass: wp.array2d[float],
|
|
body_mocapid: wp.array[int],
|
|
body_parentid: wp.array[int],
|
|
body_pos: wp.array2d[wp.vec3],
|
|
body_quat: wp.array2d[wp.quat],
|
|
body_rootid: wp.array[int],
|
|
body_subtreemass: wp.array2d[float],
|
|
body_tree: tuple[wp.array[int], ...],
|
|
body_treeid: wp.array[int],
|
|
body_weldid: wp.array[int],
|
|
cam_bodyid: wp.array[int],
|
|
cam_fovy: wp.array2d[float],
|
|
cam_intrinsic: wp.array2d[wp.vec4],
|
|
cam_mat0: wp.array2d[wp.mat33],
|
|
cam_mode: wp.array[int],
|
|
cam_pos: wp.array2d[wp.vec3],
|
|
cam_pos0: wp.array2d[wp.vec3],
|
|
cam_poscom0: wp.array2d[wp.vec3],
|
|
cam_quat: wp.array2d[wp.quat],
|
|
cam_resolution: wp.array[wp.vec2i],
|
|
cam_sensorsize: wp.array[wp.vec2],
|
|
cam_targetbodyid: wp.array[int],
|
|
dof_armature: wp.array2d[float],
|
|
dof_bodyid: wp.array[int],
|
|
dof_damping: wp.array2d[float],
|
|
dof_dampingpoly: wp.array2d[wp.vec2],
|
|
dof_frictionloss: wp.array2d[float],
|
|
dof_invweight0: wp.array2d[float],
|
|
dof_jntid: wp.array[int],
|
|
dof_length: wp.array[float],
|
|
dof_parentid: wp.array[int],
|
|
dof_solimp: wp.array2d[mjwp_types.vec5],
|
|
dof_solref: wp.array2d[wp.vec2],
|
|
dof_treeid: wp.array[int],
|
|
dof_tri_col: wp.array[int],
|
|
dof_tri_row: wp.array[int],
|
|
eq_connect_adr: wp.array[int],
|
|
eq_data: wp.array2d[mjwp_types.vec11],
|
|
eq_flex_adr: wp.array[int],
|
|
eq_jnt_adr: wp.array[int],
|
|
eq_obj1id: wp.array[int],
|
|
eq_obj2id: wp.array[int],
|
|
eq_objtype: wp.array[int],
|
|
eq_solimp: wp.array2d[mjwp_types.vec5],
|
|
eq_solref: wp.array2d[wp.vec2],
|
|
eq_ten_adr: wp.array[int],
|
|
eq_type: wp.array[int],
|
|
eq_wld_adr: wp.array[int],
|
|
flex_bending: wp.array[float],
|
|
flex_bendingadr: wp.array[int],
|
|
flex_centered: wp.array[bool],
|
|
flex_conaffinity: wp.array[int],
|
|
flex_condim: wp.array[int],
|
|
flex_contype: wp.array[int],
|
|
flex_damping: wp.array[float],
|
|
flex_dim: wp.array[int],
|
|
flex_edge: wp.array[wp.vec2i],
|
|
flex_edgeadr: wp.array[int],
|
|
flex_edgeflap: wp.array[wp.vec2i],
|
|
flex_edgenum: wp.array[int],
|
|
flex_elem: wp.array[int],
|
|
flex_elemadr: wp.array[int],
|
|
flex_elemdataadr: wp.array[int],
|
|
flex_elemedge: wp.array[int],
|
|
flex_elemedgeadr: 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_stiffness: wp.array[float],
|
|
flex_stiffnessadr: wp.array[int],
|
|
flex_vert: wp.array[wp.vec3],
|
|
flex_vertadr: wp.array[int],
|
|
flex_vertbodyid: wp.array[int],
|
|
flex_vertflexid: wp.array[int],
|
|
flex_vertnum: wp.array[int],
|
|
flexedge_J_colind: wp.array[int],
|
|
flexedge_J_rowadr: wp.array[int],
|
|
flexedge_J_rownnz: wp.array[int],
|
|
flexedge_invweight0: wp.array[float],
|
|
flexedge_length0: wp.array[float],
|
|
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_fluid: wp.array2d[float],
|
|
geom_friction: wp.array2d[wp.vec3],
|
|
geom_gap: wp.array2d[float],
|
|
geom_group: wp.array[int],
|
|
geom_margin: wp.array2d[float],
|
|
geom_matid: wp.array2d[int],
|
|
geom_pair_type_count: tuple[int, ...],
|
|
geom_plugin_index: wp.array[int],
|
|
geom_pos: wp.array2d[wp.vec3],
|
|
geom_priority: wp.array[int],
|
|
geom_quat: wp.array2d[wp.quat],
|
|
geom_rbound: wp.array2d[float],
|
|
geom_rgba: wp.array2d[wp.vec4],
|
|
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_fluid: bool,
|
|
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],
|
|
is_sparse: bool,
|
|
jnt_actfrclimited: wp.array[bool],
|
|
jnt_actfrcrange: wp.array2d[wp.vec2],
|
|
jnt_actgravcomp: wp.array[int],
|
|
jnt_axis: wp.array2d[wp.vec3],
|
|
jnt_bodyid: wp.array[int],
|
|
jnt_dofadr: wp.array[int],
|
|
jnt_limited_ball_adr: wp.array[int],
|
|
jnt_limited_slide_hinge_adr: wp.array[int],
|
|
jnt_margin: wp.array2d[float],
|
|
jnt_pos: wp.array2d[wp.vec3],
|
|
jnt_qposadr: wp.array[int],
|
|
jnt_range: wp.array2d[wp.vec2],
|
|
jnt_solimp: wp.array2d[mjwp_types.vec5],
|
|
jnt_solref: wp.array2d[wp.vec2],
|
|
jnt_stiffness: wp.array2d[float],
|
|
jnt_stiffnesspoly: wp.array2d[wp.vec2],
|
|
jnt_type: wp.array[int],
|
|
light_bodyid: wp.array[int],
|
|
light_dir: wp.array2d[wp.vec3],
|
|
light_dir0: wp.array2d[wp.vec3],
|
|
light_mode: wp.array[int],
|
|
light_pos: wp.array2d[wp.vec3],
|
|
light_pos0: wp.array2d[wp.vec3],
|
|
light_poscom0: wp.array2d[wp.vec3],
|
|
light_targetbodyid: wp.array[int],
|
|
mat_rgba: wp.array2d[wp.vec4],
|
|
max_ten_J_rownnz: int,
|
|
mesh_face: wp.array[wp.vec3i],
|
|
mesh_faceadr: wp.array[int],
|
|
mesh_graph: wp.array[int],
|
|
mesh_graphadr: wp.array[int],
|
|
mesh_normal: wp.array[wp.vec3],
|
|
mesh_normaladr: wp.array[int],
|
|
mesh_normalnum: 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_quat: wp.array[wp.quat],
|
|
mesh_vert: wp.array[wp.vec3],
|
|
mesh_vertadr: wp.array[int],
|
|
mesh_vertnum: wp.array[int],
|
|
nJten: int,
|
|
na: int,
|
|
nacttrnbody: int,
|
|
nbody: int,
|
|
nbranch: int,
|
|
ncam: int,
|
|
neq: int,
|
|
nflex: int,
|
|
nflexedge: int,
|
|
nflexelem: int,
|
|
nflexshelldata: int,
|
|
nflexvert: int,
|
|
ngeom: int,
|
|
ngravcomp: int,
|
|
nhistory: int,
|
|
njnt: int,
|
|
nlight: int,
|
|
nmaxcondim: int,
|
|
nmaxmeshdeg: int,
|
|
nmaxpolygon: int,
|
|
nmaxpyramid: int,
|
|
nmeshface: int,
|
|
nrangefinder: int,
|
|
nsensorcollision: int,
|
|
nsensorcontact: int,
|
|
nsensortaxel: int,
|
|
nsite: int,
|
|
ntendon: int,
|
|
ntree: int,
|
|
nu: int,
|
|
nv: int,
|
|
nv_pad: int,
|
|
nwrap: 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],
|
|
qLD_all_updates: wp.array[wp.vec3i],
|
|
qLD_level_offsets: wp.array[int],
|
|
qLD_updates: tuple[wp.array[wp.vec3i], ...],
|
|
qpos0: wp.array2d[float],
|
|
qpos_spring: wp.array2d[float],
|
|
rangefinder_sensor_adr: wp.array[int],
|
|
sensor_acc_adr: wp.array[int],
|
|
sensor_adr: wp.array[int],
|
|
sensor_adr_to_contact_adr: wp.array[int],
|
|
sensor_contact_adr: wp.array[int],
|
|
sensor_cutoff: wp.array[float],
|
|
sensor_datatype: wp.array[int],
|
|
sensor_delay: wp.array[float],
|
|
sensor_dim: wp.array[int],
|
|
sensor_e_kinetic: bool,
|
|
sensor_e_potential: bool,
|
|
sensor_history: wp.array[wp.vec2i],
|
|
sensor_historyadr: wp.array[int],
|
|
sensor_interval: wp.array[wp.vec2],
|
|
sensor_intprm: wp.array2d[int],
|
|
sensor_limitfrc_adr: wp.array[int],
|
|
sensor_limitpos_adr: wp.array[int],
|
|
sensor_limitvel_adr: wp.array[int],
|
|
sensor_objid: wp.array[int],
|
|
sensor_objtype: wp.array[int],
|
|
sensor_pos_adr: wp.array[int],
|
|
sensor_rangefinder_adr: wp.array[int],
|
|
sensor_rangefinder_bodyid: wp.array[int],
|
|
sensor_refid: wp.array[int],
|
|
sensor_reftype: wp.array[int],
|
|
sensor_rne_postconstraint: bool,
|
|
sensor_subtree_vel: bool,
|
|
sensor_tendonactfrc_adr: wp.array[int],
|
|
sensor_touch_adr: wp.array[int],
|
|
sensor_type: wp.array[int],
|
|
sensor_vel_adr: wp.array[int],
|
|
site_bodyid: wp.array[int],
|
|
site_pos: wp.array2d[wp.vec3],
|
|
site_quat: wp.array2d[wp.quat],
|
|
site_size: wp.array[wp.vec3],
|
|
site_type: wp.array[int],
|
|
taxel_sensorid: wp.array[int],
|
|
taxel_vertadr: wp.array[int],
|
|
ten_J_colind: wp.array[int],
|
|
ten_J_rowadr: wp.array[int],
|
|
ten_J_rownnz: wp.array[int],
|
|
tendon_actfrclimited: wp.array[bool],
|
|
tendon_actfrcrange: wp.array2d[wp.vec2],
|
|
tendon_adr: wp.array[int],
|
|
tendon_armature: wp.array2d[float],
|
|
tendon_damping: wp.array2d[float],
|
|
tendon_dampingpoly: wp.array2d[wp.vec2],
|
|
tendon_frictionloss: wp.array2d[float],
|
|
tendon_geom_adr: wp.array[int],
|
|
tendon_invweight0: wp.array2d[float],
|
|
tendon_jnt_adr: wp.array[int],
|
|
tendon_length0: wp.array2d[float],
|
|
tendon_lengthspring: wp.array2d[wp.vec2],
|
|
tendon_limited: wp.array[int],
|
|
tendon_limited_adr: wp.array[int],
|
|
tendon_margin: wp.array2d[float],
|
|
tendon_num: wp.array[int],
|
|
tendon_range: wp.array2d[wp.vec2],
|
|
tendon_site_pair_adr: wp.array[int],
|
|
tendon_solimp_fri: wp.array2d[mjwp_types.vec5],
|
|
tendon_solimp_lim: wp.array2d[mjwp_types.vec5],
|
|
tendon_solref_fri: wp.array2d[wp.vec2],
|
|
tendon_solref_lim: wp.array2d[wp.vec2],
|
|
tendon_stiffness: wp.array2d[float],
|
|
tendon_stiffnesspoly: wp.array2d[wp.vec2],
|
|
tree_dofadr: wp.array[int],
|
|
tree_dofnum: wp.array[int],
|
|
tree_sleep_policy: wp.array[int],
|
|
wrap_geom_adr: wp.array[int],
|
|
wrap_jnt_adr: wp.array[int],
|
|
wrap_objid: wp.array[int],
|
|
wrap_prm: wp.array[float],
|
|
wrap_pulley_scale: wp.array[float],
|
|
wrap_site_pair_adr: wp.array[int],
|
|
wrap_type: wp.array[int],
|
|
opt__broadphase: int,
|
|
opt__broadphase_filter: int,
|
|
opt__ccd_iterations: int,
|
|
opt__ccd_tolerance: wp.array[float],
|
|
opt__cone: int,
|
|
opt__contact_sensor_maxmatch: int,
|
|
opt__density: wp.array[float],
|
|
opt__disableflags: int,
|
|
opt__enableflags: int,
|
|
opt__graph_conditional: bool,
|
|
opt__gravity: wp.array[wp.vec3],
|
|
opt__impratio_invsqrt: wp.array[float],
|
|
opt__integrator: int,
|
|
opt__iterations: int,
|
|
opt__ls_iterations: int,
|
|
opt__ls_parallel: bool,
|
|
opt__ls_parallel_min_step: float,
|
|
opt__ls_tolerance: wp.array[float],
|
|
opt__magnetic: wp.array[wp.vec3],
|
|
opt__run_collision_detection: bool,
|
|
opt__sdf_initpoints: int,
|
|
opt__sdf_iterations: int,
|
|
opt__solver: int,
|
|
opt__timestep: wp.array[float],
|
|
opt__tolerance: wp.array[float],
|
|
opt__viscosity: wp.array[float],
|
|
opt__wind: wp.array[wp.vec3],
|
|
stat__meaninertia: wp.array[float],
|
|
# Data
|
|
naccdmax: int,
|
|
naconmax: int,
|
|
njmax: int,
|
|
njmax_nnz: int,
|
|
M: wp.array3d[float],
|
|
act: wp.array2d[float],
|
|
act_dot: wp.array2d[float],
|
|
actuator_force: wp.array2d[float],
|
|
actuator_length: wp.array2d[float],
|
|
actuator_moment: wp.array2d[float],
|
|
actuator_velocity: wp.array2d[float],
|
|
body_awake: wp.array2d[int],
|
|
body_awake_ind: wp.array2d[int],
|
|
cacc: wp.array2d[wp.spatial_vector],
|
|
cam_xmat: wp.array2d[wp.mat33],
|
|
cam_xpos: wp.array2d[wp.vec3],
|
|
cdof: wp.array2d[wp.spatial_vector],
|
|
cdof_dot: wp.array2d[wp.spatial_vector],
|
|
cfrc_ext: wp.array2d[wp.spatial_vector],
|
|
cfrc_int: wp.array2d[wp.spatial_vector],
|
|
cinert: wp.array2d[mjwp_types.vec10],
|
|
crb: wp.array2d[mjwp_types.vec10],
|
|
ctrl: wp.array2d[float],
|
|
cvel: wp.array2d[wp.spatial_vector],
|
|
dof_awake_ind: wp.array2d[int],
|
|
dof_island: wp.array2d[int],
|
|
dof_islandid: wp.array2d[int],
|
|
efc_islandid: wp.array2d[int],
|
|
energy: wp.array[wp.vec2],
|
|
eq_active: wp.array2d[bool],
|
|
flexedge_J: wp.array2d[float],
|
|
flexedge_length: wp.array2d[float],
|
|
flexedge_velocity: wp.array2d[float],
|
|
flexvert_xpos: wp.array2d[wp.vec3],
|
|
geom_xmat: wp.array2d[wp.mat33],
|
|
geom_xpos: wp.array2d[wp.vec3],
|
|
history: wp.array2d[float],
|
|
iqacc: wp.array2d[float],
|
|
iqacc_smooth: wp.array2d[float],
|
|
iqfrc_constraint: wp.array2d[float],
|
|
iqfrc_smooth: wp.array2d[float],
|
|
island_dofadr: wp.array2d[int],
|
|
island_efcadr: wp.array2d[int],
|
|
island_idofadr: wp.array2d[int],
|
|
island_ne: wp.array2d[int],
|
|
island_nefc: wp.array2d[int],
|
|
island_nf: wp.array2d[int],
|
|
island_nv: wp.array2d[int],
|
|
light_xdir: wp.array2d[wp.vec3],
|
|
light_xpos: wp.array2d[wp.vec3],
|
|
map_dof2idof: wp.array2d[int],
|
|
map_efc2iefc: wp.array2d[int],
|
|
map_idof2dof: wp.array2d[int],
|
|
map_iefc2efc: wp.array2d[int],
|
|
mocap_pos: wp.array2d[wp.vec3],
|
|
mocap_quat: wp.array2d[wp.quat],
|
|
moment_colind: wp.array2d[int],
|
|
moment_rowadr: wp.array2d[int],
|
|
moment_rownnz: wp.array2d[int],
|
|
nacon: wp.array[int],
|
|
nbody_awake: wp.array[int],
|
|
ncollision: wp.array[int],
|
|
ne: wp.array[int],
|
|
nefc: wp.array[int],
|
|
nf: wp.array[int],
|
|
nidof: wp.array[int],
|
|
nisland: wp.array[int],
|
|
nl: wp.array[int],
|
|
ntree_awake: wp.array[int],
|
|
nv_awake: wp.array[int],
|
|
qLD: wp.array3d[float],
|
|
qLDiagInv: wp.array2d[float],
|
|
qacc: wp.array2d[float],
|
|
qacc_smooth: wp.array2d[float],
|
|
qacc_warmstart: wp.array2d[float],
|
|
qfrc_actuator: wp.array2d[float],
|
|
qfrc_applied: wp.array2d[float],
|
|
qfrc_bias: wp.array2d[float],
|
|
qfrc_constraint: wp.array2d[float],
|
|
qfrc_damper: wp.array2d[float],
|
|
qfrc_fluid: wp.array2d[float],
|
|
qfrc_gravcomp: wp.array2d[float],
|
|
qfrc_passive: wp.array2d[float],
|
|
qfrc_smooth: wp.array2d[float],
|
|
qfrc_spring: wp.array2d[float],
|
|
qpos: wp.array2d[float],
|
|
qvel: wp.array2d[float],
|
|
sensordata: wp.array2d[float],
|
|
site_xmat: wp.array2d[wp.mat33],
|
|
site_xpos: wp.array2d[wp.vec3],
|
|
solver_niter: wp.array[int],
|
|
subtree_angmom: wp.array2d[wp.vec3],
|
|
subtree_com: wp.array2d[wp.vec3],
|
|
subtree_linvel: wp.array2d[wp.vec3],
|
|
ten_J: wp.array2d[float],
|
|
ten_length: wp.array2d[float],
|
|
ten_velocity: wp.array2d[float],
|
|
ten_wrapadr: wp.array2d[int],
|
|
ten_wrapnum: wp.array2d[int],
|
|
time: wp.array[float],
|
|
tree_asleep: wp.array2d[int],
|
|
tree_awake: wp.array2d[int],
|
|
tree_island: wp.array2d[int],
|
|
wrap_obj: wp.array2d[wp.vec2i],
|
|
wrap_xpos: wp.array2d[wp.spatial_vector],
|
|
xanchor: wp.array2d[wp.vec3],
|
|
xaxis: wp.array2d[wp.vec3],
|
|
xfrc_applied: wp.array2d[wp.spatial_vector],
|
|
ximat: wp.array2d[wp.mat33],
|
|
xipos: wp.array2d[wp.vec3],
|
|
xmat: wp.array2d[wp.mat33],
|
|
xpos: wp.array2d[wp.vec3],
|
|
xquat: wp.array2d[wp.quat],
|
|
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],
|
|
efc__D: wp.array2d[float],
|
|
efc__J: wp.array3d[float],
|
|
efc__J_colind: wp.array3d[int],
|
|
efc__J_rowadr: wp.array2d[int],
|
|
efc__J_rownnz: wp.array2d[int],
|
|
efc__Jqvel: wp.array2d[float],
|
|
efc__Ma: wp.array2d[float],
|
|
efc__aref: wp.array2d[float],
|
|
efc__force: wp.array2d[float],
|
|
efc__frictionloss: wp.array2d[float],
|
|
efc__iD: wp.array2d[float],
|
|
efc__iJ: wp.array3d[float],
|
|
efc__iJ_colind: wp.array3d[int],
|
|
efc__iJ_rowadr: wp.array2d[int],
|
|
efc__iJ_rownnz: wp.array2d[int],
|
|
efc__iaref: wp.array2d[float],
|
|
efc__id: wp.array2d[int],
|
|
efc__iforce: wp.array2d[float],
|
|
efc__ifrictionloss: wp.array2d[float],
|
|
efc__iid: wp.array2d[int],
|
|
efc__island: wp.array2d[int],
|
|
efc__istate: wp.array2d[int],
|
|
efc__itype: wp.array2d[int],
|
|
efc__margin: wp.array2d[float],
|
|
efc__pos: wp.array2d[float],
|
|
efc__state: wp.array2d[int],
|
|
efc__type: wp.array2d[int],
|
|
efc__vel: wp.array2d[float],
|
|
):
|
|
_m.stat = _s
|
|
_m.opt = _o
|
|
_m.callback = _cb
|
|
_d.efc = _e
|
|
_d.contact = _c
|
|
_m.M_elemid = M_elemid
|
|
_m.M_fullm_i = M_fullm_i
|
|
_m.M_fullm_j = M_fullm_j
|
|
_m.M_mulm_col = M_mulm_col
|
|
_m.M_mulm_madr = M_mulm_madr
|
|
_m.M_mulm_rowadr = M_mulm_rowadr
|
|
_m.M_rowadr = M_rowadr
|
|
_m.M_rownnz = M_rownnz
|
|
_m.M_tiles = M_tiles
|
|
_m.actuator_acc0 = actuator_acc0
|
|
_m.actuator_actadr = actuator_actadr
|
|
_m.actuator_actearly = actuator_actearly
|
|
_m.actuator_actlimited = actuator_actlimited
|
|
_m.actuator_actnum = actuator_actnum
|
|
_m.actuator_actrange = actuator_actrange
|
|
_m.actuator_biasprm = actuator_biasprm
|
|
_m.actuator_biastype = actuator_biastype
|
|
_m.actuator_cranklength = actuator_cranklength
|
|
_m.actuator_ctrllimited = actuator_ctrllimited
|
|
_m.actuator_ctrlrange = actuator_ctrlrange
|
|
_m.actuator_delay = actuator_delay
|
|
_m.actuator_dynprm = actuator_dynprm
|
|
_m.actuator_dyntype = actuator_dyntype
|
|
_m.actuator_forcelimited = actuator_forcelimited
|
|
_m.actuator_forcerange = actuator_forcerange
|
|
_m.actuator_gainprm = actuator_gainprm
|
|
_m.actuator_gaintype = actuator_gaintype
|
|
_m.actuator_gear = actuator_gear
|
|
_m.actuator_history = actuator_history
|
|
_m.actuator_historyadr = actuator_historyadr
|
|
_m.actuator_lengthrange = actuator_lengthrange
|
|
_m.actuator_trnid = actuator_trnid
|
|
_m.actuator_trntype = actuator_trntype
|
|
_m.actuator_trntype_body_adr = actuator_trntype_body_adr
|
|
_m.block_dim = block_dim
|
|
_m.body_branch_start = body_branch_start
|
|
_m.body_branches = body_branches
|
|
_m.body_dofadr = body_dofadr
|
|
_m.body_dofnum = body_dofnum
|
|
_m.body_fluid_ellipsoid = body_fluid_ellipsoid
|
|
_m.body_geomadr = body_geomadr
|
|
_m.body_geomnum = body_geomnum
|
|
_m.body_gravcomp = body_gravcomp
|
|
_m.body_inertia = body_inertia
|
|
_m.body_invweight0 = body_invweight0
|
|
_m.body_ipos = body_ipos
|
|
_m.body_iquat = body_iquat
|
|
_m.body_isdofancestor = body_isdofancestor
|
|
_m.body_jntadr = body_jntadr
|
|
_m.body_jntnum = body_jntnum
|
|
_m.body_mass = body_mass
|
|
_m.body_mocapid = body_mocapid
|
|
_m.body_parentid = body_parentid
|
|
_m.body_pos = body_pos
|
|
_m.body_quat = body_quat
|
|
_m.body_rootid = body_rootid
|
|
_m.body_subtreemass = body_subtreemass
|
|
_m.body_tree = body_tree
|
|
_m.body_treeid = body_treeid
|
|
_m.body_weldid = body_weldid
|
|
_m.cam_bodyid = cam_bodyid
|
|
_m.cam_fovy = cam_fovy
|
|
_m.cam_intrinsic = cam_intrinsic
|
|
_m.cam_mat0 = cam_mat0
|
|
_m.cam_mode = cam_mode
|
|
_m.cam_pos = cam_pos
|
|
_m.cam_pos0 = cam_pos0
|
|
_m.cam_poscom0 = cam_poscom0
|
|
_m.cam_quat = cam_quat
|
|
_m.cam_resolution = cam_resolution
|
|
_m.cam_sensorsize = cam_sensorsize
|
|
_m.cam_targetbodyid = cam_targetbodyid
|
|
_m.dof_armature = dof_armature
|
|
_m.dof_bodyid = dof_bodyid
|
|
_m.dof_damping = dof_damping
|
|
_m.dof_dampingpoly = dof_dampingpoly
|
|
_m.dof_frictionloss = dof_frictionloss
|
|
_m.dof_invweight0 = dof_invweight0
|
|
_m.dof_jntid = dof_jntid
|
|
_m.dof_length = dof_length
|
|
_m.dof_parentid = dof_parentid
|
|
_m.dof_solimp = dof_solimp
|
|
_m.dof_solref = dof_solref
|
|
_m.dof_treeid = dof_treeid
|
|
_m.dof_tri_col = dof_tri_col
|
|
_m.dof_tri_row = dof_tri_row
|
|
_m.eq_connect_adr = eq_connect_adr
|
|
_m.eq_data = eq_data
|
|
_m.eq_flex_adr = eq_flex_adr
|
|
_m.eq_jnt_adr = eq_jnt_adr
|
|
_m.eq_obj1id = eq_obj1id
|
|
_m.eq_obj2id = eq_obj2id
|
|
_m.eq_objtype = eq_objtype
|
|
_m.eq_solimp = eq_solimp
|
|
_m.eq_solref = eq_solref
|
|
_m.eq_ten_adr = eq_ten_adr
|
|
_m.eq_type = eq_type
|
|
_m.eq_wld_adr = eq_wld_adr
|
|
_m.flex_bending = flex_bending
|
|
_m.flex_bendingadr = flex_bendingadr
|
|
_m.flex_centered = flex_centered
|
|
_m.flex_conaffinity = flex_conaffinity
|
|
_m.flex_condim = flex_condim
|
|
_m.flex_contype = flex_contype
|
|
_m.flex_damping = flex_damping
|
|
_m.flex_dim = flex_dim
|
|
_m.flex_edge = flex_edge
|
|
_m.flex_edgeadr = flex_edgeadr
|
|
_m.flex_edgeflap = flex_edgeflap
|
|
_m.flex_edgenum = flex_edgenum
|
|
_m.flex_elem = flex_elem
|
|
_m.flex_elemadr = flex_elemadr
|
|
_m.flex_elemdataadr = flex_elemdataadr
|
|
_m.flex_elemedge = flex_elemedge
|
|
_m.flex_elemedgeadr = flex_elemedgeadr
|
|
_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_stiffness = flex_stiffness
|
|
_m.flex_stiffnessadr = flex_stiffnessadr
|
|
_m.flex_vert = flex_vert
|
|
_m.flex_vertadr = flex_vertadr
|
|
_m.flex_vertbodyid = flex_vertbodyid
|
|
_m.flex_vertflexid = flex_vertflexid
|
|
_m.flex_vertnum = flex_vertnum
|
|
_m.flexedge_J_colind = flexedge_J_colind
|
|
_m.flexedge_J_rowadr = flexedge_J_rowadr
|
|
_m.flexedge_J_rownnz = flexedge_J_rownnz
|
|
_m.flexedge_invweight0 = flexedge_invweight0
|
|
_m.flexedge_length0 = flexedge_length0
|
|
_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_fluid = geom_fluid
|
|
_m.geom_friction = geom_friction
|
|
_m.geom_gap = geom_gap
|
|
_m.geom_group = geom_group
|
|
_m.geom_margin = geom_margin
|
|
_m.geom_matid = geom_matid
|
|
_m.geom_pair_type_count = geom_pair_type_count
|
|
_m.geom_plugin_index = geom_plugin_index
|
|
_m.geom_pos = geom_pos
|
|
_m.geom_priority = geom_priority
|
|
_m.geom_quat = geom_quat
|
|
_m.geom_rbound = geom_rbound
|
|
_m.geom_rgba = geom_rgba
|
|
_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_fluid = has_fluid
|
|
_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.is_sparse = is_sparse
|
|
_m.jnt_actfrclimited = jnt_actfrclimited
|
|
_m.jnt_actfrcrange = jnt_actfrcrange
|
|
_m.jnt_actgravcomp = jnt_actgravcomp
|
|
_m.jnt_axis = jnt_axis
|
|
_m.jnt_bodyid = jnt_bodyid
|
|
_m.jnt_dofadr = jnt_dofadr
|
|
_m.jnt_limited_ball_adr = jnt_limited_ball_adr
|
|
_m.jnt_limited_slide_hinge_adr = jnt_limited_slide_hinge_adr
|
|
_m.jnt_margin = jnt_margin
|
|
_m.jnt_pos = jnt_pos
|
|
_m.jnt_qposadr = jnt_qposadr
|
|
_m.jnt_range = jnt_range
|
|
_m.jnt_solimp = jnt_solimp
|
|
_m.jnt_solref = jnt_solref
|
|
_m.jnt_stiffness = jnt_stiffness
|
|
_m.jnt_stiffnesspoly = jnt_stiffnesspoly
|
|
_m.jnt_type = jnt_type
|
|
_m.light_bodyid = light_bodyid
|
|
_m.light_dir = light_dir
|
|
_m.light_dir0 = light_dir0
|
|
_m.light_mode = light_mode
|
|
_m.light_pos = light_pos
|
|
_m.light_pos0 = light_pos0
|
|
_m.light_poscom0 = light_poscom0
|
|
_m.light_targetbodyid = light_targetbodyid
|
|
_m.mat_rgba = mat_rgba
|
|
_m.max_ten_J_rownnz = max_ten_J_rownnz
|
|
_m.mesh_face = mesh_face
|
|
_m.mesh_faceadr = mesh_faceadr
|
|
_m.mesh_graph = mesh_graph
|
|
_m.mesh_graphadr = mesh_graphadr
|
|
_m.mesh_normal = mesh_normal
|
|
_m.mesh_normaladr = mesh_normaladr
|
|
_m.mesh_normalnum = mesh_normalnum
|
|
_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_quat = mesh_quat
|
|
_m.mesh_vert = mesh_vert
|
|
_m.mesh_vertadr = mesh_vertadr
|
|
_m.mesh_vertnum = mesh_vertnum
|
|
_m.nJten = nJten
|
|
_m.na = na
|
|
_m.nacttrnbody = nacttrnbody
|
|
_m.nbody = nbody
|
|
_m.nbranch = nbranch
|
|
_m.ncam = ncam
|
|
_m.neq = neq
|
|
_m.nflex = nflex
|
|
_m.nflexedge = nflexedge
|
|
_m.nflexelem = nflexelem
|
|
_m.nflexshelldata = nflexshelldata
|
|
_m.nflexvert = nflexvert
|
|
_m.ngeom = ngeom
|
|
_m.ngravcomp = ngravcomp
|
|
_m.nhistory = nhistory
|
|
_m.njnt = njnt
|
|
_m.nlight = nlight
|
|
_m.nmaxcondim = nmaxcondim
|
|
_m.nmaxmeshdeg = nmaxmeshdeg
|
|
_m.nmaxpolygon = nmaxpolygon
|
|
_m.nmaxpyramid = nmaxpyramid
|
|
_m.nmeshface = nmeshface
|
|
_m.nrangefinder = nrangefinder
|
|
_m.nsensorcollision = nsensorcollision
|
|
_m.nsensorcontact = nsensorcontact
|
|
_m.nsensortaxel = nsensortaxel
|
|
_m.nsite = nsite
|
|
_m.ntendon = ntendon
|
|
_m.ntree = ntree
|
|
_m.nu = nu
|
|
_m.nv = nv
|
|
_m.nv_pad = nv_pad
|
|
_m.nwrap = nwrap
|
|
_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.cone = opt__cone
|
|
_m.opt.contact_sensor_maxmatch = opt__contact_sensor_maxmatch
|
|
_m.opt.density = opt__density
|
|
_m.opt.disableflags = opt__disableflags
|
|
_m.opt.enableflags = opt__enableflags
|
|
_m.opt.graph_conditional = opt__graph_conditional
|
|
_m.opt.gravity = opt__gravity
|
|
_m.opt.impratio_invsqrt = opt__impratio_invsqrt
|
|
_m.opt.integrator = opt__integrator
|
|
_m.opt.iterations = opt__iterations
|
|
_m.opt.ls_iterations = opt__ls_iterations
|
|
_m.opt.ls_parallel = opt__ls_parallel
|
|
_m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
|
|
_m.opt.ls_tolerance = opt__ls_tolerance
|
|
_m.opt.magnetic = opt__magnetic
|
|
_m.opt.run_collision_detection = opt__run_collision_detection
|
|
_m.opt.sdf_initpoints = opt__sdf_initpoints
|
|
_m.opt.sdf_iterations = opt__sdf_iterations
|
|
_m.opt.solver = opt__solver
|
|
_m.opt.timestep = opt__timestep
|
|
_m.opt.tolerance = opt__tolerance
|
|
_m.opt.viscosity = opt__viscosity
|
|
_m.opt.wind = opt__wind
|
|
_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
|
|
_m.qLD_all_updates = qLD_all_updates
|
|
_m.qLD_level_offsets = qLD_level_offsets
|
|
_m.qLD_updates = qLD_updates
|
|
_m.qpos0 = qpos0
|
|
_m.qpos_spring = qpos_spring
|
|
_m.rangefinder_sensor_adr = rangefinder_sensor_adr
|
|
_m.sensor_acc_adr = sensor_acc_adr
|
|
_m.sensor_adr = sensor_adr
|
|
_m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr
|
|
_m.sensor_contact_adr = sensor_contact_adr
|
|
_m.sensor_cutoff = sensor_cutoff
|
|
_m.sensor_datatype = sensor_datatype
|
|
_m.sensor_delay = sensor_delay
|
|
_m.sensor_dim = sensor_dim
|
|
_m.sensor_e_kinetic = sensor_e_kinetic
|
|
_m.sensor_e_potential = sensor_e_potential
|
|
_m.sensor_history = sensor_history
|
|
_m.sensor_historyadr = sensor_historyadr
|
|
_m.sensor_interval = sensor_interval
|
|
_m.sensor_intprm = sensor_intprm
|
|
_m.sensor_limitfrc_adr = sensor_limitfrc_adr
|
|
_m.sensor_limitpos_adr = sensor_limitpos_adr
|
|
_m.sensor_limitvel_adr = sensor_limitvel_adr
|
|
_m.sensor_objid = sensor_objid
|
|
_m.sensor_objtype = sensor_objtype
|
|
_m.sensor_pos_adr = sensor_pos_adr
|
|
_m.sensor_rangefinder_adr = sensor_rangefinder_adr
|
|
_m.sensor_rangefinder_bodyid = sensor_rangefinder_bodyid
|
|
_m.sensor_refid = sensor_refid
|
|
_m.sensor_reftype = sensor_reftype
|
|
_m.sensor_rne_postconstraint = sensor_rne_postconstraint
|
|
_m.sensor_subtree_vel = sensor_subtree_vel
|
|
_m.sensor_tendonactfrc_adr = sensor_tendonactfrc_adr
|
|
_m.sensor_touch_adr = sensor_touch_adr
|
|
_m.sensor_type = sensor_type
|
|
_m.sensor_vel_adr = sensor_vel_adr
|
|
_m.site_bodyid = site_bodyid
|
|
_m.site_pos = site_pos
|
|
_m.site_quat = site_quat
|
|
_m.site_size = site_size
|
|
_m.site_type = site_type
|
|
_m.stat.meaninertia = stat__meaninertia
|
|
_m.taxel_sensorid = taxel_sensorid
|
|
_m.taxel_vertadr = taxel_vertadr
|
|
_m.ten_J_colind = ten_J_colind
|
|
_m.ten_J_rowadr = ten_J_rowadr
|
|
_m.ten_J_rownnz = ten_J_rownnz
|
|
_m.tendon_actfrclimited = tendon_actfrclimited
|
|
_m.tendon_actfrcrange = tendon_actfrcrange
|
|
_m.tendon_adr = tendon_adr
|
|
_m.tendon_armature = tendon_armature
|
|
_m.tendon_damping = tendon_damping
|
|
_m.tendon_dampingpoly = tendon_dampingpoly
|
|
_m.tendon_frictionloss = tendon_frictionloss
|
|
_m.tendon_geom_adr = tendon_geom_adr
|
|
_m.tendon_invweight0 = tendon_invweight0
|
|
_m.tendon_jnt_adr = tendon_jnt_adr
|
|
_m.tendon_length0 = tendon_length0
|
|
_m.tendon_lengthspring = tendon_lengthspring
|
|
_m.tendon_limited = tendon_limited
|
|
_m.tendon_limited_adr = tendon_limited_adr
|
|
_m.tendon_margin = tendon_margin
|
|
_m.tendon_num = tendon_num
|
|
_m.tendon_range = tendon_range
|
|
_m.tendon_site_pair_adr = tendon_site_pair_adr
|
|
_m.tendon_solimp_fri = tendon_solimp_fri
|
|
_m.tendon_solimp_lim = tendon_solimp_lim
|
|
_m.tendon_solref_fri = tendon_solref_fri
|
|
_m.tendon_solref_lim = tendon_solref_lim
|
|
_m.tendon_stiffness = tendon_stiffness
|
|
_m.tendon_stiffnesspoly = tendon_stiffnesspoly
|
|
_m.tree_dofadr = tree_dofadr
|
|
_m.tree_dofnum = tree_dofnum
|
|
_m.tree_sleep_policy = tree_sleep_policy
|
|
_m.wrap_geom_adr = wrap_geom_adr
|
|
_m.wrap_jnt_adr = wrap_jnt_adr
|
|
_m.wrap_objid = wrap_objid
|
|
_m.wrap_prm = wrap_prm
|
|
_m.wrap_pulley_scale = wrap_pulley_scale
|
|
_m.wrap_site_pair_adr = wrap_site_pair_adr
|
|
_m.wrap_type = wrap_type
|
|
_d.M = M
|
|
_d.act = act
|
|
_d.act_dot = act_dot
|
|
_d.actuator_force = actuator_force
|
|
_d.actuator_length = actuator_length
|
|
_d.actuator_moment = actuator_moment
|
|
_d.actuator_velocity = actuator_velocity
|
|
_d.body_awake = body_awake
|
|
_d.body_awake_ind = body_awake_ind
|
|
_d.cacc = cacc
|
|
_d.cam_xmat = cam_xmat
|
|
_d.cam_xpos = cam_xpos
|
|
_d.cdof = cdof
|
|
_d.cdof_dot = cdof_dot
|
|
_d.cfrc_ext = cfrc_ext
|
|
_d.cfrc_int = cfrc_int
|
|
_d.cinert = cinert
|
|
_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.crb = crb
|
|
_d.ctrl = ctrl
|
|
_d.cvel = cvel
|
|
_d.dof_awake_ind = dof_awake_ind
|
|
_d.dof_island = dof_island
|
|
_d.dof_islandid = dof_islandid
|
|
_d.efc.D = efc__D
|
|
_d.efc.J = efc__J
|
|
_d.efc.J_colind = efc__J_colind
|
|
_d.efc.J_rowadr = efc__J_rowadr
|
|
_d.efc.J_rownnz = efc__J_rownnz
|
|
_d.efc.Jqvel = efc__Jqvel
|
|
_d.efc.Ma = efc__Ma
|
|
_d.efc.aref = efc__aref
|
|
_d.efc.force = efc__force
|
|
_d.efc.frictionloss = efc__frictionloss
|
|
_d.efc.iD = efc__iD
|
|
_d.efc.iJ = efc__iJ
|
|
_d.efc.iJ_colind = efc__iJ_colind
|
|
_d.efc.iJ_rowadr = efc__iJ_rowadr
|
|
_d.efc.iJ_rownnz = efc__iJ_rownnz
|
|
_d.efc.iaref = efc__iaref
|
|
_d.efc.id = efc__id
|
|
_d.efc.iforce = efc__iforce
|
|
_d.efc.ifrictionloss = efc__ifrictionloss
|
|
_d.efc.iid = efc__iid
|
|
_d.efc.island = efc__island
|
|
_d.efc.istate = efc__istate
|
|
_d.efc.itype = efc__itype
|
|
_d.efc.margin = efc__margin
|
|
_d.efc.pos = efc__pos
|
|
_d.efc.state = efc__state
|
|
_d.efc.type = efc__type
|
|
_d.efc.vel = efc__vel
|
|
_d.efc_islandid = efc_islandid
|
|
_d.energy = energy
|
|
_d.eq_active = eq_active
|
|
_d.flexedge_J = flexedge_J
|
|
_d.flexedge_length = flexedge_length
|
|
_d.flexedge_velocity = flexedge_velocity
|
|
_d.flexvert_xpos = flexvert_xpos
|
|
_d.geom_xmat = geom_xmat
|
|
_d.geom_xpos = geom_xpos
|
|
_d.history = history
|
|
_d.iqacc = iqacc
|
|
_d.iqacc_smooth = iqacc_smooth
|
|
_d.iqfrc_constraint = iqfrc_constraint
|
|
_d.iqfrc_smooth = iqfrc_smooth
|
|
_d.island_dofadr = island_dofadr
|
|
_d.island_efcadr = island_efcadr
|
|
_d.island_idofadr = island_idofadr
|
|
_d.island_ne = island_ne
|
|
_d.island_nefc = island_nefc
|
|
_d.island_nf = island_nf
|
|
_d.island_nv = island_nv
|
|
_d.light_xdir = light_xdir
|
|
_d.light_xpos = light_xpos
|
|
_d.map_dof2idof = map_dof2idof
|
|
_d.map_efc2iefc = map_efc2iefc
|
|
_d.map_idof2dof = map_idof2dof
|
|
_d.map_iefc2efc = map_iefc2efc
|
|
_d.mocap_pos = mocap_pos
|
|
_d.mocap_quat = mocap_quat
|
|
_d.moment_colind = moment_colind
|
|
_d.moment_rowadr = moment_rowadr
|
|
_d.moment_rownnz = moment_rownnz
|
|
_d.naccdmax = naccdmax
|
|
_d.nacon = nacon
|
|
_d.naconmax = naconmax
|
|
_d.nbody_awake = nbody_awake
|
|
_d.ncollision = ncollision
|
|
_d.ne = ne
|
|
_d.nefc = nefc
|
|
_d.nf = nf
|
|
_d.nidof = nidof
|
|
_d.nisland = nisland
|
|
_d.njmax = njmax
|
|
_d.njmax_nnz = njmax_nnz
|
|
_d.nl = nl
|
|
_d.ntree_awake = ntree_awake
|
|
_d.nv_awake = nv_awake
|
|
_d.qLD = qLD
|
|
_d.qLDiagInv = qLDiagInv
|
|
_d.qacc = qacc
|
|
_d.qacc_smooth = qacc_smooth
|
|
_d.qacc_warmstart = qacc_warmstart
|
|
_d.qfrc_actuator = qfrc_actuator
|
|
_d.qfrc_applied = qfrc_applied
|
|
_d.qfrc_bias = qfrc_bias
|
|
_d.qfrc_constraint = qfrc_constraint
|
|
_d.qfrc_damper = qfrc_damper
|
|
_d.qfrc_fluid = qfrc_fluid
|
|
_d.qfrc_gravcomp = qfrc_gravcomp
|
|
_d.qfrc_passive = qfrc_passive
|
|
_d.qfrc_smooth = qfrc_smooth
|
|
_d.qfrc_spring = qfrc_spring
|
|
_d.qpos = qpos
|
|
_d.qvel = qvel
|
|
_d.sensordata = sensordata
|
|
_d.site_xmat = site_xmat
|
|
_d.site_xpos = site_xpos
|
|
_d.solver_niter = solver_niter
|
|
_d.subtree_angmom = subtree_angmom
|
|
_d.subtree_com = subtree_com
|
|
_d.subtree_linvel = subtree_linvel
|
|
_d.ten_J = ten_J
|
|
_d.ten_length = ten_length
|
|
_d.ten_velocity = ten_velocity
|
|
_d.ten_wrapadr = ten_wrapadr
|
|
_d.ten_wrapnum = ten_wrapnum
|
|
_d.time = time
|
|
_d.tree_asleep = tree_asleep
|
|
_d.tree_awake = tree_awake
|
|
_d.tree_island = tree_island
|
|
_d.wrap_obj = wrap_obj
|
|
_d.wrap_xpos = wrap_xpos
|
|
_d.xanchor = xanchor
|
|
_d.xaxis = xaxis
|
|
_d.xfrc_applied = xfrc_applied
|
|
_d.ximat = ximat
|
|
_d.xipos = xipos
|
|
_d.xmat = xmat
|
|
_d.xpos = xpos
|
|
_d.xquat = xquat
|
|
_d.nworld = nworld
|
|
mjwarp.forward(_m, _d)
|
|
|
|
|
|
def _forward_jax_impl(m: types.Model, d: types.Data):
|
|
output_dims = {
|
|
'M': d._impl.M.shape,
|
|
'act_dot': d.act_dot.shape,
|
|
'actuator_force': d.actuator_force.shape,
|
|
'actuator_length': d.actuator_length.shape,
|
|
'actuator_moment': d._impl.actuator_moment.shape,
|
|
'actuator_velocity': d._impl.actuator_velocity.shape,
|
|
'body_awake': d._impl.body_awake.shape,
|
|
'body_awake_ind': d._impl.body_awake_ind.shape,
|
|
'cacc': d._impl.cacc.shape,
|
|
'cam_xmat': d.cam_xmat.shape,
|
|
'cam_xpos': d.cam_xpos.shape,
|
|
'cdof': d.cdof.shape,
|
|
'cdof_dot': d.cdof_dot.shape,
|
|
'cfrc_ext': d._impl.cfrc_ext.shape,
|
|
'cfrc_int': d._impl.cfrc_int.shape,
|
|
'cinert': d._impl.cinert.shape,
|
|
'crb': d._impl.crb.shape,
|
|
'cvel': d.cvel.shape,
|
|
'dof_awake_ind': d._impl.dof_awake_ind.shape,
|
|
'dof_island': d._impl.dof_island.shape,
|
|
'dof_islandid': d._impl.dof_islandid.shape,
|
|
'efc_islandid': d._impl.efc_islandid.shape,
|
|
'energy': d._impl.energy.shape,
|
|
'flexedge_J': d._impl.flexedge_J.shape,
|
|
'flexedge_length': d._impl.flexedge_length.shape,
|
|
'flexedge_velocity': d._impl.flexedge_velocity.shape,
|
|
'flexvert_xpos': d._impl.flexvert_xpos.shape,
|
|
'geom_xmat': d.geom_xmat.shape,
|
|
'geom_xpos': d.geom_xpos.shape,
|
|
'history': d.history.shape,
|
|
'iqacc': d._impl.iqacc.shape,
|
|
'iqacc_smooth': d._impl.iqacc_smooth.shape,
|
|
'iqfrc_constraint': d._impl.iqfrc_constraint.shape,
|
|
'iqfrc_smooth': d._impl.iqfrc_smooth.shape,
|
|
'island_dofadr': d._impl.island_dofadr.shape,
|
|
'island_efcadr': d._impl.island_efcadr.shape,
|
|
'island_idofadr': d._impl.island_idofadr.shape,
|
|
'island_ne': d._impl.island_ne.shape,
|
|
'island_nefc': d._impl.island_nefc.shape,
|
|
'island_nf': d._impl.island_nf.shape,
|
|
'island_nv': d._impl.island_nv.shape,
|
|
'light_xdir': d._impl.light_xdir.shape,
|
|
'light_xpos': d._impl.light_xpos.shape,
|
|
'map_dof2idof': d._impl.map_dof2idof.shape,
|
|
'map_efc2iefc': d._impl.map_efc2iefc.shape,
|
|
'map_idof2dof': d._impl.map_idof2dof.shape,
|
|
'map_iefc2efc': d._impl.map_iefc2efc.shape,
|
|
'moment_colind': d._impl.moment_colind.shape,
|
|
'moment_rowadr': d._impl.moment_rowadr.shape,
|
|
'moment_rownnz': d._impl.moment_rownnz.shape,
|
|
'nacon': d._impl.nacon.shape,
|
|
'nbody_awake': d._impl.nbody_awake.shape,
|
|
'ncollision': d._impl.ncollision.shape,
|
|
'ne': d._impl.ne.shape,
|
|
'nefc': d._impl.nefc.shape,
|
|
'nf': d._impl.nf.shape,
|
|
'nidof': d._impl.nidof.shape,
|
|
'nisland': d._impl.nisland.shape,
|
|
'nl': d._impl.nl.shape,
|
|
'ntree_awake': d._impl.ntree_awake.shape,
|
|
'nv_awake': d._impl.nv_awake.shape,
|
|
'qLD': d._impl.qLD.shape,
|
|
'qLDiagInv': d._impl.qLDiagInv.shape,
|
|
'qacc': d.qacc.shape,
|
|
'qacc_smooth': d.qacc_smooth.shape,
|
|
'qfrc_actuator': d.qfrc_actuator.shape,
|
|
'qfrc_bias': d.qfrc_bias.shape,
|
|
'qfrc_constraint': d.qfrc_constraint.shape,
|
|
'qfrc_damper': d._impl.qfrc_damper.shape,
|
|
'qfrc_fluid': d.qfrc_fluid.shape,
|
|
'qfrc_gravcomp': d.qfrc_gravcomp.shape,
|
|
'qfrc_passive': d.qfrc_passive.shape,
|
|
'qfrc_smooth': d.qfrc_smooth.shape,
|
|
'qfrc_spring': d._impl.qfrc_spring.shape,
|
|
'qvel': d.qvel.shape,
|
|
'sensordata': d.sensordata.shape,
|
|
'site_xmat': d.site_xmat.shape,
|
|
'site_xpos': d.site_xpos.shape,
|
|
'solver_niter': d._impl.solver_niter.shape,
|
|
'subtree_angmom': d._impl.subtree_angmom.shape,
|
|
'subtree_com': d.subtree_com.shape,
|
|
'subtree_linvel': d._impl.subtree_linvel.shape,
|
|
'ten_J': d._impl.ten_J.shape,
|
|
'ten_length': d.ten_length.shape,
|
|
'ten_velocity': d._impl.ten_velocity.shape,
|
|
'ten_wrapadr': d._impl.ten_wrapadr.shape,
|
|
'ten_wrapnum': d._impl.ten_wrapnum.shape,
|
|
'tree_asleep': d._impl.tree_asleep.shape,
|
|
'tree_awake': d._impl.tree_awake.shape,
|
|
'tree_island': d._impl.tree_island.shape,
|
|
'wrap_obj': d._impl.wrap_obj.shape,
|
|
'wrap_xpos': d._impl.wrap_xpos.shape,
|
|
'xanchor': d.xanchor.shape,
|
|
'xaxis': d.xaxis.shape,
|
|
'ximat': d.ximat.shape,
|
|
'xipos': d.xipos.shape,
|
|
'xmat': d.xmat.shape,
|
|
'xpos': d.xpos.shape,
|
|
'xquat': d.xquat.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,
|
|
'efc__D': d._impl.efc__D.shape,
|
|
'efc__J': d._impl.efc__J.shape,
|
|
'efc__J_colind': d._impl.efc__J_colind.shape,
|
|
'efc__J_rowadr': d._impl.efc__J_rowadr.shape,
|
|
'efc__J_rownnz': d._impl.efc__J_rownnz.shape,
|
|
'efc__Jqvel': d._impl.efc__Jqvel.shape,
|
|
'efc__Ma': d._impl.efc__Ma.shape,
|
|
'efc__aref': d._impl.efc__aref.shape,
|
|
'efc__force': d._impl.efc__force.shape,
|
|
'efc__frictionloss': d._impl.efc__frictionloss.shape,
|
|
'efc__iD': d._impl.efc__iD.shape,
|
|
'efc__iJ': d._impl.efc__iJ.shape,
|
|
'efc__iJ_colind': d._impl.efc__iJ_colind.shape,
|
|
'efc__iJ_rowadr': d._impl.efc__iJ_rowadr.shape,
|
|
'efc__iJ_rownnz': d._impl.efc__iJ_rownnz.shape,
|
|
'efc__iaref': d._impl.efc__iaref.shape,
|
|
'efc__id': d._impl.efc__id.shape,
|
|
'efc__iforce': d._impl.efc__iforce.shape,
|
|
'efc__ifrictionloss': d._impl.efc__ifrictionloss.shape,
|
|
'efc__iid': d._impl.efc__iid.shape,
|
|
'efc__island': d._impl.efc__island.shape,
|
|
'efc__istate': d._impl.efc__istate.shape,
|
|
'efc__itype': d._impl.efc__itype.shape,
|
|
'efc__margin': d._impl.efc__margin.shape,
|
|
'efc__pos': d._impl.efc__pos.shape,
|
|
'efc__state': d._impl.efc__state.shape,
|
|
'efc__type': d._impl.efc__type.shape,
|
|
'efc__vel': d._impl.efc__vel.shape,
|
|
}
|
|
jf = ffi.jax_callable_variadic_tuple(
|
|
_forward_shim,
|
|
num_outputs=143,
|
|
output_dims=output_dims,
|
|
vmap_method=None,
|
|
in_out_argnames=set([
|
|
'M',
|
|
'act_dot',
|
|
'actuator_force',
|
|
'actuator_length',
|
|
'actuator_moment',
|
|
'actuator_velocity',
|
|
'body_awake',
|
|
'body_awake_ind',
|
|
'cacc',
|
|
'cam_xmat',
|
|
'cam_xpos',
|
|
'cdof',
|
|
'cdof_dot',
|
|
'cfrc_ext',
|
|
'cfrc_int',
|
|
'cinert',
|
|
'crb',
|
|
'cvel',
|
|
'dof_awake_ind',
|
|
'dof_island',
|
|
'dof_islandid',
|
|
'efc_islandid',
|
|
'energy',
|
|
'flexedge_J',
|
|
'flexedge_length',
|
|
'flexedge_velocity',
|
|
'flexvert_xpos',
|
|
'geom_xmat',
|
|
'geom_xpos',
|
|
'history',
|
|
'iqacc',
|
|
'iqacc_smooth',
|
|
'iqfrc_constraint',
|
|
'iqfrc_smooth',
|
|
'island_dofadr',
|
|
'island_efcadr',
|
|
'island_idofadr',
|
|
'island_ne',
|
|
'island_nefc',
|
|
'island_nf',
|
|
'island_nv',
|
|
'light_xdir',
|
|
'light_xpos',
|
|
'map_dof2idof',
|
|
'map_efc2iefc',
|
|
'map_idof2dof',
|
|
'map_iefc2efc',
|
|
'moment_colind',
|
|
'moment_rowadr',
|
|
'moment_rownnz',
|
|
'nacon',
|
|
'nbody_awake',
|
|
'ncollision',
|
|
'ne',
|
|
'nefc',
|
|
'nf',
|
|
'nidof',
|
|
'nisland',
|
|
'nl',
|
|
'ntree_awake',
|
|
'nv_awake',
|
|
'qLD',
|
|
'qLDiagInv',
|
|
'qacc',
|
|
'qacc_smooth',
|
|
'qfrc_actuator',
|
|
'qfrc_bias',
|
|
'qfrc_constraint',
|
|
'qfrc_damper',
|
|
'qfrc_fluid',
|
|
'qfrc_gravcomp',
|
|
'qfrc_passive',
|
|
'qfrc_smooth',
|
|
'qfrc_spring',
|
|
'qvel',
|
|
'sensordata',
|
|
'site_xmat',
|
|
'site_xpos',
|
|
'solver_niter',
|
|
'subtree_angmom',
|
|
'subtree_com',
|
|
'subtree_linvel',
|
|
'ten_J',
|
|
'ten_length',
|
|
'ten_velocity',
|
|
'ten_wrapadr',
|
|
'ten_wrapnum',
|
|
'tree_asleep',
|
|
'tree_awake',
|
|
'tree_island',
|
|
'wrap_obj',
|
|
'wrap_xpos',
|
|
'xanchor',
|
|
'xaxis',
|
|
'ximat',
|
|
'xipos',
|
|
'xmat',
|
|
'xpos',
|
|
'xquat',
|
|
'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',
|
|
'efc__D',
|
|
'efc__J',
|
|
'efc__J_colind',
|
|
'efc__J_rowadr',
|
|
'efc__J_rownnz',
|
|
'efc__Jqvel',
|
|
'efc__Ma',
|
|
'efc__aref',
|
|
'efc__force',
|
|
'efc__frictionloss',
|
|
'efc__iD',
|
|
'efc__iJ',
|
|
'efc__iJ_colind',
|
|
'efc__iJ_rowadr',
|
|
'efc__iJ_rownnz',
|
|
'efc__iaref',
|
|
'efc__id',
|
|
'efc__iforce',
|
|
'efc__ifrictionloss',
|
|
'efc__iid',
|
|
'efc__island',
|
|
'efc__istate',
|
|
'efc__itype',
|
|
'efc__margin',
|
|
'efc__pos',
|
|
'efc__state',
|
|
'efc__type',
|
|
'efc__vel',
|
|
]),
|
|
stage_in_argnames=set([
|
|
'act',
|
|
'act_dot',
|
|
'actuator_acc0',
|
|
'actuator_actrange',
|
|
'actuator_biasprm',
|
|
'actuator_cranklength',
|
|
'actuator_ctrlrange',
|
|
'actuator_dynprm',
|
|
'actuator_force',
|
|
'actuator_forcerange',
|
|
'actuator_gainprm',
|
|
'actuator_gear',
|
|
'actuator_length',
|
|
'actuator_lengthrange',
|
|
'body_gravcomp',
|
|
'body_inertia',
|
|
'body_invweight0',
|
|
'body_ipos',
|
|
'body_iquat',
|
|
'body_mass',
|
|
'body_pos',
|
|
'body_quat',
|
|
'body_subtreemass',
|
|
'cam_fovy',
|
|
'cam_intrinsic',
|
|
'cam_mat0',
|
|
'cam_pos',
|
|
'cam_pos0',
|
|
'cam_poscom0',
|
|
'cam_quat',
|
|
'cam_xmat',
|
|
'cam_xpos',
|
|
'cdof',
|
|
'cdof_dot',
|
|
'ctrl',
|
|
'cvel',
|
|
'dof_armature',
|
|
'dof_damping',
|
|
'dof_dampingpoly',
|
|
'dof_frictionloss',
|
|
'dof_invweight0',
|
|
'dof_solimp',
|
|
'dof_solref',
|
|
'eq_active',
|
|
'eq_data',
|
|
'eq_solimp',
|
|
'eq_solref',
|
|
'geom_aabb',
|
|
'geom_friction',
|
|
'geom_gap',
|
|
'geom_margin',
|
|
'geom_matid',
|
|
'geom_pos',
|
|
'geom_quat',
|
|
'geom_rbound',
|
|
'geom_rgba',
|
|
'geom_size',
|
|
'geom_solimp',
|
|
'geom_solmix',
|
|
'geom_solref',
|
|
'geom_xmat',
|
|
'geom_xpos',
|
|
'hfield_data',
|
|
'history',
|
|
'jnt_actfrcrange',
|
|
'jnt_axis',
|
|
'jnt_margin',
|
|
'jnt_pos',
|
|
'jnt_range',
|
|
'jnt_solimp',
|
|
'jnt_solref',
|
|
'jnt_stiffness',
|
|
'jnt_stiffnesspoly',
|
|
'light_dir',
|
|
'light_dir0',
|
|
'light_pos',
|
|
'light_pos0',
|
|
'light_poscom0',
|
|
'mat_rgba',
|
|
'mocap_pos',
|
|
'mocap_quat',
|
|
'opt__density',
|
|
'opt__gravity',
|
|
'opt__ls_tolerance',
|
|
'opt__magnetic',
|
|
'opt__timestep',
|
|
'opt__tolerance',
|
|
'opt__viscosity',
|
|
'opt__wind',
|
|
'pair_friction',
|
|
'pair_gap',
|
|
'pair_margin',
|
|
'pair_solimp',
|
|
'pair_solref',
|
|
'pair_solreffriction',
|
|
'qacc',
|
|
'qacc_smooth',
|
|
'qacc_warmstart',
|
|
'qfrc_actuator',
|
|
'qfrc_applied',
|
|
'qfrc_bias',
|
|
'qfrc_constraint',
|
|
'qfrc_fluid',
|
|
'qfrc_gravcomp',
|
|
'qfrc_passive',
|
|
'qfrc_smooth',
|
|
'qpos',
|
|
'qpos0',
|
|
'qpos_spring',
|
|
'qvel',
|
|
'sensordata',
|
|
'site_pos',
|
|
'site_quat',
|
|
'site_xmat',
|
|
'site_xpos',
|
|
'subtree_com',
|
|
'ten_length',
|
|
'tendon_actfrcrange',
|
|
'tendon_armature',
|
|
'tendon_damping',
|
|
'tendon_dampingpoly',
|
|
'tendon_frictionloss',
|
|
'tendon_invweight0',
|
|
'tendon_length0',
|
|
'tendon_lengthspring',
|
|
'tendon_margin',
|
|
'tendon_range',
|
|
'tendon_solimp_fri',
|
|
'tendon_solimp_lim',
|
|
'tendon_solref_fri',
|
|
'tendon_solref_lim',
|
|
'tendon_stiffness',
|
|
'tendon_stiffnesspoly',
|
|
'time',
|
|
'xanchor',
|
|
'xaxis',
|
|
'xfrc_applied',
|
|
'ximat',
|
|
'xipos',
|
|
'xmat',
|
|
'xpos',
|
|
'xquat',
|
|
]),
|
|
stage_out_argnames=set([
|
|
'act_dot',
|
|
'actuator_force',
|
|
'actuator_length',
|
|
'cam_xmat',
|
|
'cam_xpos',
|
|
'cdof',
|
|
'cdof_dot',
|
|
'cvel',
|
|
'geom_xmat',
|
|
'geom_xpos',
|
|
'history',
|
|
'qacc',
|
|
'qacc_smooth',
|
|
'qfrc_actuator',
|
|
'qfrc_bias',
|
|
'qfrc_constraint',
|
|
'qfrc_fluid',
|
|
'qfrc_gravcomp',
|
|
'qfrc_passive',
|
|
'qfrc_smooth',
|
|
'qvel',
|
|
'sensordata',
|
|
'site_xmat',
|
|
'site_xpos',
|
|
'subtree_com',
|
|
'ten_length',
|
|
'xanchor',
|
|
'xaxis',
|
|
'ximat',
|
|
'xipos',
|
|
'xmat',
|
|
'xpos',
|
|
'xquat',
|
|
]),
|
|
graph_mode=m.opt._impl.graph_mode,
|
|
has_side_effect=False,
|
|
)
|
|
out = jf(
|
|
d.qpos.shape[0],
|
|
m._impl.M_elemid,
|
|
m._impl.M_fullm_i,
|
|
m._impl.M_fullm_j,
|
|
m._impl.M_mulm_col,
|
|
m._impl.M_mulm_madr,
|
|
m._impl.M_mulm_rowadr,
|
|
m._impl.M_rowadr,
|
|
m._impl.M_rownnz,
|
|
m._impl.M_tiles,
|
|
m.actuator_acc0,
|
|
m.actuator_actadr,
|
|
m.actuator_actearly,
|
|
m.actuator_actlimited,
|
|
m.actuator_actnum,
|
|
m.actuator_actrange,
|
|
m.actuator_biasprm,
|
|
m.actuator_biastype,
|
|
m.actuator_cranklength,
|
|
m.actuator_ctrllimited,
|
|
m.actuator_ctrlrange,
|
|
m._impl.actuator_delay,
|
|
m.actuator_dynprm,
|
|
m.actuator_dyntype,
|
|
m.actuator_forcelimited,
|
|
m.actuator_forcerange,
|
|
m.actuator_gainprm,
|
|
m.actuator_gaintype,
|
|
m.actuator_gear,
|
|
m._impl.actuator_history,
|
|
m._impl.actuator_historyadr,
|
|
m.actuator_lengthrange,
|
|
m.actuator_trnid,
|
|
m.actuator_trntype,
|
|
m._impl.actuator_trntype_body_adr,
|
|
m._impl.block_dim,
|
|
m._impl.body_branch_start,
|
|
m._impl.body_branches,
|
|
m.body_dofadr,
|
|
m.body_dofnum,
|
|
m._impl.body_fluid_ellipsoid,
|
|
m.body_geomadr,
|
|
m.body_geomnum,
|
|
m.body_gravcomp,
|
|
m.body_inertia,
|
|
m.body_invweight0,
|
|
m.body_ipos,
|
|
m.body_iquat,
|
|
m._impl.body_isdofancestor,
|
|
m.body_jntadr,
|
|
m.body_jntnum,
|
|
m.body_mass,
|
|
m.body_mocapid,
|
|
m.body_parentid,
|
|
m.body_pos,
|
|
m.body_quat,
|
|
m.body_rootid,
|
|
m.body_subtreemass,
|
|
m._impl.body_tree,
|
|
m.body_treeid,
|
|
m.body_weldid,
|
|
m.cam_bodyid,
|
|
m.cam_fovy,
|
|
m.cam_intrinsic,
|
|
m.cam_mat0,
|
|
m.cam_mode,
|
|
m.cam_pos,
|
|
m.cam_pos0,
|
|
m.cam_poscom0,
|
|
m.cam_quat,
|
|
m.cam_resolution,
|
|
m.cam_sensorsize,
|
|
m.cam_targetbodyid,
|
|
m.dof_armature,
|
|
m.dof_bodyid,
|
|
m.dof_damping,
|
|
m.dof_dampingpoly,
|
|
m.dof_frictionloss,
|
|
m.dof_invweight0,
|
|
m.dof_jntid,
|
|
m._impl.dof_length,
|
|
m.dof_parentid,
|
|
m.dof_solimp,
|
|
m.dof_solref,
|
|
m.dof_treeid,
|
|
m._impl.dof_tri_col,
|
|
m._impl.dof_tri_row,
|
|
m._impl.eq_connect_adr,
|
|
m.eq_data,
|
|
m._impl.eq_flex_adr,
|
|
m._impl.eq_jnt_adr,
|
|
m.eq_obj1id,
|
|
m.eq_obj2id,
|
|
m.eq_objtype,
|
|
m.eq_solimp,
|
|
m.eq_solref,
|
|
m._impl.eq_ten_adr,
|
|
m.eq_type,
|
|
m._impl.eq_wld_adr,
|
|
m._impl.flex_bending,
|
|
m._impl.flex_bendingadr,
|
|
m._impl.flex_centered,
|
|
m._impl.flex_conaffinity,
|
|
m._impl.flex_condim,
|
|
m._impl.flex_contype,
|
|
m._impl.flex_damping,
|
|
m._impl.flex_dim,
|
|
m._impl.flex_edge,
|
|
m._impl.flex_edgeadr,
|
|
m._impl.flex_edgeflap,
|
|
m._impl.flex_edgenum,
|
|
m._impl.flex_elem,
|
|
m._impl.flex_elemadr,
|
|
m._impl.flex_elemdataadr,
|
|
m._impl.flex_elemedge,
|
|
m._impl.flex_elemedgeadr,
|
|
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._impl.flex_stiffness,
|
|
m._impl.flex_stiffnessadr,
|
|
m._impl.flex_vert,
|
|
m.flex_vertadr,
|
|
m._impl.flex_vertbodyid,
|
|
m._impl.flex_vertflexid,
|
|
m.flex_vertnum,
|
|
m._impl.flexedge_J_colind,
|
|
m._impl.flexedge_J_rowadr,
|
|
m._impl.flexedge_J_rownnz,
|
|
m._impl.flexedge_invweight0,
|
|
m._impl.flexedge_length0,
|
|
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_fluid,
|
|
m.geom_friction,
|
|
m.geom_gap,
|
|
m.geom_group,
|
|
m.geom_margin,
|
|
m.geom_matid,
|
|
m._impl.geom_pair_type_count,
|
|
m._impl.geom_plugin_index,
|
|
m.geom_pos,
|
|
m.geom_priority,
|
|
m.geom_quat,
|
|
m.geom_rbound,
|
|
m.geom_rgba,
|
|
m.geom_size,
|
|
m.geom_solimp,
|
|
m.geom_solmix,
|
|
m.geom_solref,
|
|
m.geom_type,
|
|
m._impl.has_fluid,
|
|
m._impl.has_sdf_geom,
|
|
m.hfield_adr,
|
|
m.hfield_data,
|
|
m.hfield_ncol,
|
|
m.hfield_nrow,
|
|
m.hfield_size,
|
|
m._impl.is_sparse,
|
|
m.jnt_actfrclimited,
|
|
m.jnt_actfrcrange,
|
|
m.jnt_actgravcomp,
|
|
m.jnt_axis,
|
|
m.jnt_bodyid,
|
|
m.jnt_dofadr,
|
|
m._impl.jnt_limited_ball_adr,
|
|
m._impl.jnt_limited_slide_hinge_adr,
|
|
m.jnt_margin,
|
|
m.jnt_pos,
|
|
m.jnt_qposadr,
|
|
m.jnt_range,
|
|
m.jnt_solimp,
|
|
m.jnt_solref,
|
|
m.jnt_stiffness,
|
|
m.jnt_stiffnesspoly,
|
|
m.jnt_type,
|
|
m._impl.light_bodyid,
|
|
m.light_dir,
|
|
m.light_dir0,
|
|
m.light_mode,
|
|
m.light_pos,
|
|
m.light_pos0,
|
|
m.light_poscom0,
|
|
m._impl.light_targetbodyid,
|
|
m.mat_rgba,
|
|
m._impl.max_ten_J_rownnz,
|
|
m.mesh_face,
|
|
m.mesh_faceadr,
|
|
m.mesh_graph,
|
|
m.mesh_graphadr,
|
|
m.mesh_normal,
|
|
m.mesh_normaladr,
|
|
m.mesh_normalnum,
|
|
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_quat,
|
|
m.mesh_vert,
|
|
m.mesh_vertadr,
|
|
m.mesh_vertnum,
|
|
m.nJten,
|
|
m.na,
|
|
m._impl.nacttrnbody,
|
|
m.nbody,
|
|
m._impl.nbranch,
|
|
m.ncam,
|
|
m.neq,
|
|
m.nflex,
|
|
m._impl.nflexedge,
|
|
m._impl.nflexelem,
|
|
m._impl.nflexshelldata,
|
|
m._impl.nflexvert,
|
|
m.ngeom,
|
|
m.ngravcomp,
|
|
m.nhistory,
|
|
m.njnt,
|
|
m.nlight,
|
|
m._impl.nmaxcondim,
|
|
m._impl.nmaxmeshdeg,
|
|
m._impl.nmaxpolygon,
|
|
m._impl.nmaxpyramid,
|
|
m.nmeshface,
|
|
m._impl.nrangefinder,
|
|
m._impl.nsensorcollision,
|
|
m._impl.nsensorcontact,
|
|
m._impl.nsensortaxel,
|
|
m.nsite,
|
|
m.ntendon,
|
|
m._impl.ntree,
|
|
m.nu,
|
|
m.nv,
|
|
m._impl.nv_pad,
|
|
m.nwrap,
|
|
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._impl.qLD_all_updates,
|
|
m._impl.qLD_level_offsets,
|
|
m._impl.qLD_updates,
|
|
m.qpos0,
|
|
m.qpos_spring,
|
|
m._impl.rangefinder_sensor_adr,
|
|
m._impl.sensor_acc_adr,
|
|
m.sensor_adr,
|
|
m._impl.sensor_adr_to_contact_adr,
|
|
m._impl.sensor_contact_adr,
|
|
m.sensor_cutoff,
|
|
m.sensor_datatype,
|
|
m._impl.sensor_delay,
|
|
m.sensor_dim,
|
|
m._impl.sensor_e_kinetic,
|
|
m._impl.sensor_e_potential,
|
|
m._impl.sensor_history,
|
|
m._impl.sensor_historyadr,
|
|
m._impl.sensor_interval,
|
|
m.sensor_intprm,
|
|
m._impl.sensor_limitfrc_adr,
|
|
m._impl.sensor_limitpos_adr,
|
|
m._impl.sensor_limitvel_adr,
|
|
m.sensor_objid,
|
|
m.sensor_objtype,
|
|
m._impl.sensor_pos_adr,
|
|
m._impl.sensor_rangefinder_adr,
|
|
m._impl.sensor_rangefinder_bodyid,
|
|
m.sensor_refid,
|
|
m.sensor_reftype,
|
|
m._impl.sensor_rne_postconstraint,
|
|
m._impl.sensor_subtree_vel,
|
|
m._impl.sensor_tendonactfrc_adr,
|
|
m._impl.sensor_touch_adr,
|
|
m.sensor_type,
|
|
m._impl.sensor_vel_adr,
|
|
m.site_bodyid,
|
|
m.site_pos,
|
|
m.site_quat,
|
|
m.site_size,
|
|
m.site_type,
|
|
m._impl.taxel_sensorid,
|
|
m._impl.taxel_vertadr,
|
|
m._impl.ten_J_colind,
|
|
m._impl.ten_J_rowadr,
|
|
m._impl.ten_J_rownnz,
|
|
m.tendon_actfrclimited,
|
|
m.tendon_actfrcrange,
|
|
m.tendon_adr,
|
|
m.tendon_armature,
|
|
m.tendon_damping,
|
|
m.tendon_dampingpoly,
|
|
m.tendon_frictionloss,
|
|
m._impl.tendon_geom_adr,
|
|
m.tendon_invweight0,
|
|
m._impl.tendon_jnt_adr,
|
|
m.tendon_length0,
|
|
m.tendon_lengthspring,
|
|
m.tendon_limited,
|
|
m._impl.tendon_limited_adr,
|
|
m.tendon_margin,
|
|
m.tendon_num,
|
|
m.tendon_range,
|
|
m._impl.tendon_site_pair_adr,
|
|
m.tendon_solimp_fri,
|
|
m.tendon_solimp_lim,
|
|
m.tendon_solref_fri,
|
|
m.tendon_solref_lim,
|
|
m.tendon_stiffness,
|
|
m.tendon_stiffnesspoly,
|
|
m._impl.tree_dofadr,
|
|
m._impl.tree_dofnum,
|
|
m._impl.tree_sleep_policy,
|
|
m._impl.wrap_geom_adr,
|
|
m._impl.wrap_jnt_adr,
|
|
m.wrap_objid,
|
|
m.wrap_prm,
|
|
m._impl.wrap_pulley_scale,
|
|
m._impl.wrap_site_pair_adr,
|
|
m.wrap_type,
|
|
m.opt._impl.broadphase,
|
|
m.opt._impl.broadphase_filter,
|
|
m.opt._impl.ccd_iterations,
|
|
m.opt._impl.ccd_tolerance,
|
|
m.opt.cone,
|
|
m.opt._impl.contact_sensor_maxmatch,
|
|
m.opt.density,
|
|
m.opt.disableflags,
|
|
m.opt.enableflags,
|
|
m.opt._impl.graph_conditional,
|
|
m.opt.gravity,
|
|
m.opt._impl.impratio_invsqrt,
|
|
m.opt.integrator,
|
|
m.opt.iterations,
|
|
m.opt.ls_iterations,
|
|
m.opt._impl.ls_parallel,
|
|
m.opt._impl.ls_parallel_min_step,
|
|
m.opt.ls_tolerance,
|
|
m.opt.magnetic,
|
|
m.opt._impl.run_collision_detection,
|
|
m.opt._impl.sdf_initpoints,
|
|
m.opt._impl.sdf_iterations,
|
|
m.opt.solver,
|
|
m.opt.timestep,
|
|
m.opt.tolerance,
|
|
m.opt.viscosity,
|
|
m.opt.wind,
|
|
m.stat.meaninertia,
|
|
d._impl.naccdmax,
|
|
d._impl.naconmax,
|
|
d._impl.njmax,
|
|
d._impl.njmax_nnz,
|
|
d._impl.M,
|
|
d.act,
|
|
d.act_dot,
|
|
d.actuator_force,
|
|
d.actuator_length,
|
|
d._impl.actuator_moment,
|
|
d._impl.actuator_velocity,
|
|
d._impl.body_awake,
|
|
d._impl.body_awake_ind,
|
|
d._impl.cacc,
|
|
d.cam_xmat,
|
|
d.cam_xpos,
|
|
d.cdof,
|
|
d.cdof_dot,
|
|
d._impl.cfrc_ext,
|
|
d._impl.cfrc_int,
|
|
d._impl.cinert,
|
|
d._impl.crb,
|
|
d.ctrl,
|
|
d.cvel,
|
|
d._impl.dof_awake_ind,
|
|
d._impl.dof_island,
|
|
d._impl.dof_islandid,
|
|
d._impl.efc_islandid,
|
|
d._impl.energy,
|
|
d.eq_active,
|
|
d._impl.flexedge_J,
|
|
d._impl.flexedge_length,
|
|
d._impl.flexedge_velocity,
|
|
d._impl.flexvert_xpos,
|
|
d.geom_xmat,
|
|
d.geom_xpos,
|
|
d.history,
|
|
d._impl.iqacc,
|
|
d._impl.iqacc_smooth,
|
|
d._impl.iqfrc_constraint,
|
|
d._impl.iqfrc_smooth,
|
|
d._impl.island_dofadr,
|
|
d._impl.island_efcadr,
|
|
d._impl.island_idofadr,
|
|
d._impl.island_ne,
|
|
d._impl.island_nefc,
|
|
d._impl.island_nf,
|
|
d._impl.island_nv,
|
|
d._impl.light_xdir,
|
|
d._impl.light_xpos,
|
|
d._impl.map_dof2idof,
|
|
d._impl.map_efc2iefc,
|
|
d._impl.map_idof2dof,
|
|
d._impl.map_iefc2efc,
|
|
d.mocap_pos,
|
|
d.mocap_quat,
|
|
d._impl.moment_colind,
|
|
d._impl.moment_rowadr,
|
|
d._impl.moment_rownnz,
|
|
d._impl.nacon,
|
|
d._impl.nbody_awake,
|
|
d._impl.ncollision,
|
|
d._impl.ne,
|
|
d._impl.nefc,
|
|
d._impl.nf,
|
|
d._impl.nidof,
|
|
d._impl.nisland,
|
|
d._impl.nl,
|
|
d._impl.ntree_awake,
|
|
d._impl.nv_awake,
|
|
d._impl.qLD,
|
|
d._impl.qLDiagInv,
|
|
d.qacc,
|
|
d.qacc_smooth,
|
|
d.qacc_warmstart,
|
|
d.qfrc_actuator,
|
|
d.qfrc_applied,
|
|
d.qfrc_bias,
|
|
d.qfrc_constraint,
|
|
d._impl.qfrc_damper,
|
|
d.qfrc_fluid,
|
|
d.qfrc_gravcomp,
|
|
d.qfrc_passive,
|
|
d.qfrc_smooth,
|
|
d._impl.qfrc_spring,
|
|
d.qpos,
|
|
d.qvel,
|
|
d.sensordata,
|
|
d.site_xmat,
|
|
d.site_xpos,
|
|
d._impl.solver_niter,
|
|
d._impl.subtree_angmom,
|
|
d.subtree_com,
|
|
d._impl.subtree_linvel,
|
|
d._impl.ten_J,
|
|
d.ten_length,
|
|
d._impl.ten_velocity,
|
|
d._impl.ten_wrapadr,
|
|
d._impl.ten_wrapnum,
|
|
d.time,
|
|
d._impl.tree_asleep,
|
|
d._impl.tree_awake,
|
|
d._impl.tree_island,
|
|
d._impl.wrap_obj,
|
|
d._impl.wrap_xpos,
|
|
d.xanchor,
|
|
d.xaxis,
|
|
d.xfrc_applied,
|
|
d.ximat,
|
|
d.xipos,
|
|
d.xmat,
|
|
d.xpos,
|
|
d.xquat,
|
|
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._impl.efc__D,
|
|
d._impl.efc__J,
|
|
d._impl.efc__J_colind,
|
|
d._impl.efc__J_rowadr,
|
|
d._impl.efc__J_rownnz,
|
|
d._impl.efc__Jqvel,
|
|
d._impl.efc__Ma,
|
|
d._impl.efc__aref,
|
|
d._impl.efc__force,
|
|
d._impl.efc__frictionloss,
|
|
d._impl.efc__iD,
|
|
d._impl.efc__iJ,
|
|
d._impl.efc__iJ_colind,
|
|
d._impl.efc__iJ_rowadr,
|
|
d._impl.efc__iJ_rownnz,
|
|
d._impl.efc__iaref,
|
|
d._impl.efc__id,
|
|
d._impl.efc__iforce,
|
|
d._impl.efc__ifrictionloss,
|
|
d._impl.efc__iid,
|
|
d._impl.efc__island,
|
|
d._impl.efc__istate,
|
|
d._impl.efc__itype,
|
|
d._impl.efc__margin,
|
|
d._impl.efc__pos,
|
|
d._impl.efc__state,
|
|
d._impl.efc__type,
|
|
d._impl.efc__vel,
|
|
)
|
|
d = d.tree_replace({
|
|
'_impl.M': out[0],
|
|
'act_dot': out[1],
|
|
'actuator_force': out[2],
|
|
'actuator_length': out[3],
|
|
'_impl.actuator_moment': out[4],
|
|
'_impl.actuator_velocity': out[5],
|
|
'_impl.body_awake': out[6],
|
|
'_impl.body_awake_ind': out[7],
|
|
'_impl.cacc': out[8],
|
|
'cam_xmat': out[9],
|
|
'cam_xpos': out[10],
|
|
'cdof': out[11],
|
|
'cdof_dot': out[12],
|
|
'_impl.cfrc_ext': out[13],
|
|
'_impl.cfrc_int': out[14],
|
|
'_impl.cinert': out[15],
|
|
'_impl.crb': out[16],
|
|
'cvel': out[17],
|
|
'_impl.dof_awake_ind': out[18],
|
|
'_impl.dof_island': out[19],
|
|
'_impl.dof_islandid': out[20],
|
|
'_impl.efc_islandid': out[21],
|
|
'_impl.energy': out[22],
|
|
'_impl.flexedge_J': out[23],
|
|
'_impl.flexedge_length': out[24],
|
|
'_impl.flexedge_velocity': out[25],
|
|
'_impl.flexvert_xpos': out[26],
|
|
'geom_xmat': out[27],
|
|
'geom_xpos': out[28],
|
|
'history': out[29],
|
|
'_impl.iqacc': out[30],
|
|
'_impl.iqacc_smooth': out[31],
|
|
'_impl.iqfrc_constraint': out[32],
|
|
'_impl.iqfrc_smooth': out[33],
|
|
'_impl.island_dofadr': out[34],
|
|
'_impl.island_efcadr': out[35],
|
|
'_impl.island_idofadr': out[36],
|
|
'_impl.island_ne': out[37],
|
|
'_impl.island_nefc': out[38],
|
|
'_impl.island_nf': out[39],
|
|
'_impl.island_nv': out[40],
|
|
'_impl.light_xdir': out[41],
|
|
'_impl.light_xpos': out[42],
|
|
'_impl.map_dof2idof': out[43],
|
|
'_impl.map_efc2iefc': out[44],
|
|
'_impl.map_idof2dof': out[45],
|
|
'_impl.map_iefc2efc': out[46],
|
|
'_impl.moment_colind': out[47],
|
|
'_impl.moment_rowadr': out[48],
|
|
'_impl.moment_rownnz': out[49],
|
|
'_impl.nacon': out[50],
|
|
'_impl.nbody_awake': out[51],
|
|
'_impl.ncollision': out[52],
|
|
'_impl.ne': out[53],
|
|
'_impl.nefc': out[54],
|
|
'_impl.nf': out[55],
|
|
'_impl.nidof': out[56],
|
|
'_impl.nisland': out[57],
|
|
'_impl.nl': out[58],
|
|
'_impl.ntree_awake': out[59],
|
|
'_impl.nv_awake': out[60],
|
|
'_impl.qLD': out[61],
|
|
'_impl.qLDiagInv': out[62],
|
|
'qacc': out[63],
|
|
'qacc_smooth': out[64],
|
|
'qfrc_actuator': out[65],
|
|
'qfrc_bias': out[66],
|
|
'qfrc_constraint': out[67],
|
|
'_impl.qfrc_damper': out[68],
|
|
'qfrc_fluid': out[69],
|
|
'qfrc_gravcomp': out[70],
|
|
'qfrc_passive': out[71],
|
|
'qfrc_smooth': out[72],
|
|
'_impl.qfrc_spring': out[73],
|
|
'qvel': out[74],
|
|
'sensordata': out[75],
|
|
'site_xmat': out[76],
|
|
'site_xpos': out[77],
|
|
'_impl.solver_niter': out[78],
|
|
'_impl.subtree_angmom': out[79],
|
|
'subtree_com': out[80],
|
|
'_impl.subtree_linvel': out[81],
|
|
'_impl.ten_J': out[82],
|
|
'ten_length': out[83],
|
|
'_impl.ten_velocity': out[84],
|
|
'_impl.ten_wrapadr': out[85],
|
|
'_impl.ten_wrapnum': out[86],
|
|
'_impl.tree_asleep': out[87],
|
|
'_impl.tree_awake': out[88],
|
|
'_impl.tree_island': out[89],
|
|
'_impl.wrap_obj': out[90],
|
|
'_impl.wrap_xpos': out[91],
|
|
'xanchor': out[92],
|
|
'xaxis': out[93],
|
|
'ximat': out[94],
|
|
'xipos': out[95],
|
|
'xmat': out[96],
|
|
'xpos': out[97],
|
|
'xquat': out[98],
|
|
'_impl.contact__dim': out[99],
|
|
'_impl.contact__dist': out[100],
|
|
'_impl.contact__efc_address': out[101],
|
|
'_impl.contact__flex': out[102],
|
|
'_impl.contact__frame': out[103],
|
|
'_impl.contact__friction': out[104],
|
|
'_impl.contact__geom': out[105],
|
|
'_impl.contact__geomcollisionid': out[106],
|
|
'_impl.contact__includemargin': out[107],
|
|
'_impl.contact__pos': out[108],
|
|
'_impl.contact__solimp': out[109],
|
|
'_impl.contact__solref': out[110],
|
|
'_impl.contact__solreffriction': out[111],
|
|
'_impl.contact__type': out[112],
|
|
'_impl.contact__vert': out[113],
|
|
'_impl.contact__worldid': out[114],
|
|
'_impl.efc__D': out[115],
|
|
'_impl.efc__J': out[116],
|
|
'_impl.efc__J_colind': out[117],
|
|
'_impl.efc__J_rowadr': out[118],
|
|
'_impl.efc__J_rownnz': out[119],
|
|
'_impl.efc__Jqvel': out[120],
|
|
'_impl.efc__Ma': out[121],
|
|
'_impl.efc__aref': out[122],
|
|
'_impl.efc__force': out[123],
|
|
'_impl.efc__frictionloss': out[124],
|
|
'_impl.efc__iD': out[125],
|
|
'_impl.efc__iJ': out[126],
|
|
'_impl.efc__iJ_colind': out[127],
|
|
'_impl.efc__iJ_rowadr': out[128],
|
|
'_impl.efc__iJ_rownnz': out[129],
|
|
'_impl.efc__iaref': out[130],
|
|
'_impl.efc__id': out[131],
|
|
'_impl.efc__iforce': out[132],
|
|
'_impl.efc__ifrictionloss': out[133],
|
|
'_impl.efc__iid': out[134],
|
|
'_impl.efc__island': out[135],
|
|
'_impl.efc__istate': out[136],
|
|
'_impl.efc__itype': out[137],
|
|
'_impl.efc__margin': out[138],
|
|
'_impl.efc__pos': out[139],
|
|
'_impl.efc__state': out[140],
|
|
'_impl.efc__type': out[141],
|
|
'_impl.efc__vel': out[142],
|
|
})
|
|
return d
|
|
|
|
|
|
@jax.custom_batching.custom_vmap
|
|
@ffi.marshal_jax_warp_callable
|
|
def forward(m: types.Model, d: types.Data):
|
|
return _forward_jax_impl(m, d)
|
|
|
|
|
|
@forward.def_vmap
|
|
@ffi.marshal_custom_vmap
|
|
def forward_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
|
|
d = forward(m, d)
|
|
return d, is_batched[1]
|
|
|
|
|
|
@ffi.format_args_for_warp
|
|
def _step_shim(
|
|
# Model
|
|
nworld: int,
|
|
D_colind: wp.array[int],
|
|
D_diag: wp.array[int],
|
|
D_rowadr: wp.array[int],
|
|
D_rownnz: wp.array[int],
|
|
M_elemid: wp.array2d[int],
|
|
M_fullm_i: wp.array[int],
|
|
M_fullm_j: wp.array[int],
|
|
M_mulm_col: wp.array[int],
|
|
M_mulm_madr: wp.array[int],
|
|
M_mulm_rowadr: wp.array[int],
|
|
M_rowadr: wp.array[int],
|
|
M_rownnz: wp.array[int],
|
|
M_tiles: tuple[mjwp_types.TileSet, ...],
|
|
actuator_acc0: wp.array2d[float],
|
|
actuator_actadr: wp.array[int],
|
|
actuator_actearly: wp.array[bool],
|
|
actuator_actlimited: wp.array[bool],
|
|
actuator_actnum: wp.array[int],
|
|
actuator_actrange: wp.array2d[wp.vec2],
|
|
actuator_biasprm: wp.array2d[mjwp_types.vec10f],
|
|
actuator_biastype: wp.array[int],
|
|
actuator_cranklength: wp.array2d[float],
|
|
actuator_ctrllimited: wp.array[bool],
|
|
actuator_ctrlrange: wp.array2d[wp.vec2],
|
|
actuator_delay: wp.array[float],
|
|
actuator_dynprm: wp.array2d[mjwp_types.vec10f],
|
|
actuator_dyntype: wp.array[int],
|
|
actuator_forcelimited: wp.array[bool],
|
|
actuator_forcerange: wp.array2d[wp.vec2],
|
|
actuator_gainprm: wp.array2d[mjwp_types.vec10f],
|
|
actuator_gaintype: wp.array[int],
|
|
actuator_gear: wp.array2d[wp.spatial_vector],
|
|
actuator_history: wp.array[wp.vec2i],
|
|
actuator_historyadr: wp.array[int],
|
|
actuator_lengthrange: wp.array2d[wp.vec2],
|
|
actuator_trnid: wp.array[wp.vec2i],
|
|
actuator_trntype: wp.array[int],
|
|
actuator_trntype_body_adr: wp.array[int],
|
|
block_dim: mjwp_types.BlockDim,
|
|
body_branch_start: wp.array[int],
|
|
body_branches: wp.array[int],
|
|
body_dofadr: wp.array[int],
|
|
body_dofnum: wp.array[int],
|
|
body_fluid_ellipsoid: wp.array[bool],
|
|
body_geomadr: wp.array[int],
|
|
body_geomnum: wp.array[int],
|
|
body_gravcomp: wp.array2d[float],
|
|
body_inertia: wp.array2d[wp.vec3],
|
|
body_invweight0: wp.array2d[wp.vec2],
|
|
body_ipos: wp.array2d[wp.vec3],
|
|
body_iquat: wp.array2d[wp.quat],
|
|
body_isdofancestor: wp.array2d[int],
|
|
body_jntadr: wp.array[int],
|
|
body_jntnum: wp.array[int],
|
|
body_mass: wp.array2d[float],
|
|
body_mocapid: wp.array[int],
|
|
body_parentid: wp.array[int],
|
|
body_pos: wp.array2d[wp.vec3],
|
|
body_quat: wp.array2d[wp.quat],
|
|
body_rootid: wp.array[int],
|
|
body_subtreemass: wp.array2d[float],
|
|
body_tree: tuple[wp.array[int], ...],
|
|
body_treeid: wp.array[int],
|
|
body_weldid: wp.array[int],
|
|
cam_bodyid: wp.array[int],
|
|
cam_fovy: wp.array2d[float],
|
|
cam_intrinsic: wp.array2d[wp.vec4],
|
|
cam_mat0: wp.array2d[wp.mat33],
|
|
cam_mode: wp.array[int],
|
|
cam_pos: wp.array2d[wp.vec3],
|
|
cam_pos0: wp.array2d[wp.vec3],
|
|
cam_poscom0: wp.array2d[wp.vec3],
|
|
cam_quat: wp.array2d[wp.quat],
|
|
cam_resolution: wp.array[wp.vec2i],
|
|
cam_sensorsize: wp.array[wp.vec2],
|
|
cam_targetbodyid: wp.array[int],
|
|
dof_armature: wp.array2d[float],
|
|
dof_bodyid: wp.array[int],
|
|
dof_damping: wp.array2d[float],
|
|
dof_dampingpoly: wp.array2d[wp.vec2],
|
|
dof_frictionloss: wp.array2d[float],
|
|
dof_invweight0: wp.array2d[float],
|
|
dof_jntid: wp.array[int],
|
|
dof_length: wp.array[float],
|
|
dof_parentid: wp.array[int],
|
|
dof_solimp: wp.array2d[mjwp_types.vec5],
|
|
dof_solref: wp.array2d[wp.vec2],
|
|
dof_treeid: wp.array[int],
|
|
dof_tri_col: wp.array[int],
|
|
dof_tri_row: wp.array[int],
|
|
eq_connect_adr: wp.array[int],
|
|
eq_data: wp.array2d[mjwp_types.vec11],
|
|
eq_flex_adr: wp.array[int],
|
|
eq_jnt_adr: wp.array[int],
|
|
eq_obj1id: wp.array[int],
|
|
eq_obj2id: wp.array[int],
|
|
eq_objtype: wp.array[int],
|
|
eq_solimp: wp.array2d[mjwp_types.vec5],
|
|
eq_solref: wp.array2d[wp.vec2],
|
|
eq_ten_adr: wp.array[int],
|
|
eq_type: wp.array[int],
|
|
eq_wld_adr: wp.array[int],
|
|
flex_bending: wp.array[float],
|
|
flex_bendingadr: wp.array[int],
|
|
flex_centered: wp.array[bool],
|
|
flex_conaffinity: wp.array[int],
|
|
flex_condim: wp.array[int],
|
|
flex_contype: wp.array[int],
|
|
flex_damping: wp.array[float],
|
|
flex_dim: wp.array[int],
|
|
flex_edge: wp.array[wp.vec2i],
|
|
flex_edgeadr: wp.array[int],
|
|
flex_edgeflap: wp.array[wp.vec2i],
|
|
flex_edgenum: wp.array[int],
|
|
flex_elem: wp.array[int],
|
|
flex_elemadr: wp.array[int],
|
|
flex_elemdataadr: wp.array[int],
|
|
flex_elemedge: wp.array[int],
|
|
flex_elemedgeadr: 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_stiffness: wp.array[float],
|
|
flex_stiffnessadr: wp.array[int],
|
|
flex_vert: wp.array[wp.vec3],
|
|
flex_vertadr: wp.array[int],
|
|
flex_vertbodyid: wp.array[int],
|
|
flex_vertflexid: wp.array[int],
|
|
flex_vertnum: wp.array[int],
|
|
flexedge_J_colind: wp.array[int],
|
|
flexedge_J_rowadr: wp.array[int],
|
|
flexedge_J_rownnz: wp.array[int],
|
|
flexedge_invweight0: wp.array[float],
|
|
flexedge_length0: wp.array[float],
|
|
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_fluid: wp.array2d[float],
|
|
geom_friction: wp.array2d[wp.vec3],
|
|
geom_gap: wp.array2d[float],
|
|
geom_group: wp.array[int],
|
|
geom_margin: wp.array2d[float],
|
|
geom_matid: wp.array2d[int],
|
|
geom_pair_type_count: tuple[int, ...],
|
|
geom_plugin_index: wp.array[int],
|
|
geom_pos: wp.array2d[wp.vec3],
|
|
geom_priority: wp.array[int],
|
|
geom_quat: wp.array2d[wp.quat],
|
|
geom_rbound: wp.array2d[float],
|
|
geom_rgba: wp.array2d[wp.vec4],
|
|
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_fluid: bool,
|
|
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],
|
|
is_sparse: bool,
|
|
jnt_actfrclimited: wp.array[bool],
|
|
jnt_actfrcrange: wp.array2d[wp.vec2],
|
|
jnt_actgravcomp: wp.array[int],
|
|
jnt_axis: wp.array2d[wp.vec3],
|
|
jnt_bodyid: wp.array[int],
|
|
jnt_dofadr: wp.array[int],
|
|
jnt_limited_ball_adr: wp.array[int],
|
|
jnt_limited_slide_hinge_adr: wp.array[int],
|
|
jnt_margin: wp.array2d[float],
|
|
jnt_pos: wp.array2d[wp.vec3],
|
|
jnt_qposadr: wp.array[int],
|
|
jnt_range: wp.array2d[wp.vec2],
|
|
jnt_solimp: wp.array2d[mjwp_types.vec5],
|
|
jnt_solref: wp.array2d[wp.vec2],
|
|
jnt_stiffness: wp.array2d[float],
|
|
jnt_stiffnesspoly: wp.array2d[wp.vec2],
|
|
jnt_type: wp.array[int],
|
|
light_bodyid: wp.array[int],
|
|
light_dir: wp.array2d[wp.vec3],
|
|
light_dir0: wp.array2d[wp.vec3],
|
|
light_mode: wp.array[int],
|
|
light_pos: wp.array2d[wp.vec3],
|
|
light_pos0: wp.array2d[wp.vec3],
|
|
light_poscom0: wp.array2d[wp.vec3],
|
|
light_targetbodyid: wp.array[int],
|
|
mapM2D: wp.array[int],
|
|
mat_rgba: wp.array2d[wp.vec4],
|
|
max_ten_J_rownnz: int,
|
|
mesh_face: wp.array[wp.vec3i],
|
|
mesh_faceadr: wp.array[int],
|
|
mesh_graph: wp.array[int],
|
|
mesh_graphadr: wp.array[int],
|
|
mesh_normal: wp.array[wp.vec3],
|
|
mesh_normaladr: wp.array[int],
|
|
mesh_normalnum: 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_quat: wp.array[wp.quat],
|
|
mesh_vert: wp.array[wp.vec3],
|
|
mesh_vertadr: wp.array[int],
|
|
mesh_vertnum: wp.array[int],
|
|
nC: int,
|
|
nD: int,
|
|
nJten: int,
|
|
na: int,
|
|
nacttrnbody: int,
|
|
nbody: int,
|
|
nbranch: int,
|
|
ncam: int,
|
|
neq: int,
|
|
nflex: int,
|
|
nflexedge: int,
|
|
nflexelem: int,
|
|
nflexshelldata: int,
|
|
nflexvert: int,
|
|
ngeom: int,
|
|
ngravcomp: int,
|
|
nhistory: int,
|
|
njnt: int,
|
|
nlight: int,
|
|
nmaxcondim: int,
|
|
nmaxmeshdeg: int,
|
|
nmaxpolygon: int,
|
|
nmaxpyramid: int,
|
|
nmeshface: int,
|
|
nrangefinder: int,
|
|
nsensorcollision: int,
|
|
nsensorcontact: int,
|
|
nsensortaxel: int,
|
|
nsite: int,
|
|
ntendon: int,
|
|
ntree: int,
|
|
nu: int,
|
|
nv: int,
|
|
nv_pad: int,
|
|
nwrap: 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],
|
|
qD_fullm_i: wp.array[int],
|
|
qD_fullm_j: wp.array[int],
|
|
qLD_all_updates: wp.array[wp.vec3i],
|
|
qLD_level_offsets: wp.array[int],
|
|
qLD_updates: tuple[wp.array[wp.vec3i], ...],
|
|
qpos0: wp.array2d[float],
|
|
qpos_spring: wp.array2d[float],
|
|
rangefinder_sensor_adr: wp.array[int],
|
|
sensor_acc_adr: wp.array[int],
|
|
sensor_adr: wp.array[int],
|
|
sensor_adr_to_contact_adr: wp.array[int],
|
|
sensor_contact_adr: wp.array[int],
|
|
sensor_cutoff: wp.array[float],
|
|
sensor_datatype: wp.array[int],
|
|
sensor_delay: wp.array[float],
|
|
sensor_dim: wp.array[int],
|
|
sensor_e_kinetic: bool,
|
|
sensor_e_potential: bool,
|
|
sensor_history: wp.array[wp.vec2i],
|
|
sensor_historyadr: wp.array[int],
|
|
sensor_interval: wp.array[wp.vec2],
|
|
sensor_intprm: wp.array2d[int],
|
|
sensor_limitfrc_adr: wp.array[int],
|
|
sensor_limitpos_adr: wp.array[int],
|
|
sensor_limitvel_adr: wp.array[int],
|
|
sensor_objid: wp.array[int],
|
|
sensor_objtype: wp.array[int],
|
|
sensor_pos_adr: wp.array[int],
|
|
sensor_rangefinder_adr: wp.array[int],
|
|
sensor_rangefinder_bodyid: wp.array[int],
|
|
sensor_refid: wp.array[int],
|
|
sensor_reftype: wp.array[int],
|
|
sensor_rne_postconstraint: bool,
|
|
sensor_subtree_vel: bool,
|
|
sensor_tendonactfrc_adr: wp.array[int],
|
|
sensor_touch_adr: wp.array[int],
|
|
sensor_type: wp.array[int],
|
|
sensor_vel_adr: wp.array[int],
|
|
site_bodyid: wp.array[int],
|
|
site_pos: wp.array2d[wp.vec3],
|
|
site_quat: wp.array2d[wp.quat],
|
|
site_size: wp.array[wp.vec3],
|
|
site_type: wp.array[int],
|
|
taxel_sensorid: wp.array[int],
|
|
taxel_vertadr: wp.array[int],
|
|
ten_J_colind: wp.array[int],
|
|
ten_J_rowadr: wp.array[int],
|
|
ten_J_rownnz: wp.array[int],
|
|
tendon_actfrclimited: wp.array[bool],
|
|
tendon_actfrcrange: wp.array2d[wp.vec2],
|
|
tendon_adr: wp.array[int],
|
|
tendon_armature: wp.array2d[float],
|
|
tendon_damping: wp.array2d[float],
|
|
tendon_dampingpoly: wp.array2d[wp.vec2],
|
|
tendon_frictionloss: wp.array2d[float],
|
|
tendon_geom_adr: wp.array[int],
|
|
tendon_invweight0: wp.array2d[float],
|
|
tendon_jnt_adr: wp.array[int],
|
|
tendon_length0: wp.array2d[float],
|
|
tendon_lengthspring: wp.array2d[wp.vec2],
|
|
tendon_limited: wp.array[int],
|
|
tendon_limited_adr: wp.array[int],
|
|
tendon_margin: wp.array2d[float],
|
|
tendon_num: wp.array[int],
|
|
tendon_range: wp.array2d[wp.vec2],
|
|
tendon_site_pair_adr: wp.array[int],
|
|
tendon_solimp_fri: wp.array2d[mjwp_types.vec5],
|
|
tendon_solimp_lim: wp.array2d[mjwp_types.vec5],
|
|
tendon_solref_fri: wp.array2d[wp.vec2],
|
|
tendon_solref_lim: wp.array2d[wp.vec2],
|
|
tendon_stiffness: wp.array2d[float],
|
|
tendon_stiffnesspoly: wp.array2d[wp.vec2],
|
|
tree_dofadr: wp.array[int],
|
|
tree_dofnum: wp.array[int],
|
|
tree_sleep_policy: wp.array[int],
|
|
wrap_geom_adr: wp.array[int],
|
|
wrap_jnt_adr: wp.array[int],
|
|
wrap_objid: wp.array[int],
|
|
wrap_prm: wp.array[float],
|
|
wrap_pulley_scale: wp.array[float],
|
|
wrap_site_pair_adr: wp.array[int],
|
|
wrap_type: wp.array[int],
|
|
opt__broadphase: int,
|
|
opt__broadphase_filter: int,
|
|
opt__ccd_iterations: int,
|
|
opt__ccd_tolerance: wp.array[float],
|
|
opt__cone: int,
|
|
opt__contact_sensor_maxmatch: int,
|
|
opt__density: wp.array[float],
|
|
opt__disableflags: int,
|
|
opt__enableflags: int,
|
|
opt__graph_conditional: bool,
|
|
opt__gravity: wp.array[wp.vec3],
|
|
opt__impratio_invsqrt: wp.array[float],
|
|
opt__integrator: int,
|
|
opt__iterations: int,
|
|
opt__ls_iterations: int,
|
|
opt__ls_parallel: bool,
|
|
opt__ls_parallel_min_step: float,
|
|
opt__ls_tolerance: wp.array[float],
|
|
opt__magnetic: wp.array[wp.vec3],
|
|
opt__run_collision_detection: bool,
|
|
opt__sdf_initpoints: int,
|
|
opt__sdf_iterations: int,
|
|
opt__sleep_tolerance: wp.array[float],
|
|
opt__solver: int,
|
|
opt__timestep: wp.array[float],
|
|
opt__tolerance: wp.array[float],
|
|
opt__viscosity: wp.array[float],
|
|
opt__wind: wp.array[wp.vec3],
|
|
stat__meaninertia: wp.array[float],
|
|
# Data
|
|
naccdmax: int,
|
|
naconmax: int,
|
|
njmax: int,
|
|
njmax_nnz: int,
|
|
M: wp.array3d[float],
|
|
act: wp.array2d[float],
|
|
act_dot: wp.array2d[float],
|
|
actuator_force: wp.array2d[float],
|
|
actuator_length: wp.array2d[float],
|
|
actuator_moment: wp.array2d[float],
|
|
actuator_velocity: wp.array2d[float],
|
|
body_awake: wp.array2d[int],
|
|
body_awake_ind: wp.array2d[int],
|
|
cacc: wp.array2d[wp.spatial_vector],
|
|
cam_xmat: wp.array2d[wp.mat33],
|
|
cam_xpos: wp.array2d[wp.vec3],
|
|
cdof: wp.array2d[wp.spatial_vector],
|
|
cdof_dot: wp.array2d[wp.spatial_vector],
|
|
cfrc_ext: wp.array2d[wp.spatial_vector],
|
|
cfrc_int: wp.array2d[wp.spatial_vector],
|
|
cinert: wp.array2d[mjwp_types.vec10],
|
|
crb: wp.array2d[mjwp_types.vec10],
|
|
ctrl: wp.array2d[float],
|
|
cvel: wp.array2d[wp.spatial_vector],
|
|
dof_awake_ind: wp.array2d[int],
|
|
dof_island: wp.array2d[int],
|
|
dof_islandid: wp.array2d[int],
|
|
efc_islandid: wp.array2d[int],
|
|
energy: wp.array[wp.vec2],
|
|
eq_active: wp.array2d[bool],
|
|
flexedge_J: wp.array2d[float],
|
|
flexedge_length: wp.array2d[float],
|
|
flexedge_velocity: wp.array2d[float],
|
|
flexvert_xpos: wp.array2d[wp.vec3],
|
|
geom_xmat: wp.array2d[wp.mat33],
|
|
geom_xpos: wp.array2d[wp.vec3],
|
|
history: wp.array2d[float],
|
|
iqacc: wp.array2d[float],
|
|
iqacc_smooth: wp.array2d[float],
|
|
iqfrc_constraint: wp.array2d[float],
|
|
iqfrc_smooth: wp.array2d[float],
|
|
island_dofadr: wp.array2d[int],
|
|
island_efcadr: wp.array2d[int],
|
|
island_idofadr: wp.array2d[int],
|
|
island_ne: wp.array2d[int],
|
|
island_nefc: wp.array2d[int],
|
|
island_nf: wp.array2d[int],
|
|
island_nv: wp.array2d[int],
|
|
light_xdir: wp.array2d[wp.vec3],
|
|
light_xpos: wp.array2d[wp.vec3],
|
|
map_dof2idof: wp.array2d[int],
|
|
map_efc2iefc: wp.array2d[int],
|
|
map_idof2dof: wp.array2d[int],
|
|
map_iefc2efc: wp.array2d[int],
|
|
mocap_pos: wp.array2d[wp.vec3],
|
|
mocap_quat: wp.array2d[wp.quat],
|
|
moment_colind: wp.array2d[int],
|
|
moment_rowadr: wp.array2d[int],
|
|
moment_rownnz: wp.array2d[int],
|
|
nacon: wp.array[int],
|
|
nbody_awake: wp.array[int],
|
|
ncollision: wp.array[int],
|
|
ne: wp.array[int],
|
|
nefc: wp.array[int],
|
|
nf: wp.array[int],
|
|
nidof: wp.array[int],
|
|
nisland: wp.array[int],
|
|
nl: wp.array[int],
|
|
ntree_awake: wp.array[int],
|
|
nv_awake: wp.array[int],
|
|
qLD: wp.array3d[float],
|
|
qLDiagInv: wp.array2d[float],
|
|
qLU: wp.array3d[float],
|
|
qacc: wp.array2d[float],
|
|
qacc_smooth: wp.array2d[float],
|
|
qacc_warmstart: wp.array2d[float],
|
|
qfrc_actuator: wp.array2d[float],
|
|
qfrc_applied: wp.array2d[float],
|
|
qfrc_bias: wp.array2d[float],
|
|
qfrc_constraint: wp.array2d[float],
|
|
qfrc_damper: wp.array2d[float],
|
|
qfrc_fluid: wp.array2d[float],
|
|
qfrc_gravcomp: wp.array2d[float],
|
|
qfrc_passive: wp.array2d[float],
|
|
qfrc_smooth: wp.array2d[float],
|
|
qfrc_spring: wp.array2d[float],
|
|
qpos: wp.array2d[float],
|
|
qvel: wp.array2d[float],
|
|
sensordata: wp.array2d[float],
|
|
site_xmat: wp.array2d[wp.mat33],
|
|
site_xpos: wp.array2d[wp.vec3],
|
|
solver_niter: wp.array[int],
|
|
subtree_angmom: wp.array2d[wp.vec3],
|
|
subtree_com: wp.array2d[wp.vec3],
|
|
subtree_linvel: wp.array2d[wp.vec3],
|
|
ten_J: wp.array2d[float],
|
|
ten_length: wp.array2d[float],
|
|
ten_velocity: wp.array2d[float],
|
|
ten_wrapadr: wp.array2d[int],
|
|
ten_wrapnum: wp.array2d[int],
|
|
time: wp.array[float],
|
|
tree_asleep: wp.array2d[int],
|
|
tree_awake: wp.array2d[int],
|
|
tree_island: wp.array2d[int],
|
|
wrap_obj: wp.array2d[wp.vec2i],
|
|
wrap_xpos: wp.array2d[wp.spatial_vector],
|
|
xanchor: wp.array2d[wp.vec3],
|
|
xaxis: wp.array2d[wp.vec3],
|
|
xfrc_applied: wp.array2d[wp.spatial_vector],
|
|
ximat: wp.array2d[wp.mat33],
|
|
xipos: wp.array2d[wp.vec3],
|
|
xmat: wp.array2d[wp.mat33],
|
|
xpos: wp.array2d[wp.vec3],
|
|
xquat: wp.array2d[wp.quat],
|
|
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],
|
|
efc__D: wp.array2d[float],
|
|
efc__J: wp.array3d[float],
|
|
efc__J_colind: wp.array3d[int],
|
|
efc__J_rowadr: wp.array2d[int],
|
|
efc__J_rownnz: wp.array2d[int],
|
|
efc__Jqvel: wp.array2d[float],
|
|
efc__Ma: wp.array2d[float],
|
|
efc__aref: wp.array2d[float],
|
|
efc__force: wp.array2d[float],
|
|
efc__frictionloss: wp.array2d[float],
|
|
efc__iD: wp.array2d[float],
|
|
efc__iJ: wp.array3d[float],
|
|
efc__iJ_colind: wp.array3d[int],
|
|
efc__iJ_rowadr: wp.array2d[int],
|
|
efc__iJ_rownnz: wp.array2d[int],
|
|
efc__iaref: wp.array2d[float],
|
|
efc__id: wp.array2d[int],
|
|
efc__iforce: wp.array2d[float],
|
|
efc__ifrictionloss: wp.array2d[float],
|
|
efc__iid: wp.array2d[int],
|
|
efc__island: wp.array2d[int],
|
|
efc__istate: wp.array2d[int],
|
|
efc__itype: wp.array2d[int],
|
|
efc__margin: wp.array2d[float],
|
|
efc__pos: wp.array2d[float],
|
|
efc__state: wp.array2d[int],
|
|
efc__type: wp.array2d[int],
|
|
efc__vel: wp.array2d[float],
|
|
):
|
|
_m.stat = _s
|
|
_m.opt = _o
|
|
_m.callback = _cb
|
|
_d.efc = _e
|
|
_d.contact = _c
|
|
_m.D_colind = D_colind
|
|
_m.D_diag = D_diag
|
|
_m.D_rowadr = D_rowadr
|
|
_m.D_rownnz = D_rownnz
|
|
_m.M_elemid = M_elemid
|
|
_m.M_fullm_i = M_fullm_i
|
|
_m.M_fullm_j = M_fullm_j
|
|
_m.M_mulm_col = M_mulm_col
|
|
_m.M_mulm_madr = M_mulm_madr
|
|
_m.M_mulm_rowadr = M_mulm_rowadr
|
|
_m.M_rowadr = M_rowadr
|
|
_m.M_rownnz = M_rownnz
|
|
_m.M_tiles = M_tiles
|
|
_m.actuator_acc0 = actuator_acc0
|
|
_m.actuator_actadr = actuator_actadr
|
|
_m.actuator_actearly = actuator_actearly
|
|
_m.actuator_actlimited = actuator_actlimited
|
|
_m.actuator_actnum = actuator_actnum
|
|
_m.actuator_actrange = actuator_actrange
|
|
_m.actuator_biasprm = actuator_biasprm
|
|
_m.actuator_biastype = actuator_biastype
|
|
_m.actuator_cranklength = actuator_cranklength
|
|
_m.actuator_ctrllimited = actuator_ctrllimited
|
|
_m.actuator_ctrlrange = actuator_ctrlrange
|
|
_m.actuator_delay = actuator_delay
|
|
_m.actuator_dynprm = actuator_dynprm
|
|
_m.actuator_dyntype = actuator_dyntype
|
|
_m.actuator_forcelimited = actuator_forcelimited
|
|
_m.actuator_forcerange = actuator_forcerange
|
|
_m.actuator_gainprm = actuator_gainprm
|
|
_m.actuator_gaintype = actuator_gaintype
|
|
_m.actuator_gear = actuator_gear
|
|
_m.actuator_history = actuator_history
|
|
_m.actuator_historyadr = actuator_historyadr
|
|
_m.actuator_lengthrange = actuator_lengthrange
|
|
_m.actuator_trnid = actuator_trnid
|
|
_m.actuator_trntype = actuator_trntype
|
|
_m.actuator_trntype_body_adr = actuator_trntype_body_adr
|
|
_m.block_dim = block_dim
|
|
_m.body_branch_start = body_branch_start
|
|
_m.body_branches = body_branches
|
|
_m.body_dofadr = body_dofadr
|
|
_m.body_dofnum = body_dofnum
|
|
_m.body_fluid_ellipsoid = body_fluid_ellipsoid
|
|
_m.body_geomadr = body_geomadr
|
|
_m.body_geomnum = body_geomnum
|
|
_m.body_gravcomp = body_gravcomp
|
|
_m.body_inertia = body_inertia
|
|
_m.body_invweight0 = body_invweight0
|
|
_m.body_ipos = body_ipos
|
|
_m.body_iquat = body_iquat
|
|
_m.body_isdofancestor = body_isdofancestor
|
|
_m.body_jntadr = body_jntadr
|
|
_m.body_jntnum = body_jntnum
|
|
_m.body_mass = body_mass
|
|
_m.body_mocapid = body_mocapid
|
|
_m.body_parentid = body_parentid
|
|
_m.body_pos = body_pos
|
|
_m.body_quat = body_quat
|
|
_m.body_rootid = body_rootid
|
|
_m.body_subtreemass = body_subtreemass
|
|
_m.body_tree = body_tree
|
|
_m.body_treeid = body_treeid
|
|
_m.body_weldid = body_weldid
|
|
_m.cam_bodyid = cam_bodyid
|
|
_m.cam_fovy = cam_fovy
|
|
_m.cam_intrinsic = cam_intrinsic
|
|
_m.cam_mat0 = cam_mat0
|
|
_m.cam_mode = cam_mode
|
|
_m.cam_pos = cam_pos
|
|
_m.cam_pos0 = cam_pos0
|
|
_m.cam_poscom0 = cam_poscom0
|
|
_m.cam_quat = cam_quat
|
|
_m.cam_resolution = cam_resolution
|
|
_m.cam_sensorsize = cam_sensorsize
|
|
_m.cam_targetbodyid = cam_targetbodyid
|
|
_m.dof_armature = dof_armature
|
|
_m.dof_bodyid = dof_bodyid
|
|
_m.dof_damping = dof_damping
|
|
_m.dof_dampingpoly = dof_dampingpoly
|
|
_m.dof_frictionloss = dof_frictionloss
|
|
_m.dof_invweight0 = dof_invweight0
|
|
_m.dof_jntid = dof_jntid
|
|
_m.dof_length = dof_length
|
|
_m.dof_parentid = dof_parentid
|
|
_m.dof_solimp = dof_solimp
|
|
_m.dof_solref = dof_solref
|
|
_m.dof_treeid = dof_treeid
|
|
_m.dof_tri_col = dof_tri_col
|
|
_m.dof_tri_row = dof_tri_row
|
|
_m.eq_connect_adr = eq_connect_adr
|
|
_m.eq_data = eq_data
|
|
_m.eq_flex_adr = eq_flex_adr
|
|
_m.eq_jnt_adr = eq_jnt_adr
|
|
_m.eq_obj1id = eq_obj1id
|
|
_m.eq_obj2id = eq_obj2id
|
|
_m.eq_objtype = eq_objtype
|
|
_m.eq_solimp = eq_solimp
|
|
_m.eq_solref = eq_solref
|
|
_m.eq_ten_adr = eq_ten_adr
|
|
_m.eq_type = eq_type
|
|
_m.eq_wld_adr = eq_wld_adr
|
|
_m.flex_bending = flex_bending
|
|
_m.flex_bendingadr = flex_bendingadr
|
|
_m.flex_centered = flex_centered
|
|
_m.flex_conaffinity = flex_conaffinity
|
|
_m.flex_condim = flex_condim
|
|
_m.flex_contype = flex_contype
|
|
_m.flex_damping = flex_damping
|
|
_m.flex_dim = flex_dim
|
|
_m.flex_edge = flex_edge
|
|
_m.flex_edgeadr = flex_edgeadr
|
|
_m.flex_edgeflap = flex_edgeflap
|
|
_m.flex_edgenum = flex_edgenum
|
|
_m.flex_elem = flex_elem
|
|
_m.flex_elemadr = flex_elemadr
|
|
_m.flex_elemdataadr = flex_elemdataadr
|
|
_m.flex_elemedge = flex_elemedge
|
|
_m.flex_elemedgeadr = flex_elemedgeadr
|
|
_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_stiffness = flex_stiffness
|
|
_m.flex_stiffnessadr = flex_stiffnessadr
|
|
_m.flex_vert = flex_vert
|
|
_m.flex_vertadr = flex_vertadr
|
|
_m.flex_vertbodyid = flex_vertbodyid
|
|
_m.flex_vertflexid = flex_vertflexid
|
|
_m.flex_vertnum = flex_vertnum
|
|
_m.flexedge_J_colind = flexedge_J_colind
|
|
_m.flexedge_J_rowadr = flexedge_J_rowadr
|
|
_m.flexedge_J_rownnz = flexedge_J_rownnz
|
|
_m.flexedge_invweight0 = flexedge_invweight0
|
|
_m.flexedge_length0 = flexedge_length0
|
|
_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_fluid = geom_fluid
|
|
_m.geom_friction = geom_friction
|
|
_m.geom_gap = geom_gap
|
|
_m.geom_group = geom_group
|
|
_m.geom_margin = geom_margin
|
|
_m.geom_matid = geom_matid
|
|
_m.geom_pair_type_count = geom_pair_type_count
|
|
_m.geom_plugin_index = geom_plugin_index
|
|
_m.geom_pos = geom_pos
|
|
_m.geom_priority = geom_priority
|
|
_m.geom_quat = geom_quat
|
|
_m.geom_rbound = geom_rbound
|
|
_m.geom_rgba = geom_rgba
|
|
_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_fluid = has_fluid
|
|
_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.is_sparse = is_sparse
|
|
_m.jnt_actfrclimited = jnt_actfrclimited
|
|
_m.jnt_actfrcrange = jnt_actfrcrange
|
|
_m.jnt_actgravcomp = jnt_actgravcomp
|
|
_m.jnt_axis = jnt_axis
|
|
_m.jnt_bodyid = jnt_bodyid
|
|
_m.jnt_dofadr = jnt_dofadr
|
|
_m.jnt_limited_ball_adr = jnt_limited_ball_adr
|
|
_m.jnt_limited_slide_hinge_adr = jnt_limited_slide_hinge_adr
|
|
_m.jnt_margin = jnt_margin
|
|
_m.jnt_pos = jnt_pos
|
|
_m.jnt_qposadr = jnt_qposadr
|
|
_m.jnt_range = jnt_range
|
|
_m.jnt_solimp = jnt_solimp
|
|
_m.jnt_solref = jnt_solref
|
|
_m.jnt_stiffness = jnt_stiffness
|
|
_m.jnt_stiffnesspoly = jnt_stiffnesspoly
|
|
_m.jnt_type = jnt_type
|
|
_m.light_bodyid = light_bodyid
|
|
_m.light_dir = light_dir
|
|
_m.light_dir0 = light_dir0
|
|
_m.light_mode = light_mode
|
|
_m.light_pos = light_pos
|
|
_m.light_pos0 = light_pos0
|
|
_m.light_poscom0 = light_poscom0
|
|
_m.light_targetbodyid = light_targetbodyid
|
|
_m.mapM2D = mapM2D
|
|
_m.mat_rgba = mat_rgba
|
|
_m.max_ten_J_rownnz = max_ten_J_rownnz
|
|
_m.mesh_face = mesh_face
|
|
_m.mesh_faceadr = mesh_faceadr
|
|
_m.mesh_graph = mesh_graph
|
|
_m.mesh_graphadr = mesh_graphadr
|
|
_m.mesh_normal = mesh_normal
|
|
_m.mesh_normaladr = mesh_normaladr
|
|
_m.mesh_normalnum = mesh_normalnum
|
|
_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_quat = mesh_quat
|
|
_m.mesh_vert = mesh_vert
|
|
_m.mesh_vertadr = mesh_vertadr
|
|
_m.mesh_vertnum = mesh_vertnum
|
|
_m.nC = nC
|
|
_m.nD = nD
|
|
_m.nJten = nJten
|
|
_m.na = na
|
|
_m.nacttrnbody = nacttrnbody
|
|
_m.nbody = nbody
|
|
_m.nbranch = nbranch
|
|
_m.ncam = ncam
|
|
_m.neq = neq
|
|
_m.nflex = nflex
|
|
_m.nflexedge = nflexedge
|
|
_m.nflexelem = nflexelem
|
|
_m.nflexshelldata = nflexshelldata
|
|
_m.nflexvert = nflexvert
|
|
_m.ngeom = ngeom
|
|
_m.ngravcomp = ngravcomp
|
|
_m.nhistory = nhistory
|
|
_m.njnt = njnt
|
|
_m.nlight = nlight
|
|
_m.nmaxcondim = nmaxcondim
|
|
_m.nmaxmeshdeg = nmaxmeshdeg
|
|
_m.nmaxpolygon = nmaxpolygon
|
|
_m.nmaxpyramid = nmaxpyramid
|
|
_m.nmeshface = nmeshface
|
|
_m.nrangefinder = nrangefinder
|
|
_m.nsensorcollision = nsensorcollision
|
|
_m.nsensorcontact = nsensorcontact
|
|
_m.nsensortaxel = nsensortaxel
|
|
_m.nsite = nsite
|
|
_m.ntendon = ntendon
|
|
_m.ntree = ntree
|
|
_m.nu = nu
|
|
_m.nv = nv
|
|
_m.nv_pad = nv_pad
|
|
_m.nwrap = nwrap
|
|
_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.cone = opt__cone
|
|
_m.opt.contact_sensor_maxmatch = opt__contact_sensor_maxmatch
|
|
_m.opt.density = opt__density
|
|
_m.opt.disableflags = opt__disableflags
|
|
_m.opt.enableflags = opt__enableflags
|
|
_m.opt.graph_conditional = opt__graph_conditional
|
|
_m.opt.gravity = opt__gravity
|
|
_m.opt.impratio_invsqrt = opt__impratio_invsqrt
|
|
_m.opt.integrator = opt__integrator
|
|
_m.opt.iterations = opt__iterations
|
|
_m.opt.ls_iterations = opt__ls_iterations
|
|
_m.opt.ls_parallel = opt__ls_parallel
|
|
_m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
|
|
_m.opt.ls_tolerance = opt__ls_tolerance
|
|
_m.opt.magnetic = opt__magnetic
|
|
_m.opt.run_collision_detection = opt__run_collision_detection
|
|
_m.opt.sdf_initpoints = opt__sdf_initpoints
|
|
_m.opt.sdf_iterations = opt__sdf_iterations
|
|
_m.opt.sleep_tolerance = opt__sleep_tolerance
|
|
_m.opt.solver = opt__solver
|
|
_m.opt.timestep = opt__timestep
|
|
_m.opt.tolerance = opt__tolerance
|
|
_m.opt.viscosity = opt__viscosity
|
|
_m.opt.wind = opt__wind
|
|
_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
|
|
_m.qD_fullm_i = qD_fullm_i
|
|
_m.qD_fullm_j = qD_fullm_j
|
|
_m.qLD_all_updates = qLD_all_updates
|
|
_m.qLD_level_offsets = qLD_level_offsets
|
|
_m.qLD_updates = qLD_updates
|
|
_m.qpos0 = qpos0
|
|
_m.qpos_spring = qpos_spring
|
|
_m.rangefinder_sensor_adr = rangefinder_sensor_adr
|
|
_m.sensor_acc_adr = sensor_acc_adr
|
|
_m.sensor_adr = sensor_adr
|
|
_m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr
|
|
_m.sensor_contact_adr = sensor_contact_adr
|
|
_m.sensor_cutoff = sensor_cutoff
|
|
_m.sensor_datatype = sensor_datatype
|
|
_m.sensor_delay = sensor_delay
|
|
_m.sensor_dim = sensor_dim
|
|
_m.sensor_e_kinetic = sensor_e_kinetic
|
|
_m.sensor_e_potential = sensor_e_potential
|
|
_m.sensor_history = sensor_history
|
|
_m.sensor_historyadr = sensor_historyadr
|
|
_m.sensor_interval = sensor_interval
|
|
_m.sensor_intprm = sensor_intprm
|
|
_m.sensor_limitfrc_adr = sensor_limitfrc_adr
|
|
_m.sensor_limitpos_adr = sensor_limitpos_adr
|
|
_m.sensor_limitvel_adr = sensor_limitvel_adr
|
|
_m.sensor_objid = sensor_objid
|
|
_m.sensor_objtype = sensor_objtype
|
|
_m.sensor_pos_adr = sensor_pos_adr
|
|
_m.sensor_rangefinder_adr = sensor_rangefinder_adr
|
|
_m.sensor_rangefinder_bodyid = sensor_rangefinder_bodyid
|
|
_m.sensor_refid = sensor_refid
|
|
_m.sensor_reftype = sensor_reftype
|
|
_m.sensor_rne_postconstraint = sensor_rne_postconstraint
|
|
_m.sensor_subtree_vel = sensor_subtree_vel
|
|
_m.sensor_tendonactfrc_adr = sensor_tendonactfrc_adr
|
|
_m.sensor_touch_adr = sensor_touch_adr
|
|
_m.sensor_type = sensor_type
|
|
_m.sensor_vel_adr = sensor_vel_adr
|
|
_m.site_bodyid = site_bodyid
|
|
_m.site_pos = site_pos
|
|
_m.site_quat = site_quat
|
|
_m.site_size = site_size
|
|
_m.site_type = site_type
|
|
_m.stat.meaninertia = stat__meaninertia
|
|
_m.taxel_sensorid = taxel_sensorid
|
|
_m.taxel_vertadr = taxel_vertadr
|
|
_m.ten_J_colind = ten_J_colind
|
|
_m.ten_J_rowadr = ten_J_rowadr
|
|
_m.ten_J_rownnz = ten_J_rownnz
|
|
_m.tendon_actfrclimited = tendon_actfrclimited
|
|
_m.tendon_actfrcrange = tendon_actfrcrange
|
|
_m.tendon_adr = tendon_adr
|
|
_m.tendon_armature = tendon_armature
|
|
_m.tendon_damping = tendon_damping
|
|
_m.tendon_dampingpoly = tendon_dampingpoly
|
|
_m.tendon_frictionloss = tendon_frictionloss
|
|
_m.tendon_geom_adr = tendon_geom_adr
|
|
_m.tendon_invweight0 = tendon_invweight0
|
|
_m.tendon_jnt_adr = tendon_jnt_adr
|
|
_m.tendon_length0 = tendon_length0
|
|
_m.tendon_lengthspring = tendon_lengthspring
|
|
_m.tendon_limited = tendon_limited
|
|
_m.tendon_limited_adr = tendon_limited_adr
|
|
_m.tendon_margin = tendon_margin
|
|
_m.tendon_num = tendon_num
|
|
_m.tendon_range = tendon_range
|
|
_m.tendon_site_pair_adr = tendon_site_pair_adr
|
|
_m.tendon_solimp_fri = tendon_solimp_fri
|
|
_m.tendon_solimp_lim = tendon_solimp_lim
|
|
_m.tendon_solref_fri = tendon_solref_fri
|
|
_m.tendon_solref_lim = tendon_solref_lim
|
|
_m.tendon_stiffness = tendon_stiffness
|
|
_m.tendon_stiffnesspoly = tendon_stiffnesspoly
|
|
_m.tree_dofadr = tree_dofadr
|
|
_m.tree_dofnum = tree_dofnum
|
|
_m.tree_sleep_policy = tree_sleep_policy
|
|
_m.wrap_geom_adr = wrap_geom_adr
|
|
_m.wrap_jnt_adr = wrap_jnt_adr
|
|
_m.wrap_objid = wrap_objid
|
|
_m.wrap_prm = wrap_prm
|
|
_m.wrap_pulley_scale = wrap_pulley_scale
|
|
_m.wrap_site_pair_adr = wrap_site_pair_adr
|
|
_m.wrap_type = wrap_type
|
|
_d.M = M
|
|
_d.act = act
|
|
_d.act_dot = act_dot
|
|
_d.actuator_force = actuator_force
|
|
_d.actuator_length = actuator_length
|
|
_d.actuator_moment = actuator_moment
|
|
_d.actuator_velocity = actuator_velocity
|
|
_d.body_awake = body_awake
|
|
_d.body_awake_ind = body_awake_ind
|
|
_d.cacc = cacc
|
|
_d.cam_xmat = cam_xmat
|
|
_d.cam_xpos = cam_xpos
|
|
_d.cdof = cdof
|
|
_d.cdof_dot = cdof_dot
|
|
_d.cfrc_ext = cfrc_ext
|
|
_d.cfrc_int = cfrc_int
|
|
_d.cinert = cinert
|
|
_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.crb = crb
|
|
_d.ctrl = ctrl
|
|
_d.cvel = cvel
|
|
_d.dof_awake_ind = dof_awake_ind
|
|
_d.dof_island = dof_island
|
|
_d.dof_islandid = dof_islandid
|
|
_d.efc.D = efc__D
|
|
_d.efc.J = efc__J
|
|
_d.efc.J_colind = efc__J_colind
|
|
_d.efc.J_rowadr = efc__J_rowadr
|
|
_d.efc.J_rownnz = efc__J_rownnz
|
|
_d.efc.Jqvel = efc__Jqvel
|
|
_d.efc.Ma = efc__Ma
|
|
_d.efc.aref = efc__aref
|
|
_d.efc.force = efc__force
|
|
_d.efc.frictionloss = efc__frictionloss
|
|
_d.efc.iD = efc__iD
|
|
_d.efc.iJ = efc__iJ
|
|
_d.efc.iJ_colind = efc__iJ_colind
|
|
_d.efc.iJ_rowadr = efc__iJ_rowadr
|
|
_d.efc.iJ_rownnz = efc__iJ_rownnz
|
|
_d.efc.iaref = efc__iaref
|
|
_d.efc.id = efc__id
|
|
_d.efc.iforce = efc__iforce
|
|
_d.efc.ifrictionloss = efc__ifrictionloss
|
|
_d.efc.iid = efc__iid
|
|
_d.efc.island = efc__island
|
|
_d.efc.istate = efc__istate
|
|
_d.efc.itype = efc__itype
|
|
_d.efc.margin = efc__margin
|
|
_d.efc.pos = efc__pos
|
|
_d.efc.state = efc__state
|
|
_d.efc.type = efc__type
|
|
_d.efc.vel = efc__vel
|
|
_d.efc_islandid = efc_islandid
|
|
_d.energy = energy
|
|
_d.eq_active = eq_active
|
|
_d.flexedge_J = flexedge_J
|
|
_d.flexedge_length = flexedge_length
|
|
_d.flexedge_velocity = flexedge_velocity
|
|
_d.flexvert_xpos = flexvert_xpos
|
|
_d.geom_xmat = geom_xmat
|
|
_d.geom_xpos = geom_xpos
|
|
_d.history = history
|
|
_d.iqacc = iqacc
|
|
_d.iqacc_smooth = iqacc_smooth
|
|
_d.iqfrc_constraint = iqfrc_constraint
|
|
_d.iqfrc_smooth = iqfrc_smooth
|
|
_d.island_dofadr = island_dofadr
|
|
_d.island_efcadr = island_efcadr
|
|
_d.island_idofadr = island_idofadr
|
|
_d.island_ne = island_ne
|
|
_d.island_nefc = island_nefc
|
|
_d.island_nf = island_nf
|
|
_d.island_nv = island_nv
|
|
_d.light_xdir = light_xdir
|
|
_d.light_xpos = light_xpos
|
|
_d.map_dof2idof = map_dof2idof
|
|
_d.map_efc2iefc = map_efc2iefc
|
|
_d.map_idof2dof = map_idof2dof
|
|
_d.map_iefc2efc = map_iefc2efc
|
|
_d.mocap_pos = mocap_pos
|
|
_d.mocap_quat = mocap_quat
|
|
_d.moment_colind = moment_colind
|
|
_d.moment_rowadr = moment_rowadr
|
|
_d.moment_rownnz = moment_rownnz
|
|
_d.naccdmax = naccdmax
|
|
_d.nacon = nacon
|
|
_d.naconmax = naconmax
|
|
_d.nbody_awake = nbody_awake
|
|
_d.ncollision = ncollision
|
|
_d.ne = ne
|
|
_d.nefc = nefc
|
|
_d.nf = nf
|
|
_d.nidof = nidof
|
|
_d.nisland = nisland
|
|
_d.njmax = njmax
|
|
_d.njmax_nnz = njmax_nnz
|
|
_d.nl = nl
|
|
_d.ntree_awake = ntree_awake
|
|
_d.nv_awake = nv_awake
|
|
_d.qLD = qLD
|
|
_d.qLDiagInv = qLDiagInv
|
|
_d.qLU = qLU
|
|
_d.qacc = qacc
|
|
_d.qacc_smooth = qacc_smooth
|
|
_d.qacc_warmstart = qacc_warmstart
|
|
_d.qfrc_actuator = qfrc_actuator
|
|
_d.qfrc_applied = qfrc_applied
|
|
_d.qfrc_bias = qfrc_bias
|
|
_d.qfrc_constraint = qfrc_constraint
|
|
_d.qfrc_damper = qfrc_damper
|
|
_d.qfrc_fluid = qfrc_fluid
|
|
_d.qfrc_gravcomp = qfrc_gravcomp
|
|
_d.qfrc_passive = qfrc_passive
|
|
_d.qfrc_smooth = qfrc_smooth
|
|
_d.qfrc_spring = qfrc_spring
|
|
_d.qpos = qpos
|
|
_d.qvel = qvel
|
|
_d.sensordata = sensordata
|
|
_d.site_xmat = site_xmat
|
|
_d.site_xpos = site_xpos
|
|
_d.solver_niter = solver_niter
|
|
_d.subtree_angmom = subtree_angmom
|
|
_d.subtree_com = subtree_com
|
|
_d.subtree_linvel = subtree_linvel
|
|
_d.ten_J = ten_J
|
|
_d.ten_length = ten_length
|
|
_d.ten_velocity = ten_velocity
|
|
_d.ten_wrapadr = ten_wrapadr
|
|
_d.ten_wrapnum = ten_wrapnum
|
|
_d.time = time
|
|
_d.tree_asleep = tree_asleep
|
|
_d.tree_awake = tree_awake
|
|
_d.tree_island = tree_island
|
|
_d.wrap_obj = wrap_obj
|
|
_d.wrap_xpos = wrap_xpos
|
|
_d.xanchor = xanchor
|
|
_d.xaxis = xaxis
|
|
_d.xfrc_applied = xfrc_applied
|
|
_d.ximat = ximat
|
|
_d.xipos = xipos
|
|
_d.xmat = xmat
|
|
_d.xpos = xpos
|
|
_d.xquat = xquat
|
|
_d.nworld = nworld
|
|
mjwarp.step(_m, _d)
|
|
|
|
|
|
def _step_jax_impl(m: types.Model, d: types.Data):
|
|
output_dims = {
|
|
'M': d._impl.M.shape,
|
|
'act': d.act.shape,
|
|
'act_dot': d.act_dot.shape,
|
|
'actuator_force': d.actuator_force.shape,
|
|
'actuator_length': d.actuator_length.shape,
|
|
'actuator_moment': d._impl.actuator_moment.shape,
|
|
'actuator_velocity': d._impl.actuator_velocity.shape,
|
|
'body_awake': d._impl.body_awake.shape,
|
|
'body_awake_ind': d._impl.body_awake_ind.shape,
|
|
'cacc': d._impl.cacc.shape,
|
|
'cam_xmat': d.cam_xmat.shape,
|
|
'cam_xpos': d.cam_xpos.shape,
|
|
'cdof': d.cdof.shape,
|
|
'cdof_dot': d.cdof_dot.shape,
|
|
'cfrc_ext': d._impl.cfrc_ext.shape,
|
|
'cfrc_int': d._impl.cfrc_int.shape,
|
|
'cinert': d._impl.cinert.shape,
|
|
'crb': d._impl.crb.shape,
|
|
'cvel': d.cvel.shape,
|
|
'dof_awake_ind': d._impl.dof_awake_ind.shape,
|
|
'dof_island': d._impl.dof_island.shape,
|
|
'dof_islandid': d._impl.dof_islandid.shape,
|
|
'efc_islandid': d._impl.efc_islandid.shape,
|
|
'energy': d._impl.energy.shape,
|
|
'flexedge_J': d._impl.flexedge_J.shape,
|
|
'flexedge_length': d._impl.flexedge_length.shape,
|
|
'flexedge_velocity': d._impl.flexedge_velocity.shape,
|
|
'flexvert_xpos': d._impl.flexvert_xpos.shape,
|
|
'geom_xmat': d.geom_xmat.shape,
|
|
'geom_xpos': d.geom_xpos.shape,
|
|
'history': d.history.shape,
|
|
'iqacc': d._impl.iqacc.shape,
|
|
'iqacc_smooth': d._impl.iqacc_smooth.shape,
|
|
'iqfrc_constraint': d._impl.iqfrc_constraint.shape,
|
|
'iqfrc_smooth': d._impl.iqfrc_smooth.shape,
|
|
'island_dofadr': d._impl.island_dofadr.shape,
|
|
'island_efcadr': d._impl.island_efcadr.shape,
|
|
'island_idofadr': d._impl.island_idofadr.shape,
|
|
'island_ne': d._impl.island_ne.shape,
|
|
'island_nefc': d._impl.island_nefc.shape,
|
|
'island_nf': d._impl.island_nf.shape,
|
|
'island_nv': d._impl.island_nv.shape,
|
|
'light_xdir': d._impl.light_xdir.shape,
|
|
'light_xpos': d._impl.light_xpos.shape,
|
|
'map_dof2idof': d._impl.map_dof2idof.shape,
|
|
'map_efc2iefc': d._impl.map_efc2iefc.shape,
|
|
'map_idof2dof': d._impl.map_idof2dof.shape,
|
|
'map_iefc2efc': d._impl.map_iefc2efc.shape,
|
|
'moment_colind': d._impl.moment_colind.shape,
|
|
'moment_rowadr': d._impl.moment_rowadr.shape,
|
|
'moment_rownnz': d._impl.moment_rownnz.shape,
|
|
'nacon': d._impl.nacon.shape,
|
|
'nbody_awake': d._impl.nbody_awake.shape,
|
|
'ncollision': d._impl.ncollision.shape,
|
|
'ne': d._impl.ne.shape,
|
|
'nefc': d._impl.nefc.shape,
|
|
'nf': d._impl.nf.shape,
|
|
'nidof': d._impl.nidof.shape,
|
|
'nisland': d._impl.nisland.shape,
|
|
'nl': d._impl.nl.shape,
|
|
'ntree_awake': d._impl.ntree_awake.shape,
|
|
'nv_awake': d._impl.nv_awake.shape,
|
|
'qLD': d._impl.qLD.shape,
|
|
'qLDiagInv': d._impl.qLDiagInv.shape,
|
|
'qLU': d._impl.qLU.shape,
|
|
'qacc': d.qacc.shape,
|
|
'qacc_smooth': d.qacc_smooth.shape,
|
|
'qacc_warmstart': d.qacc_warmstart.shape,
|
|
'qfrc_actuator': d.qfrc_actuator.shape,
|
|
'qfrc_bias': d.qfrc_bias.shape,
|
|
'qfrc_constraint': d.qfrc_constraint.shape,
|
|
'qfrc_damper': d._impl.qfrc_damper.shape,
|
|
'qfrc_fluid': d.qfrc_fluid.shape,
|
|
'qfrc_gravcomp': d.qfrc_gravcomp.shape,
|
|
'qfrc_passive': d.qfrc_passive.shape,
|
|
'qfrc_smooth': d.qfrc_smooth.shape,
|
|
'qfrc_spring': d._impl.qfrc_spring.shape,
|
|
'qpos': d.qpos.shape,
|
|
'qvel': d.qvel.shape,
|
|
'sensordata': d.sensordata.shape,
|
|
'site_xmat': d.site_xmat.shape,
|
|
'site_xpos': d.site_xpos.shape,
|
|
'solver_niter': d._impl.solver_niter.shape,
|
|
'subtree_angmom': d._impl.subtree_angmom.shape,
|
|
'subtree_com': d.subtree_com.shape,
|
|
'subtree_linvel': d._impl.subtree_linvel.shape,
|
|
'ten_J': d._impl.ten_J.shape,
|
|
'ten_length': d.ten_length.shape,
|
|
'ten_velocity': d._impl.ten_velocity.shape,
|
|
'ten_wrapadr': d._impl.ten_wrapadr.shape,
|
|
'ten_wrapnum': d._impl.ten_wrapnum.shape,
|
|
'time': d.time.shape,
|
|
'tree_asleep': d._impl.tree_asleep.shape,
|
|
'tree_awake': d._impl.tree_awake.shape,
|
|
'tree_island': d._impl.tree_island.shape,
|
|
'wrap_obj': d._impl.wrap_obj.shape,
|
|
'wrap_xpos': d._impl.wrap_xpos.shape,
|
|
'xanchor': d.xanchor.shape,
|
|
'xaxis': d.xaxis.shape,
|
|
'ximat': d.ximat.shape,
|
|
'xipos': d.xipos.shape,
|
|
'xmat': d.xmat.shape,
|
|
'xpos': d.xpos.shape,
|
|
'xquat': d.xquat.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,
|
|
'efc__D': d._impl.efc__D.shape,
|
|
'efc__J': d._impl.efc__J.shape,
|
|
'efc__J_colind': d._impl.efc__J_colind.shape,
|
|
'efc__J_rowadr': d._impl.efc__J_rowadr.shape,
|
|
'efc__J_rownnz': d._impl.efc__J_rownnz.shape,
|
|
'efc__Jqvel': d._impl.efc__Jqvel.shape,
|
|
'efc__Ma': d._impl.efc__Ma.shape,
|
|
'efc__aref': d._impl.efc__aref.shape,
|
|
'efc__force': d._impl.efc__force.shape,
|
|
'efc__frictionloss': d._impl.efc__frictionloss.shape,
|
|
'efc__iD': d._impl.efc__iD.shape,
|
|
'efc__iJ': d._impl.efc__iJ.shape,
|
|
'efc__iJ_colind': d._impl.efc__iJ_colind.shape,
|
|
'efc__iJ_rowadr': d._impl.efc__iJ_rowadr.shape,
|
|
'efc__iJ_rownnz': d._impl.efc__iJ_rownnz.shape,
|
|
'efc__iaref': d._impl.efc__iaref.shape,
|
|
'efc__id': d._impl.efc__id.shape,
|
|
'efc__iforce': d._impl.efc__iforce.shape,
|
|
'efc__ifrictionloss': d._impl.efc__ifrictionloss.shape,
|
|
'efc__iid': d._impl.efc__iid.shape,
|
|
'efc__island': d._impl.efc__island.shape,
|
|
'efc__istate': d._impl.efc__istate.shape,
|
|
'efc__itype': d._impl.efc__itype.shape,
|
|
'efc__margin': d._impl.efc__margin.shape,
|
|
'efc__pos': d._impl.efc__pos.shape,
|
|
'efc__state': d._impl.efc__state.shape,
|
|
'efc__type': d._impl.efc__type.shape,
|
|
'efc__vel': d._impl.efc__vel.shape,
|
|
}
|
|
jf = ffi.jax_callable_variadic_tuple(
|
|
_step_shim,
|
|
num_outputs=148,
|
|
output_dims=output_dims,
|
|
vmap_method=None,
|
|
in_out_argnames=set([
|
|
'M',
|
|
'act',
|
|
'act_dot',
|
|
'actuator_force',
|
|
'actuator_length',
|
|
'actuator_moment',
|
|
'actuator_velocity',
|
|
'body_awake',
|
|
'body_awake_ind',
|
|
'cacc',
|
|
'cam_xmat',
|
|
'cam_xpos',
|
|
'cdof',
|
|
'cdof_dot',
|
|
'cfrc_ext',
|
|
'cfrc_int',
|
|
'cinert',
|
|
'crb',
|
|
'cvel',
|
|
'dof_awake_ind',
|
|
'dof_island',
|
|
'dof_islandid',
|
|
'efc_islandid',
|
|
'energy',
|
|
'flexedge_J',
|
|
'flexedge_length',
|
|
'flexedge_velocity',
|
|
'flexvert_xpos',
|
|
'geom_xmat',
|
|
'geom_xpos',
|
|
'history',
|
|
'iqacc',
|
|
'iqacc_smooth',
|
|
'iqfrc_constraint',
|
|
'iqfrc_smooth',
|
|
'island_dofadr',
|
|
'island_efcadr',
|
|
'island_idofadr',
|
|
'island_ne',
|
|
'island_nefc',
|
|
'island_nf',
|
|
'island_nv',
|
|
'light_xdir',
|
|
'light_xpos',
|
|
'map_dof2idof',
|
|
'map_efc2iefc',
|
|
'map_idof2dof',
|
|
'map_iefc2efc',
|
|
'moment_colind',
|
|
'moment_rowadr',
|
|
'moment_rownnz',
|
|
'nacon',
|
|
'nbody_awake',
|
|
'ncollision',
|
|
'ne',
|
|
'nefc',
|
|
'nf',
|
|
'nidof',
|
|
'nisland',
|
|
'nl',
|
|
'ntree_awake',
|
|
'nv_awake',
|
|
'qLD',
|
|
'qLDiagInv',
|
|
'qLU',
|
|
'qacc',
|
|
'qacc_smooth',
|
|
'qacc_warmstart',
|
|
'qfrc_actuator',
|
|
'qfrc_bias',
|
|
'qfrc_constraint',
|
|
'qfrc_damper',
|
|
'qfrc_fluid',
|
|
'qfrc_gravcomp',
|
|
'qfrc_passive',
|
|
'qfrc_smooth',
|
|
'qfrc_spring',
|
|
'qpos',
|
|
'qvel',
|
|
'sensordata',
|
|
'site_xmat',
|
|
'site_xpos',
|
|
'solver_niter',
|
|
'subtree_angmom',
|
|
'subtree_com',
|
|
'subtree_linvel',
|
|
'ten_J',
|
|
'ten_length',
|
|
'ten_velocity',
|
|
'ten_wrapadr',
|
|
'ten_wrapnum',
|
|
'time',
|
|
'tree_asleep',
|
|
'tree_awake',
|
|
'tree_island',
|
|
'wrap_obj',
|
|
'wrap_xpos',
|
|
'xanchor',
|
|
'xaxis',
|
|
'ximat',
|
|
'xipos',
|
|
'xmat',
|
|
'xpos',
|
|
'xquat',
|
|
'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',
|
|
'efc__D',
|
|
'efc__J',
|
|
'efc__J_colind',
|
|
'efc__J_rowadr',
|
|
'efc__J_rownnz',
|
|
'efc__Jqvel',
|
|
'efc__Ma',
|
|
'efc__aref',
|
|
'efc__force',
|
|
'efc__frictionloss',
|
|
'efc__iD',
|
|
'efc__iJ',
|
|
'efc__iJ_colind',
|
|
'efc__iJ_rowadr',
|
|
'efc__iJ_rownnz',
|
|
'efc__iaref',
|
|
'efc__id',
|
|
'efc__iforce',
|
|
'efc__ifrictionloss',
|
|
'efc__iid',
|
|
'efc__island',
|
|
'efc__istate',
|
|
'efc__itype',
|
|
'efc__margin',
|
|
'efc__pos',
|
|
'efc__state',
|
|
'efc__type',
|
|
'efc__vel',
|
|
]),
|
|
stage_in_argnames=set([
|
|
'act',
|
|
'act_dot',
|
|
'actuator_acc0',
|
|
'actuator_actrange',
|
|
'actuator_biasprm',
|
|
'actuator_cranklength',
|
|
'actuator_ctrlrange',
|
|
'actuator_dynprm',
|
|
'actuator_force',
|
|
'actuator_forcerange',
|
|
'actuator_gainprm',
|
|
'actuator_gear',
|
|
'actuator_length',
|
|
'actuator_lengthrange',
|
|
'body_gravcomp',
|
|
'body_inertia',
|
|
'body_invweight0',
|
|
'body_ipos',
|
|
'body_iquat',
|
|
'body_mass',
|
|
'body_pos',
|
|
'body_quat',
|
|
'body_subtreemass',
|
|
'cam_fovy',
|
|
'cam_intrinsic',
|
|
'cam_mat0',
|
|
'cam_pos',
|
|
'cam_pos0',
|
|
'cam_poscom0',
|
|
'cam_quat',
|
|
'cam_xmat',
|
|
'cam_xpos',
|
|
'cdof',
|
|
'cdof_dot',
|
|
'ctrl',
|
|
'cvel',
|
|
'dof_armature',
|
|
'dof_damping',
|
|
'dof_dampingpoly',
|
|
'dof_frictionloss',
|
|
'dof_invweight0',
|
|
'dof_solimp',
|
|
'dof_solref',
|
|
'eq_active',
|
|
'eq_data',
|
|
'eq_solimp',
|
|
'eq_solref',
|
|
'geom_aabb',
|
|
'geom_friction',
|
|
'geom_gap',
|
|
'geom_margin',
|
|
'geom_matid',
|
|
'geom_pos',
|
|
'geom_quat',
|
|
'geom_rbound',
|
|
'geom_rgba',
|
|
'geom_size',
|
|
'geom_solimp',
|
|
'geom_solmix',
|
|
'geom_solref',
|
|
'geom_xmat',
|
|
'geom_xpos',
|
|
'hfield_data',
|
|
'history',
|
|
'jnt_actfrcrange',
|
|
'jnt_axis',
|
|
'jnt_margin',
|
|
'jnt_pos',
|
|
'jnt_range',
|
|
'jnt_solimp',
|
|
'jnt_solref',
|
|
'jnt_stiffness',
|
|
'jnt_stiffnesspoly',
|
|
'light_dir',
|
|
'light_dir0',
|
|
'light_pos',
|
|
'light_pos0',
|
|
'light_poscom0',
|
|
'mat_rgba',
|
|
'mocap_pos',
|
|
'mocap_quat',
|
|
'opt__density',
|
|
'opt__gravity',
|
|
'opt__ls_tolerance',
|
|
'opt__magnetic',
|
|
'opt__timestep',
|
|
'opt__tolerance',
|
|
'opt__viscosity',
|
|
'opt__wind',
|
|
'pair_friction',
|
|
'pair_gap',
|
|
'pair_margin',
|
|
'pair_solimp',
|
|
'pair_solref',
|
|
'pair_solreffriction',
|
|
'qacc',
|
|
'qacc_smooth',
|
|
'qacc_warmstart',
|
|
'qfrc_actuator',
|
|
'qfrc_applied',
|
|
'qfrc_bias',
|
|
'qfrc_constraint',
|
|
'qfrc_fluid',
|
|
'qfrc_gravcomp',
|
|
'qfrc_passive',
|
|
'qfrc_smooth',
|
|
'qpos',
|
|
'qpos0',
|
|
'qpos_spring',
|
|
'qvel',
|
|
'sensordata',
|
|
'site_pos',
|
|
'site_quat',
|
|
'site_xmat',
|
|
'site_xpos',
|
|
'subtree_com',
|
|
'ten_length',
|
|
'tendon_actfrcrange',
|
|
'tendon_armature',
|
|
'tendon_damping',
|
|
'tendon_dampingpoly',
|
|
'tendon_frictionloss',
|
|
'tendon_invweight0',
|
|
'tendon_length0',
|
|
'tendon_lengthspring',
|
|
'tendon_margin',
|
|
'tendon_range',
|
|
'tendon_solimp_fri',
|
|
'tendon_solimp_lim',
|
|
'tendon_solref_fri',
|
|
'tendon_solref_lim',
|
|
'tendon_stiffness',
|
|
'tendon_stiffnesspoly',
|
|
'time',
|
|
'xanchor',
|
|
'xaxis',
|
|
'xfrc_applied',
|
|
'ximat',
|
|
'xipos',
|
|
'xmat',
|
|
'xpos',
|
|
'xquat',
|
|
]),
|
|
stage_out_argnames=set([
|
|
'act',
|
|
'act_dot',
|
|
'actuator_force',
|
|
'actuator_length',
|
|
'cam_xmat',
|
|
'cam_xpos',
|
|
'cdof',
|
|
'cdof_dot',
|
|
'cvel',
|
|
'geom_xmat',
|
|
'geom_xpos',
|
|
'history',
|
|
'qacc',
|
|
'qacc_smooth',
|
|
'qacc_warmstart',
|
|
'qfrc_actuator',
|
|
'qfrc_bias',
|
|
'qfrc_constraint',
|
|
'qfrc_fluid',
|
|
'qfrc_gravcomp',
|
|
'qfrc_passive',
|
|
'qfrc_smooth',
|
|
'qpos',
|
|
'qvel',
|
|
'sensordata',
|
|
'site_xmat',
|
|
'site_xpos',
|
|
'subtree_com',
|
|
'ten_length',
|
|
'time',
|
|
'xanchor',
|
|
'xaxis',
|
|
'ximat',
|
|
'xipos',
|
|
'xmat',
|
|
'xpos',
|
|
'xquat',
|
|
]),
|
|
graph_mode=m.opt._impl.graph_mode,
|
|
has_side_effect=False,
|
|
)
|
|
out = jf(
|
|
d.qpos.shape[0],
|
|
m._impl.D_colind,
|
|
m._impl.D_diag,
|
|
m._impl.D_rowadr,
|
|
m._impl.D_rownnz,
|
|
m._impl.M_elemid,
|
|
m._impl.M_fullm_i,
|
|
m._impl.M_fullm_j,
|
|
m._impl.M_mulm_col,
|
|
m._impl.M_mulm_madr,
|
|
m._impl.M_mulm_rowadr,
|
|
m._impl.M_rowadr,
|
|
m._impl.M_rownnz,
|
|
m._impl.M_tiles,
|
|
m.actuator_acc0,
|
|
m.actuator_actadr,
|
|
m.actuator_actearly,
|
|
m.actuator_actlimited,
|
|
m.actuator_actnum,
|
|
m.actuator_actrange,
|
|
m.actuator_biasprm,
|
|
m.actuator_biastype,
|
|
m.actuator_cranklength,
|
|
m.actuator_ctrllimited,
|
|
m.actuator_ctrlrange,
|
|
m._impl.actuator_delay,
|
|
m.actuator_dynprm,
|
|
m.actuator_dyntype,
|
|
m.actuator_forcelimited,
|
|
m.actuator_forcerange,
|
|
m.actuator_gainprm,
|
|
m.actuator_gaintype,
|
|
m.actuator_gear,
|
|
m._impl.actuator_history,
|
|
m._impl.actuator_historyadr,
|
|
m.actuator_lengthrange,
|
|
m.actuator_trnid,
|
|
m.actuator_trntype,
|
|
m._impl.actuator_trntype_body_adr,
|
|
m._impl.block_dim,
|
|
m._impl.body_branch_start,
|
|
m._impl.body_branches,
|
|
m.body_dofadr,
|
|
m.body_dofnum,
|
|
m._impl.body_fluid_ellipsoid,
|
|
m.body_geomadr,
|
|
m.body_geomnum,
|
|
m.body_gravcomp,
|
|
m.body_inertia,
|
|
m.body_invweight0,
|
|
m.body_ipos,
|
|
m.body_iquat,
|
|
m._impl.body_isdofancestor,
|
|
m.body_jntadr,
|
|
m.body_jntnum,
|
|
m.body_mass,
|
|
m.body_mocapid,
|
|
m.body_parentid,
|
|
m.body_pos,
|
|
m.body_quat,
|
|
m.body_rootid,
|
|
m.body_subtreemass,
|
|
m._impl.body_tree,
|
|
m.body_treeid,
|
|
m.body_weldid,
|
|
m.cam_bodyid,
|
|
m.cam_fovy,
|
|
m.cam_intrinsic,
|
|
m.cam_mat0,
|
|
m.cam_mode,
|
|
m.cam_pos,
|
|
m.cam_pos0,
|
|
m.cam_poscom0,
|
|
m.cam_quat,
|
|
m.cam_resolution,
|
|
m.cam_sensorsize,
|
|
m.cam_targetbodyid,
|
|
m.dof_armature,
|
|
m.dof_bodyid,
|
|
m.dof_damping,
|
|
m.dof_dampingpoly,
|
|
m.dof_frictionloss,
|
|
m.dof_invweight0,
|
|
m.dof_jntid,
|
|
m._impl.dof_length,
|
|
m.dof_parentid,
|
|
m.dof_solimp,
|
|
m.dof_solref,
|
|
m.dof_treeid,
|
|
m._impl.dof_tri_col,
|
|
m._impl.dof_tri_row,
|
|
m._impl.eq_connect_adr,
|
|
m.eq_data,
|
|
m._impl.eq_flex_adr,
|
|
m._impl.eq_jnt_adr,
|
|
m.eq_obj1id,
|
|
m.eq_obj2id,
|
|
m.eq_objtype,
|
|
m.eq_solimp,
|
|
m.eq_solref,
|
|
m._impl.eq_ten_adr,
|
|
m.eq_type,
|
|
m._impl.eq_wld_adr,
|
|
m._impl.flex_bending,
|
|
m._impl.flex_bendingadr,
|
|
m._impl.flex_centered,
|
|
m._impl.flex_conaffinity,
|
|
m._impl.flex_condim,
|
|
m._impl.flex_contype,
|
|
m._impl.flex_damping,
|
|
m._impl.flex_dim,
|
|
m._impl.flex_edge,
|
|
m._impl.flex_edgeadr,
|
|
m._impl.flex_edgeflap,
|
|
m._impl.flex_edgenum,
|
|
m._impl.flex_elem,
|
|
m._impl.flex_elemadr,
|
|
m._impl.flex_elemdataadr,
|
|
m._impl.flex_elemedge,
|
|
m._impl.flex_elemedgeadr,
|
|
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._impl.flex_stiffness,
|
|
m._impl.flex_stiffnessadr,
|
|
m._impl.flex_vert,
|
|
m.flex_vertadr,
|
|
m._impl.flex_vertbodyid,
|
|
m._impl.flex_vertflexid,
|
|
m.flex_vertnum,
|
|
m._impl.flexedge_J_colind,
|
|
m._impl.flexedge_J_rowadr,
|
|
m._impl.flexedge_J_rownnz,
|
|
m._impl.flexedge_invweight0,
|
|
m._impl.flexedge_length0,
|
|
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_fluid,
|
|
m.geom_friction,
|
|
m.geom_gap,
|
|
m.geom_group,
|
|
m.geom_margin,
|
|
m.geom_matid,
|
|
m._impl.geom_pair_type_count,
|
|
m._impl.geom_plugin_index,
|
|
m.geom_pos,
|
|
m.geom_priority,
|
|
m.geom_quat,
|
|
m.geom_rbound,
|
|
m.geom_rgba,
|
|
m.geom_size,
|
|
m.geom_solimp,
|
|
m.geom_solmix,
|
|
m.geom_solref,
|
|
m.geom_type,
|
|
m._impl.has_fluid,
|
|
m._impl.has_sdf_geom,
|
|
m.hfield_adr,
|
|
m.hfield_data,
|
|
m.hfield_ncol,
|
|
m.hfield_nrow,
|
|
m.hfield_size,
|
|
m._impl.is_sparse,
|
|
m.jnt_actfrclimited,
|
|
m.jnt_actfrcrange,
|
|
m.jnt_actgravcomp,
|
|
m.jnt_axis,
|
|
m.jnt_bodyid,
|
|
m.jnt_dofadr,
|
|
m._impl.jnt_limited_ball_adr,
|
|
m._impl.jnt_limited_slide_hinge_adr,
|
|
m.jnt_margin,
|
|
m.jnt_pos,
|
|
m.jnt_qposadr,
|
|
m.jnt_range,
|
|
m.jnt_solimp,
|
|
m.jnt_solref,
|
|
m.jnt_stiffness,
|
|
m.jnt_stiffnesspoly,
|
|
m.jnt_type,
|
|
m._impl.light_bodyid,
|
|
m.light_dir,
|
|
m.light_dir0,
|
|
m.light_mode,
|
|
m.light_pos,
|
|
m.light_pos0,
|
|
m.light_poscom0,
|
|
m._impl.light_targetbodyid,
|
|
m._impl.mapM2D,
|
|
m.mat_rgba,
|
|
m._impl.max_ten_J_rownnz,
|
|
m.mesh_face,
|
|
m.mesh_faceadr,
|
|
m.mesh_graph,
|
|
m.mesh_graphadr,
|
|
m.mesh_normal,
|
|
m.mesh_normaladr,
|
|
m.mesh_normalnum,
|
|
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_quat,
|
|
m.mesh_vert,
|
|
m.mesh_vertadr,
|
|
m.mesh_vertnum,
|
|
m.nC,
|
|
m.nD,
|
|
m.nJten,
|
|
m.na,
|
|
m._impl.nacttrnbody,
|
|
m.nbody,
|
|
m._impl.nbranch,
|
|
m.ncam,
|
|
m.neq,
|
|
m.nflex,
|
|
m._impl.nflexedge,
|
|
m._impl.nflexelem,
|
|
m._impl.nflexshelldata,
|
|
m._impl.nflexvert,
|
|
m.ngeom,
|
|
m.ngravcomp,
|
|
m.nhistory,
|
|
m.njnt,
|
|
m.nlight,
|
|
m._impl.nmaxcondim,
|
|
m._impl.nmaxmeshdeg,
|
|
m._impl.nmaxpolygon,
|
|
m._impl.nmaxpyramid,
|
|
m.nmeshface,
|
|
m._impl.nrangefinder,
|
|
m._impl.nsensorcollision,
|
|
m._impl.nsensorcontact,
|
|
m._impl.nsensortaxel,
|
|
m.nsite,
|
|
m.ntendon,
|
|
m._impl.ntree,
|
|
m.nu,
|
|
m.nv,
|
|
m._impl.nv_pad,
|
|
m.nwrap,
|
|
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._impl.qD_fullm_i,
|
|
m._impl.qD_fullm_j,
|
|
m._impl.qLD_all_updates,
|
|
m._impl.qLD_level_offsets,
|
|
m._impl.qLD_updates,
|
|
m.qpos0,
|
|
m.qpos_spring,
|
|
m._impl.rangefinder_sensor_adr,
|
|
m._impl.sensor_acc_adr,
|
|
m.sensor_adr,
|
|
m._impl.sensor_adr_to_contact_adr,
|
|
m._impl.sensor_contact_adr,
|
|
m.sensor_cutoff,
|
|
m.sensor_datatype,
|
|
m._impl.sensor_delay,
|
|
m.sensor_dim,
|
|
m._impl.sensor_e_kinetic,
|
|
m._impl.sensor_e_potential,
|
|
m._impl.sensor_history,
|
|
m._impl.sensor_historyadr,
|
|
m._impl.sensor_interval,
|
|
m.sensor_intprm,
|
|
m._impl.sensor_limitfrc_adr,
|
|
m._impl.sensor_limitpos_adr,
|
|
m._impl.sensor_limitvel_adr,
|
|
m.sensor_objid,
|
|
m.sensor_objtype,
|
|
m._impl.sensor_pos_adr,
|
|
m._impl.sensor_rangefinder_adr,
|
|
m._impl.sensor_rangefinder_bodyid,
|
|
m.sensor_refid,
|
|
m.sensor_reftype,
|
|
m._impl.sensor_rne_postconstraint,
|
|
m._impl.sensor_subtree_vel,
|
|
m._impl.sensor_tendonactfrc_adr,
|
|
m._impl.sensor_touch_adr,
|
|
m.sensor_type,
|
|
m._impl.sensor_vel_adr,
|
|
m.site_bodyid,
|
|
m.site_pos,
|
|
m.site_quat,
|
|
m.site_size,
|
|
m.site_type,
|
|
m._impl.taxel_sensorid,
|
|
m._impl.taxel_vertadr,
|
|
m._impl.ten_J_colind,
|
|
m._impl.ten_J_rowadr,
|
|
m._impl.ten_J_rownnz,
|
|
m.tendon_actfrclimited,
|
|
m.tendon_actfrcrange,
|
|
m.tendon_adr,
|
|
m.tendon_armature,
|
|
m.tendon_damping,
|
|
m.tendon_dampingpoly,
|
|
m.tendon_frictionloss,
|
|
m._impl.tendon_geom_adr,
|
|
m.tendon_invweight0,
|
|
m._impl.tendon_jnt_adr,
|
|
m.tendon_length0,
|
|
m.tendon_lengthspring,
|
|
m.tendon_limited,
|
|
m._impl.tendon_limited_adr,
|
|
m.tendon_margin,
|
|
m.tendon_num,
|
|
m.tendon_range,
|
|
m._impl.tendon_site_pair_adr,
|
|
m.tendon_solimp_fri,
|
|
m.tendon_solimp_lim,
|
|
m.tendon_solref_fri,
|
|
m.tendon_solref_lim,
|
|
m.tendon_stiffness,
|
|
m.tendon_stiffnesspoly,
|
|
m._impl.tree_dofadr,
|
|
m._impl.tree_dofnum,
|
|
m._impl.tree_sleep_policy,
|
|
m._impl.wrap_geom_adr,
|
|
m._impl.wrap_jnt_adr,
|
|
m.wrap_objid,
|
|
m.wrap_prm,
|
|
m._impl.wrap_pulley_scale,
|
|
m._impl.wrap_site_pair_adr,
|
|
m.wrap_type,
|
|
m.opt._impl.broadphase,
|
|
m.opt._impl.broadphase_filter,
|
|
m.opt._impl.ccd_iterations,
|
|
m.opt._impl.ccd_tolerance,
|
|
m.opt.cone,
|
|
m.opt._impl.contact_sensor_maxmatch,
|
|
m.opt.density,
|
|
m.opt.disableflags,
|
|
m.opt.enableflags,
|
|
m.opt._impl.graph_conditional,
|
|
m.opt.gravity,
|
|
m.opt._impl.impratio_invsqrt,
|
|
m.opt.integrator,
|
|
m.opt.iterations,
|
|
m.opt.ls_iterations,
|
|
m.opt._impl.ls_parallel,
|
|
m.opt._impl.ls_parallel_min_step,
|
|
m.opt.ls_tolerance,
|
|
m.opt.magnetic,
|
|
m.opt._impl.run_collision_detection,
|
|
m.opt._impl.sdf_initpoints,
|
|
m.opt._impl.sdf_iterations,
|
|
m.opt._impl.sleep_tolerance,
|
|
m.opt.solver,
|
|
m.opt.timestep,
|
|
m.opt.tolerance,
|
|
m.opt.viscosity,
|
|
m.opt.wind,
|
|
m.stat.meaninertia,
|
|
d._impl.naccdmax,
|
|
d._impl.naconmax,
|
|
d._impl.njmax,
|
|
d._impl.njmax_nnz,
|
|
d._impl.M,
|
|
d.act,
|
|
d.act_dot,
|
|
d.actuator_force,
|
|
d.actuator_length,
|
|
d._impl.actuator_moment,
|
|
d._impl.actuator_velocity,
|
|
d._impl.body_awake,
|
|
d._impl.body_awake_ind,
|
|
d._impl.cacc,
|
|
d.cam_xmat,
|
|
d.cam_xpos,
|
|
d.cdof,
|
|
d.cdof_dot,
|
|
d._impl.cfrc_ext,
|
|
d._impl.cfrc_int,
|
|
d._impl.cinert,
|
|
d._impl.crb,
|
|
d.ctrl,
|
|
d.cvel,
|
|
d._impl.dof_awake_ind,
|
|
d._impl.dof_island,
|
|
d._impl.dof_islandid,
|
|
d._impl.efc_islandid,
|
|
d._impl.energy,
|
|
d.eq_active,
|
|
d._impl.flexedge_J,
|
|
d._impl.flexedge_length,
|
|
d._impl.flexedge_velocity,
|
|
d._impl.flexvert_xpos,
|
|
d.geom_xmat,
|
|
d.geom_xpos,
|
|
d.history,
|
|
d._impl.iqacc,
|
|
d._impl.iqacc_smooth,
|
|
d._impl.iqfrc_constraint,
|
|
d._impl.iqfrc_smooth,
|
|
d._impl.island_dofadr,
|
|
d._impl.island_efcadr,
|
|
d._impl.island_idofadr,
|
|
d._impl.island_ne,
|
|
d._impl.island_nefc,
|
|
d._impl.island_nf,
|
|
d._impl.island_nv,
|
|
d._impl.light_xdir,
|
|
d._impl.light_xpos,
|
|
d._impl.map_dof2idof,
|
|
d._impl.map_efc2iefc,
|
|
d._impl.map_idof2dof,
|
|
d._impl.map_iefc2efc,
|
|
d.mocap_pos,
|
|
d.mocap_quat,
|
|
d._impl.moment_colind,
|
|
d._impl.moment_rowadr,
|
|
d._impl.moment_rownnz,
|
|
d._impl.nacon,
|
|
d._impl.nbody_awake,
|
|
d._impl.ncollision,
|
|
d._impl.ne,
|
|
d._impl.nefc,
|
|
d._impl.nf,
|
|
d._impl.nidof,
|
|
d._impl.nisland,
|
|
d._impl.nl,
|
|
d._impl.ntree_awake,
|
|
d._impl.nv_awake,
|
|
d._impl.qLD,
|
|
d._impl.qLDiagInv,
|
|
d._impl.qLU,
|
|
d.qacc,
|
|
d.qacc_smooth,
|
|
d.qacc_warmstart,
|
|
d.qfrc_actuator,
|
|
d.qfrc_applied,
|
|
d.qfrc_bias,
|
|
d.qfrc_constraint,
|
|
d._impl.qfrc_damper,
|
|
d.qfrc_fluid,
|
|
d.qfrc_gravcomp,
|
|
d.qfrc_passive,
|
|
d.qfrc_smooth,
|
|
d._impl.qfrc_spring,
|
|
d.qpos,
|
|
d.qvel,
|
|
d.sensordata,
|
|
d.site_xmat,
|
|
d.site_xpos,
|
|
d._impl.solver_niter,
|
|
d._impl.subtree_angmom,
|
|
d.subtree_com,
|
|
d._impl.subtree_linvel,
|
|
d._impl.ten_J,
|
|
d.ten_length,
|
|
d._impl.ten_velocity,
|
|
d._impl.ten_wrapadr,
|
|
d._impl.ten_wrapnum,
|
|
d.time,
|
|
d._impl.tree_asleep,
|
|
d._impl.tree_awake,
|
|
d._impl.tree_island,
|
|
d._impl.wrap_obj,
|
|
d._impl.wrap_xpos,
|
|
d.xanchor,
|
|
d.xaxis,
|
|
d.xfrc_applied,
|
|
d.ximat,
|
|
d.xipos,
|
|
d.xmat,
|
|
d.xpos,
|
|
d.xquat,
|
|
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._impl.efc__D,
|
|
d._impl.efc__J,
|
|
d._impl.efc__J_colind,
|
|
d._impl.efc__J_rowadr,
|
|
d._impl.efc__J_rownnz,
|
|
d._impl.efc__Jqvel,
|
|
d._impl.efc__Ma,
|
|
d._impl.efc__aref,
|
|
d._impl.efc__force,
|
|
d._impl.efc__frictionloss,
|
|
d._impl.efc__iD,
|
|
d._impl.efc__iJ,
|
|
d._impl.efc__iJ_colind,
|
|
d._impl.efc__iJ_rowadr,
|
|
d._impl.efc__iJ_rownnz,
|
|
d._impl.efc__iaref,
|
|
d._impl.efc__id,
|
|
d._impl.efc__iforce,
|
|
d._impl.efc__ifrictionloss,
|
|
d._impl.efc__iid,
|
|
d._impl.efc__island,
|
|
d._impl.efc__istate,
|
|
d._impl.efc__itype,
|
|
d._impl.efc__margin,
|
|
d._impl.efc__pos,
|
|
d._impl.efc__state,
|
|
d._impl.efc__type,
|
|
d._impl.efc__vel,
|
|
)
|
|
d = d.tree_replace({
|
|
'_impl.M': out[0],
|
|
'act': out[1],
|
|
'act_dot': out[2],
|
|
'actuator_force': out[3],
|
|
'actuator_length': out[4],
|
|
'_impl.actuator_moment': out[5],
|
|
'_impl.actuator_velocity': out[6],
|
|
'_impl.body_awake': out[7],
|
|
'_impl.body_awake_ind': out[8],
|
|
'_impl.cacc': out[9],
|
|
'cam_xmat': out[10],
|
|
'cam_xpos': out[11],
|
|
'cdof': out[12],
|
|
'cdof_dot': out[13],
|
|
'_impl.cfrc_ext': out[14],
|
|
'_impl.cfrc_int': out[15],
|
|
'_impl.cinert': out[16],
|
|
'_impl.crb': out[17],
|
|
'cvel': out[18],
|
|
'_impl.dof_awake_ind': out[19],
|
|
'_impl.dof_island': out[20],
|
|
'_impl.dof_islandid': out[21],
|
|
'_impl.efc_islandid': out[22],
|
|
'_impl.energy': out[23],
|
|
'_impl.flexedge_J': out[24],
|
|
'_impl.flexedge_length': out[25],
|
|
'_impl.flexedge_velocity': out[26],
|
|
'_impl.flexvert_xpos': out[27],
|
|
'geom_xmat': out[28],
|
|
'geom_xpos': out[29],
|
|
'history': out[30],
|
|
'_impl.iqacc': out[31],
|
|
'_impl.iqacc_smooth': out[32],
|
|
'_impl.iqfrc_constraint': out[33],
|
|
'_impl.iqfrc_smooth': out[34],
|
|
'_impl.island_dofadr': out[35],
|
|
'_impl.island_efcadr': out[36],
|
|
'_impl.island_idofadr': out[37],
|
|
'_impl.island_ne': out[38],
|
|
'_impl.island_nefc': out[39],
|
|
'_impl.island_nf': out[40],
|
|
'_impl.island_nv': out[41],
|
|
'_impl.light_xdir': out[42],
|
|
'_impl.light_xpos': out[43],
|
|
'_impl.map_dof2idof': out[44],
|
|
'_impl.map_efc2iefc': out[45],
|
|
'_impl.map_idof2dof': out[46],
|
|
'_impl.map_iefc2efc': out[47],
|
|
'_impl.moment_colind': out[48],
|
|
'_impl.moment_rowadr': out[49],
|
|
'_impl.moment_rownnz': out[50],
|
|
'_impl.nacon': out[51],
|
|
'_impl.nbody_awake': out[52],
|
|
'_impl.ncollision': out[53],
|
|
'_impl.ne': out[54],
|
|
'_impl.nefc': out[55],
|
|
'_impl.nf': out[56],
|
|
'_impl.nidof': out[57],
|
|
'_impl.nisland': out[58],
|
|
'_impl.nl': out[59],
|
|
'_impl.ntree_awake': out[60],
|
|
'_impl.nv_awake': out[61],
|
|
'_impl.qLD': out[62],
|
|
'_impl.qLDiagInv': out[63],
|
|
'_impl.qLU': out[64],
|
|
'qacc': out[65],
|
|
'qacc_smooth': out[66],
|
|
'qacc_warmstart': out[67],
|
|
'qfrc_actuator': out[68],
|
|
'qfrc_bias': out[69],
|
|
'qfrc_constraint': out[70],
|
|
'_impl.qfrc_damper': out[71],
|
|
'qfrc_fluid': out[72],
|
|
'qfrc_gravcomp': out[73],
|
|
'qfrc_passive': out[74],
|
|
'qfrc_smooth': out[75],
|
|
'_impl.qfrc_spring': out[76],
|
|
'qpos': out[77],
|
|
'qvel': out[78],
|
|
'sensordata': out[79],
|
|
'site_xmat': out[80],
|
|
'site_xpos': out[81],
|
|
'_impl.solver_niter': out[82],
|
|
'_impl.subtree_angmom': out[83],
|
|
'subtree_com': out[84],
|
|
'_impl.subtree_linvel': out[85],
|
|
'_impl.ten_J': out[86],
|
|
'ten_length': out[87],
|
|
'_impl.ten_velocity': out[88],
|
|
'_impl.ten_wrapadr': out[89],
|
|
'_impl.ten_wrapnum': out[90],
|
|
'time': out[91],
|
|
'_impl.tree_asleep': out[92],
|
|
'_impl.tree_awake': out[93],
|
|
'_impl.tree_island': out[94],
|
|
'_impl.wrap_obj': out[95],
|
|
'_impl.wrap_xpos': out[96],
|
|
'xanchor': out[97],
|
|
'xaxis': out[98],
|
|
'ximat': out[99],
|
|
'xipos': out[100],
|
|
'xmat': out[101],
|
|
'xpos': out[102],
|
|
'xquat': out[103],
|
|
'_impl.contact__dim': out[104],
|
|
'_impl.contact__dist': out[105],
|
|
'_impl.contact__efc_address': out[106],
|
|
'_impl.contact__flex': out[107],
|
|
'_impl.contact__frame': out[108],
|
|
'_impl.contact__friction': out[109],
|
|
'_impl.contact__geom': out[110],
|
|
'_impl.contact__geomcollisionid': out[111],
|
|
'_impl.contact__includemargin': out[112],
|
|
'_impl.contact__pos': out[113],
|
|
'_impl.contact__solimp': out[114],
|
|
'_impl.contact__solref': out[115],
|
|
'_impl.contact__solreffriction': out[116],
|
|
'_impl.contact__type': out[117],
|
|
'_impl.contact__vert': out[118],
|
|
'_impl.contact__worldid': out[119],
|
|
'_impl.efc__D': out[120],
|
|
'_impl.efc__J': out[121],
|
|
'_impl.efc__J_colind': out[122],
|
|
'_impl.efc__J_rowadr': out[123],
|
|
'_impl.efc__J_rownnz': out[124],
|
|
'_impl.efc__Jqvel': out[125],
|
|
'_impl.efc__Ma': out[126],
|
|
'_impl.efc__aref': out[127],
|
|
'_impl.efc__force': out[128],
|
|
'_impl.efc__frictionloss': out[129],
|
|
'_impl.efc__iD': out[130],
|
|
'_impl.efc__iJ': out[131],
|
|
'_impl.efc__iJ_colind': out[132],
|
|
'_impl.efc__iJ_rowadr': out[133],
|
|
'_impl.efc__iJ_rownnz': out[134],
|
|
'_impl.efc__iaref': out[135],
|
|
'_impl.efc__id': out[136],
|
|
'_impl.efc__iforce': out[137],
|
|
'_impl.efc__ifrictionloss': out[138],
|
|
'_impl.efc__iid': out[139],
|
|
'_impl.efc__island': out[140],
|
|
'_impl.efc__istate': out[141],
|
|
'_impl.efc__itype': out[142],
|
|
'_impl.efc__margin': out[143],
|
|
'_impl.efc__pos': out[144],
|
|
'_impl.efc__state': out[145],
|
|
'_impl.efc__type': out[146],
|
|
'_impl.efc__vel': out[147],
|
|
})
|
|
return d
|
|
|
|
|
|
@jax.custom_batching.custom_vmap
|
|
@ffi.marshal_jax_warp_callable
|
|
def step(m: types.Model, d: types.Data):
|
|
return _step_jax_impl(m, d)
|
|
|
|
|
|
@step.def_vmap
|
|
@ffi.marshal_custom_vmap
|
|
def step_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
|
|
d = step(m, d)
|
|
return d, is_batched[1]
|