chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Import google-deepmind/mujoco_warp from GitHub.

Browse Source 
PiperOrigin-RevId: 927227210
Change-Id: Ic72c58524c8e91e6fff44b4cab7f4a0f7de2bb33
This commit is contained in:
Taylor Howell
2026-06-05 04:56:22 -07:00
committed by Copybara-Service
parent d80a847bae
commit 2a4af3246e
14 changed files with 3023 additions and 1203 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+132 -43
View File
@@ -13,7 +13,7 @@
# limitations under the License.
# ==============================================================================
from typing import Any
from typing import Any, Optional
import warp as wp
@@ -30,8 +30,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
from mujoco.mjx.third_party.mujoco_warp._src.types import CollisionType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import SleepState
from mujoco.mjx.third_party.mujoco_warp._src.types import mat23
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
@@ -77,14 +79,24 @@ MJ_COLLISION_TABLE = {
}
@wp.kernel
def _zero_nacon_ncollision(
# Data out:
nacon_out: wp.array[int],
ncollision_out: wp.array[int],
):
ncollision_out[0] = 0
nacon_out[0] = 0
@cache_kernel
def _zero_nacon_ncollision(enable_sleep: bool = False):
@wp.kernel(module="unique", enable_backward=False)
def zero_nacon_ncollision(
# In:
skip_in: wp.array[int],
# Data out:
nacon_out: wp.array[int],
ncollision_out: wp.array[int],
):
ncollision_out[0] = 0
if wp.static(enable_sleep):
if skip_in[0] != 0:
nacon_out[0] = 0
else:
nacon_out[0] = 0
return zero_nacon_ncollision
@wp.func
@@ -376,7 +388,7 @@ def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) ->
@cache_kernel
def _sap_project(opt_broadphase: int):
def _sap_project(opt_broadphase: int, enable_sleep: bool = False):
@wp.kernel(module="unique", enable_backward=False)
def sap_project(
# Model:
@@ -389,6 +401,7 @@ def _sap_project(opt_broadphase: int):
nworld_in: int,
# In:
direction_in: wp.vec3,
skip_in: wp.array[int],
# Out:
projection_lower_out: wp.array2d[float],
projection_upper_out: wp.array2d[float],
@@ -397,6 +410,10 @@ def _sap_project(opt_broadphase: int):
):
worldid, geomid = wp.tid()
if wp.static(enable_sleep):
if skip_in[0] == 0:
return
xpos = geom_xpos_in[worldid, geomid]
rbound = geom_rbound[worldid % geom_rbound.shape[0], geomid]
@@ -424,38 +441,51 @@ def _sap_project(opt_broadphase: int):
return sap_project
@wp.kernel
def _sap_range(
# Model:
ngeom: int,
# In:
projection_lower_in: wp.array2d[float],
projection_upper_in: wp.array2d[float],
sort_index_in: wp.array2d[int],
# Out:
range_out: wp.array2d[int],
):
worldid, geomid = wp.tid()
@cache_kernel
def _sap_range(enable_sleep: bool = False):
@wp.kernel(module="unique", enable_backward=False)
def sap_range(
# Model:
ngeom: int,
# In:
projection_lower_in: wp.array2d[float],
projection_upper_in: wp.array2d[float],
sort_index_in: wp.array2d[int],
skip_in: wp.array[int],
# Out:
range_out: wp.array2d[int],
):
worldid, geomid = wp.tid()
# current bounding geom
idx = sort_index_in[worldid, geomid]
if wp.static(enable_sleep):
if skip_in[0] == 0:
range_out[worldid, geomid] = 0
return
upper = projection_upper_in[worldid, idx]
# current bounding geom
idx = sort_index_in[worldid, geomid]
limit = _binary_search(projection_lower_in[worldid], upper, geomid + 1, ngeom)
limit = wp.min(ngeom - 1, limit)
upper = projection_upper_in[worldid, idx]
# range of geoms for the sweep and prune process
range_out[worldid, geomid] = limit - geomid
limit = _binary_search(projection_lower_in[worldid], upper, geomid + 1, ngeom)
limit = wp.min(ngeom - 1, limit)
# range of geoms for the sweep and prune process
range_out[worldid, geomid] = limit - geomid
return sap_range
@cache_kernel
def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
def _sap_broadphase(
opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int, enable_sleep: bool = False
):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
ngeom: int,
geom_type: wp.array[int],
geom_bodyid: wp.array[int],
geom_aabb: wp.array3d[wp.vec3],
geom_rbound: wp.array2d[float],
geom_margin: wp.array2d[float],
@@ -464,12 +494,14 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d[wp.mat33],
body_awake_in: wp.array2d[int],
nworld_in: int,
naconmax_in: int,
# In:
sort_index_in: wp.array2d[int],
cumulative_sum_in: wp.array[int],
nsweep_in: int,
skip_in: wp.array[int],
# Data out:
ncollision_out: wp.array[int],
# Out:
@@ -479,6 +511,10 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
):
worldgeomid = wp.tid()
if wp.static(enable_sleep):
if skip_in[0] == 0:
return
nworldgeom = nworld_in * ngeom
nworkpackages = cumulative_sum_in[nworldgeom - 1]
@@ -509,6 +545,16 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
if pairid[0] < -1 and pairid[1] < 0:
continue
if wp.static(enable_sleep):
b1 = geom_bodyid[geom1]
b2 = geom_bodyid[geom2]
s1 = body_awake_in[worldid, b1]
s2 = body_awake_in[worldid, b2]
if s1 == SleepState.ASLEEP and s2 == SleepState.ASLEEP:
continue
if (s1 == SleepState.ASLEEP and s2 == SleepState.STATIC) or (s2 == SleepState.ASLEEP and s1 == SleepState.STATIC):
continue
if (
wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin, ngeom_gap))(
geom_aabb, geom_rbound, geom_margin, geom_gap, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
@@ -560,7 +606,7 @@ def _segmented_sort(tile_size: int):
@event_scope
def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
def sap_broadphase(m: Model, d: Data, ctx: CollisionContext, skip: Optional[wp.array] = None):
"""Runs broadphase collision detection using a sweep-and-prune (SAP) algorithm.
This method is more efficient than the N-squared approach for large numbers of
@@ -578,6 +624,8 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
- `SAP_SEGMENTED`: Uses a segmented sort.
"""
nworldgeom = d.nworld * m.ngeom
skip_in = skip if skip is not None else wp.ones(1, dtype=int)
enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
# TODO(team): direction
@@ -593,9 +641,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
segmented_index = wp.empty(d.nworld + 1 if m.opt.broadphase == BroadphaseType.SAP_SEGMENTED else 0, dtype=int)
wp.launch(
kernel=_sap_project(m.opt.broadphase),
kernel=_sap_project(m.opt.broadphase, enable_sleep),
dim=(d.nworld, m.ngeom),
inputs=[m.ngeom, m.geom_rbound, m.geom_margin, m.geom_gap, d.geom_xpos, d.nworld, direction],
inputs=[m.ngeom, m.geom_rbound, m.geom_margin, m.geom_gap, d.geom_xpos, d.nworld, direction, skip_in],
outputs=[
projection_lower.reshape((-1, m.ngeom)),
projection_upper,
@@ -618,9 +666,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
)
wp.launch(
kernel=_sap_range,
kernel=_sap_range(enable_sleep),
dim=(d.nworld, m.ngeom),
inputs=[m.ngeom, projection_lower.reshape((-1, m.ngeom)), projection_upper, sort_index.reshape((-1, m.ngeom))],
inputs=[m.ngeom, projection_lower.reshape((-1, m.ngeom)), projection_upper, sort_index.reshape((-1, m.ngeom)), skip_in],
outputs=[range_],
)
@@ -632,12 +680,18 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
nsweep = 5 * nworldgeom
wp.launch(
kernel=_sap_broadphase(
m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
m.opt.broadphase_filter,
m.geom_aabb.shape[0],
m.geom_rbound.shape[0],
m.geom_margin.shape[0],
m.geom_gap.shape[0],
enable_sleep,
),
dim=nsweep,
inputs=[
m.ngeom,
m.geom_type,
m.geom_bodyid,
m.geom_aabb,
m.geom_rbound,
m.geom_margin,
@@ -645,22 +699,27 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
m.nxn_pairid,
d.geom_xpos,
d.geom_xmat,
d.body_awake,
d.nworld,
d.naconmax,
sort_index.reshape((-1, m.ngeom)),
cumulative_sum.reshape(-1),
nsweep,
skip_in,
],
outputs=[d.ncollision, ctx.collision_pair, ctx.collision_pairid, ctx.collision_worldid],
)
@cache_kernel
def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int):
def _nxn_broadphase(
opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: int, ngeom_margin: int, ngeom_gap: int, enable_sleep: bool = False
):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
geom_type: wp.array[int],
geom_bodyid: wp.array[int],
geom_aabb: wp.array3d[wp.vec3],
geom_rbound: wp.array2d[float],
geom_margin: wp.array2d[float],
@@ -670,7 +729,10 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d[wp.mat33],
body_awake_in: wp.array2d[int],
naconmax_in: int,
# In:
skip_in: wp.array[int],
# Data out:
ncollision_out: wp.array[int],
# Out:
@@ -680,10 +742,24 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
):
worldid, elementid = wp.tid()
if wp.static(enable_sleep):
if skip_in[0] == 0:
return
geom = nxn_geom_pair[elementid]
geom1 = geom[0]
geom2 = geom[1]
if wp.static(enable_sleep):
b1 = geom_bodyid[geom1]
b2 = geom_bodyid[geom2]
s1 = body_awake_in[worldid, b1]
s2 = body_awake_in[worldid, b2]
if s1 == SleepState.ASLEEP and s2 == SleepState.ASLEEP:
return
if (s1 == SleepState.ASLEEP and s2 == SleepState.STATIC) or (s2 == SleepState.ASLEEP and s1 == SleepState.STATIC):
return
if (
wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin, ngeom_gap))(
geom_aabb, geom_rbound, geom_margin, geom_gap, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid
@@ -708,7 +784,7 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
@event_scope
def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext, skip: Optional[wp.array] = None):
"""Runs broadphase collision detection using a brute-force N-squared approach.
This function iterates through a pre-filtered list of all possible geometry pairs and
@@ -721,13 +797,21 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
The initial list of pairs is filtered at model creation time to exclude pairs based on
`contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
"""
enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
skip_in = skip if skip is not None else wp.ones(1, dtype=int)
wp.launch(
_nxn_broadphase(
m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
m.opt.broadphase_filter,
m.geom_aabb.shape[0],
m.geom_rbound.shape[0],
m.geom_margin.shape[0],
m.geom_gap.shape[0],
enable_sleep,
),
dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
inputs=[
m.geom_type,
m.geom_bodyid,
m.geom_aabb,
m.geom_rbound,
m.geom_margin,
@@ -736,7 +820,9 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
m.nxn_pairid_filtered,
d.geom_xpos,
d.geom_xmat,
d.body_awake,
d.naconmax,
skip_in,
],
outputs=[
d.ncollision,
@@ -768,7 +854,7 @@ def _narrowphase(m: Model, d: Data, ctx: CollisionContext):
@event_scope
def collision(m: Model, d: Data):
def collision(m: Model, d: Data, skip: Optional[wp.array] = None):
"""Runs the full collision detection pipeline.
This function orchestrates the broadphase and narrowphase collision detection stages. It
@@ -789,15 +875,18 @@ def collision(m: Model, d: Data):
d.nacon.zero_()
return
# TODO(team): create context outside collision?
ctx = create_collision_context(d.naconmax)
skip_in = skip if skip is not None else wp.ones(1, dtype=int)
enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
# zero counters
wp.launch(_zero_nacon_ncollision, dim=1, outputs=[d.nacon, d.ncollision])
wp.launch(_zero_nacon_ncollision(enable_sleep), dim=1, inputs=[skip_in], outputs=[d.nacon, d.ncollision])
if m.opt.broadphase == BroadphaseType.NXN:
nxn_broadphase(m, d, ctx)
nxn_broadphase(m, d, ctx, skip)
else:
sap_broadphase(m, d, ctx)
sap_broadphase(m, d, ctx, skip)
_narrowphase(m, d, ctx)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+79 -20
View File
@@ -25,6 +25,7 @@ from mujoco.mjx.third_party.mujoco_warp._src import island
from mujoco.mjx.third_party.mujoco_warp._src import math
from mujoco.mjx.third_party.mujoco_warp._src import passive
from mujoco.mjx.third_party.mujoco_warp._src import sensor
from mujoco.mjx.third_party.mujoco_warp._src import sleep
from mujoco.mjx.third_party.mujoco_warp._src import smooth
from mujoco.mjx.third_party.mujoco_warp._src import solver
from mujoco.mjx.third_party.mujoco_warp._src import types
@@ -329,6 +330,11 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
wp.copy(d.qacc_warmstart, d.qacc)
if not (m.opt.disableflags & DisableBit.ISLAND) and (m.opt.enableflags & EnableBit.SLEEP):
sleep.sleep(m, d)
fwd_velocity(m, d)
sleep.update_sleep(m, d)
@wp.kernel
def _compute_damping_deriv(
@@ -658,13 +664,37 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
smooth.camlight(m, d)
smooth.flex(m, d)
smooth.tendon(m, d)
sleep_enabled = not (m.opt.disableflags & DisableBit.ISLAND) and (m.opt.enableflags & EnableBit.SLEEP)
if sleep_enabled and m.ntendon > 0:
sleep.wake_tendon(m, d)
sleep.update_sleep_trees(m, d)
smooth.crb(m, d)
smooth.tendon_armature(m, d)
if factorize:
smooth.factor_m(m, d)
if m.opt.run_collision_detection:
collision_driver.collision(m, d)
if sleep_enabled:
# pass 1
collision_driver.collision(m, d)
# check for newly awake
skip = wp.zeros(1, dtype=int)
sleep.wake_collision(m, d, skip)
sleep.update_sleep(m, d)
# pass 2: broadphase kernels early-return if skip[0] is 0
collision_driver.collision(m, d, skip)
else:
collision_driver.collision(m, d)
constraint.make_constraint(m, d)
if sleep_enabled:
if m.neq > 0:
sleep.wake_equality(m, d)
sleep.update_sleep(m, d)
if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
island.island(m, d)
smooth.transmission(m, d)
@@ -1242,23 +1272,39 @@ def fwd_actuation(m: Model, d: Data):
)
@wp.kernel
def _qfrc_smooth(
# Data in:
qfrc_applied_in: wp.array2d[float],
qfrc_bias_in: wp.array2d[float],
qfrc_passive_in: wp.array2d[float],
qfrc_actuator_in: wp.array2d[float],
# Data out:
qfrc_smooth_out: wp.array2d[float],
):
worldid, dofid = wp.tid()
qfrc_smooth_out[worldid, dofid] = (
qfrc_passive_in[worldid, dofid]
- qfrc_bias_in[worldid, dofid]
+ qfrc_actuator_in[worldid, dofid]
+ qfrc_applied_in[worldid, dofid]
)
@cache_kernel
def _qfrc_smooth(enable_sleep: bool):
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
body_treeid: wp.array[int],
dof_bodyid: wp.array[int],
# Data in:
qfrc_applied_in: wp.array2d[float],
tree_awake_in: wp.array2d[int],
qfrc_bias_in: wp.array2d[float],
qfrc_passive_in: wp.array2d[float],
qfrc_actuator_in: wp.array2d[float],
# Data out:
qfrc_smooth_out: wp.array2d[float],
):
worldid, dofid = wp.tid()
if wp.static(enable_sleep):
bodyid = dof_bodyid[dofid]
tree = body_treeid[bodyid]
if tree >= 0 and tree_awake_in[worldid, tree] == 0:
qfrc_smooth_out[worldid, dofid] = 0.0
return
qfrc_smooth_out[worldid, dofid] = (
qfrc_passive_in[worldid, dofid]
- qfrc_bias_in[worldid, dofid]
+ qfrc_actuator_in[worldid, dofid]
+ qfrc_applied_in[worldid, dofid]
)
return kernel
@event_scope
@@ -1270,10 +1316,19 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
d: The data object containing the current state and output arrays.
factorize: Flag to factorize inertia matrix.
"""
enable_sleep = bool(m.opt.enableflags & EnableBit.SLEEP)
wp.launch(
_qfrc_smooth,
_qfrc_smooth(enable_sleep),
dim=(d.nworld, m.nv),
inputs=[d.qfrc_applied, d.qfrc_bias, d.qfrc_passive, d.qfrc_actuator],
inputs=[
m.body_treeid,
m.dof_bodyid,
d.qfrc_applied,
d.tree_awake,
d.qfrc_bias,
d.qfrc_passive,
d.qfrc_actuator,
],
outputs=[d.qfrc_smooth],
)
xfrc_accumulate(m, d, d.qfrc_smooth)
@@ -1287,6 +1342,10 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
@event_scope
def forward(m: Model, d: Data):
"""Forward dynamics."""
if not (m.opt.disableflags & DisableBit.ISLAND) and (m.opt.enableflags & EnableBit.SLEEP):
sleep.wake(m, d)
sleep.update_sleep(m, d)
energy = m.opt.enableflags & EnableBit.ENERGY
fwd_position(m, d, factorize=False)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+125 -5
View File
@@ -154,6 +154,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
(mjm.geom_type, types.GeomType, mujoco.mjtGeom),
(mjm.sensor_type, types.SensorType, mujoco.mjtSensor),
(mjm.wrap_type, types.WrapType, mujoco.mjtWrap),
(mjm.tree_sleep_policy, types.SleepPolicy, mujoco.mjtSleepPolicy),
):
missing = ~np.isin(field, field_type)
if missing.any():
@@ -178,6 +179,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
if unsupported:
raise NotImplementedError(f"{mj_type(unsupported).name} is unsupported.")
if (mjm.opt.enableflags & mujoco.mjtEnableBit.mjENBL_SLEEP) and (mjm.eq_type == mujoco.mjtEq.mjEQ_FLEX).any():
raise NotImplementedError("Flex equality constraints are not supported with sleeping enabled.")
if mjm.opt.noslip_iterations > 0:
raise NotImplementedError(f"noslip solver not implemented.")
@@ -297,6 +301,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
m.nmaxpyramid = np.maximum(1, 2 * (m.nmaxcondim - 1))
m.has_sdf_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any()
m.block_dim = types.BlockDim()
if mjm.nv > 500:
m.block_dim.linesearch_iterative = 512
m.is_sparse = is_sparse(mjm)
m.has_fluid = mjm.opt.wind.any() or mjm.opt.density > 0 or mjm.opt.viscosity > 0
@@ -995,6 +1001,7 @@ def _allocate_island_arrays(
d.dof_island = wp.array(np.tile(mjd.dof_island, (nworld, 1 if island_enabled else 0)), dtype=int)
d.island_dofadr = wp.empty((nworld, ntree_size), dtype=int)
d.island_idofadr = wp.empty((nworld, ntree_size), dtype=int)
d.island_nv = wp.empty((nworld, ntree_size), dtype=int)
d.island_nefc = wp.empty((nworld, ntree_size), dtype=int)
d.island_ne = wp.empty((nworld, ntree_size), dtype=int)
@@ -1155,6 +1162,7 @@ def make_data(
"tree_island": None,
"dof_island": None,
"island_dofadr": None,
"island_idofadr": None,
"island_nv": None,
"island_nefc": None,
"island_ne": None,
@@ -1171,6 +1179,10 @@ def make_data(
"iqacc_smooth": None,
"iqfrc_smooth": None,
"iqfrc_constraint": None,
# sleep state: all trees start fully awake
"tree_asleep": wp.array(np.full((nworld, mjm.ntree), -(1 + types.MJ_MINAWAKE)), dtype=int),
"tree_awake": wp.array(np.ones((nworld, mjm.ntree)), dtype=int),
"body_awake": wp.array(np.ones((nworld, mjm.nbody)), dtype=int),
}
for f in dataclasses.fields(types.Data):
if f.name in d_kwargs:
@@ -1380,6 +1392,7 @@ def put_data(
"tree_island": None,
"dof_island": None,
"island_dofadr": None,
"island_idofadr": None,
"island_nv": None,
"island_nefc": None,
"island_ne": None,
@@ -1657,12 +1670,18 @@ def get_data_into(
# sensors
result.sensordata[:] = d.sensordata.numpy()[world_id]
# sleep
result.tree_asleep[:] = d.tree_asleep.numpy()[world_id]
result.tree_awake[:] = d.tree_awake.numpy()[world_id]
result.body_awake[:] = d.body_awake.numpy()[world_id]
# islands
nisland = d.nisland.numpy()[world_id]
result.nisland = nisland
if d.tree_island.shape[1] > 0 and nisland:
result.tree_island[:] = d.tree_island.numpy()[world_id]
result.dof_island[:] = d.dof_island.numpy()[world_id]
result.island_idofadr[:nisland] = d.island_idofadr.numpy()[world_id, :nisland]
result.island_dofadr[:nisland] = d.island_dofadr.numpy()[world_id, :nisland]
result.island_nv[:nisland] = d.island_nv.numpy()[world_id, :nisland]
result.island_nefc[:nisland] = d.island_nefc.numpy()[world_id, :nisland]
@@ -1743,6 +1762,8 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
nv: int,
nu: int,
na: int,
nbody: int,
ntree: int,
neq: int,
nsensordata: int,
qpos0: wp.array2d[float],
@@ -1757,6 +1778,9 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
nf_out: wp.array[int],
nl_out: wp.array[int],
nefc_out: wp.array[int],
ntree_awake_out: wp.array[int],
nbody_awake_out: wp.array[int],
nv_awake_out: wp.array[int],
time_out: wp.array[float],
energy_out: wp.array[wp.vec2],
qpos_out: wp.array2d[float],
@@ -1786,6 +1810,9 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
nefc_out[worldid] = 0
time_out[worldid] = 0.0
energy_out[worldid] = wp.vec2(0.0, 0.0)
ntree_awake_out[worldid] = ntree
nbody_awake_out[worldid] = nbody
nv_awake_out[worldid] = nv
qpos0_id = worldid % qpos0.shape[0]
for i in range(nq):
qpos_out[worldid, i] = qpos0[qpos0_id, i]
@@ -1882,6 +1909,47 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
contact_type_out[conid] = 0
contact_geomcollisionid_out[conid] = 0
@wp.kernel(module="unique", enable_backward=False)
def reset_sleep(
# Model:
nv: int,
nbody: int,
ntree: int,
body_mocapid: wp.array[int],
body_treeid: wp.array[int],
# In:
mj_minawake: int,
reset_in: wp.array[bool],
# Data out:
tree_asleep_out: wp.array2d[int],
tree_awake_out: wp.array2d[int],
body_awake_out: wp.array2d[int],
body_awake_ind_out: wp.array2d[int],
dof_awake_ind_out: wp.array2d[int],
):
worldid, elemid = wp.tid()
if wp.static(reset is not None):
if not reset_in[worldid]:
return
if elemid < ntree:
tree_asleep_out[worldid, elemid] = -(1 + mj_minawake)
tree_awake_out[worldid, elemid] = 1
if elemid < nbody:
if body_treeid[elemid] < 0:
if body_mocapid[elemid] >= 0:
body_awake_out[worldid, elemid] = int(types.SleepState.AWAKE)
else:
body_awake_out[worldid, elemid] = int(types.SleepState.STATIC)
else:
body_awake_out[worldid, elemid] = int(types.SleepState.AWAKE)
body_awake_ind_out[worldid, elemid] = elemid
if elemid < nv:
dof_awake_ind_out[worldid, elemid] = elemid
reset_input = reset or wp.ones(d.nworld, dtype=bool)
wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied])
@@ -1925,16 +1993,32 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
],
)
wp.launch(
reset_sleep,
dim=(d.nworld, max(m.ntree, m.nbody, m.nv)),
inputs=[m.nv, m.nbody, m.ntree, m.body_mocapid, m.body_treeid, types.MJ_MINAWAKE, reset_input],
outputs=[
d.tree_asleep,
d.tree_awake,
d.body_awake,
d.body_awake_ind,
d.dof_awake_ind,
],
)
wp.launch(
reset_nworld,
dim=d.nworld,
inputs=[m.nq, m.nv, m.nu, m.na, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld, reset_input],
inputs=[m.nq, m.nv, m.nu, m.na, m.nbody, m.ntree, m.neq, m.nsensordata, m.qpos0, m.eq_active0, d.nworld, reset_input],
outputs=[
d.solver_niter,
d.ne,
d.nf,
d.nl,
d.nefc,
d.ntree_awake,
d.nbody_awake,
d.nv_awake,
d.time,
d.energy,
d.qpos,
@@ -2785,7 +2869,7 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
"opt.graph_conditional",
"opt.contact_sensor_maxmatch",
}
mj_only_fields = {"opt.jacobian"}
mj_only_fields = {"opt.jacobian", "vis.quality.offsamples"}
if not isinstance(overrides, dict):
overrides_dict = {}
@@ -2951,10 +3035,14 @@ def create_render_context(
use_ambient_lighting: bool = True,
enabled_geom_groups: list[int] = [0, 1, 2],
cam_active: list[bool] | None = None,
background_color: tuple[float, float, float, float] = (0.1, 0.1, 0.2, 1.0),
flex_render_smooth: bool = True,
use_precomputed_rays: bool = True,
render_skybox: bool = False,
enable_backface_culling: bool = True,
enable_specular: bool = True,
enable_emission: bool = True,
enable_per_light_ambient: bool = True,
) -> types.RenderContext:
"""Creates a render context on device.
@@ -2969,8 +3057,9 @@ def create_render_context(
If None, uses the MuJoCo model values.
use_textures: Whether to use textures.
use_shadows: Whether to use shadows.
use_ambient_lighting: Whether to add the renderer's hemispheric ambient
lighting term before applying model lights.
use_ambient_lighting: Top-level ambient switch. When False, skips all
ambient contributions, including headlight ambient,
the no-light fallback, and per-light ambient.
enabled_geom_groups: The geom groups to render.
cam_active: List of booleans indicating which cameras to include in rendering.
If None, all cameras are included.
@@ -2983,6 +3072,19 @@ def create_render_context(
the ray (ray origin inside the geom). Matches MuJoCo's
mesh-ray rule. Default True. Disable for a small
performance gain when no camera is ever inside a geom.
background_color: The color to use for background pixels when no skybox is rendered.
enable_specular: Evaluate specular highlights per light. When False the
half-vector normalize and shininess `pow` are dropped at
compile time. Disable for performance when no specular is present.
enable_emission: Add `mat_emission * base_color` per shaded pixel. When
False the term is dropped at compile time. Disable for performance
when no emission is present.
enable_per_light_ambient: When ambient lighting is enabled, sum each
light's `ambient` color into shaded pixels
even outside its cone or in shadow. When False
the per-light ambient pass is removed at compile
time. Disable for performance when model lights
do not use ambient colors.
Returns:
The render context containing rendering fields and output arrays on device.
@@ -3163,6 +3265,13 @@ def create_render_context(
if len(flex_geom_flexid) > 0:
geom_ray_types.add(int(types.GeomType.FLEX))
geom_ray_types = tuple(sorted(geom_ray_types))
if mjm.nlight == 0:
light_attenuation_is_default = True
has_spot_lights = False
else:
atten = np.asarray(mjm.light_attenuation, dtype=np.float32).reshape(-1, 3)
light_attenuation_is_default = bool(np.allclose(atten, np.array([1.0, 0.0, 0.0], dtype=np.float32)))
has_spot_lights = bool((np.asarray(mjm.light_type) == int(mujoco.mjtLightType.mjLIGHT_SPOT)).any())
rc = types.RenderContext(
nrender=ncam,
@@ -3171,11 +3280,17 @@ def create_render_context(
use_textures=use_textures,
use_shadows=use_shadows,
use_ambient_lighting=use_ambient_lighting,
background_color=render_util.pack_rgba_to_uint32(0.1 * 255.0, 0.1 * 255.0, 0.2 * 255.0, 1.0 * 255.0),
background_color=render_util.pack_rgba_to_uint32(
background_color[0] * 255.0, background_color[1] * 255.0, background_color[2] * 255.0, background_color[3] * 255.0
),
use_precomputed_rays=use_precomputed_rays,
render_skybox=render_skybox,
skybox_tex_id=skybox_tex_id,
skybox_face_width=skybox_face_width,
headlight_active=bool(mjm.vis.headlight.active),
headlight_ambient=wp.vec3(mjm.vis.headlight.ambient),
headlight_diffuse=wp.vec3(mjm.vis.headlight.diffuse),
headlight_specular=wp.vec3(mjm.vis.headlight.specular),
bvh_ngeom=bvh_ngeom,
enabled_geom_ids=wp.array(geom_enabled_idx, dtype=int),
mesh_registry=mesh_registry,
@@ -3218,6 +3333,11 @@ def create_render_context(
total_rays=int(total),
enable_backface_culling=enable_backface_culling,
geom_ray_types=geom_ray_types,
enable_specular=enable_specular,
enable_emission=enable_emission,
enable_per_light_ambient=enable_per_light_ambient,
light_attenuation_is_default=light_attenuation_is_default,
has_spot_lights=has_spot_lights,
)
bvh.build_scene_bvh(mjm, mjd, rc, nworld)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
Vendored
+10 -5
View File
@@ -346,6 +346,7 @@ def _island_map_dofs(
island_idofadr_in: wp.array2d[int],
nidof_in: wp.array[int],
island_nv_inout: wp.array2d[int],
island_dofadr_out: wp.array2d[int],
map_dof2idof_out: wp.array2d[int],
map_idof2dof_out: wp.array2d[int],
idof_islandid_out: wp.array2d[int],
@@ -360,6 +361,7 @@ def _island_map_dofs(
local_idx = wp.atomic_add(island_nv_inout, worldid, island_id, 1)
idof = island_idofadr_in[worldid, island_id] + local_idx
idof_islandid_out[worldid, idof] = island_id
wp.atomic_min(island_dofadr_out, worldid, island_id, dofid)
else:
cnt = wp.atomic_add(unconstrained_cnt_inout, worldid, 0, 1)
idof = nidof + cnt
@@ -804,7 +806,7 @@ def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverConte
Populates island solver context arrays via ctx: nv, nefc, ne, nf,
iefcadr, nidof, map_dof2idof, map_idof2dof, dof_islandid, map_efc2iefc,
map_iefc2efc, efc_islandid. Also populates d.dof_island, d.efc.island,
and d.island_dofadr.
d.island_idofadr, and d.island_dofadr.
Args:
m: Model.
@@ -816,6 +818,8 @@ def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverConte
d.dof_islandid = wp.empty((d.nworld, m.nv), dtype=int)
if d.efc_islandid.shape[1] != d.njmax:
d.efc_islandid = wp.empty((d.nworld, d.njmax), dtype=int)
if d.island_idofadr.shape[1] != m.ntree:
d.island_idofadr = wp.empty((d.nworld, m.ntree), dtype=int)
# Ensure island-local DOF arrays are allocated at the right shape
if d.iqacc.shape[1] != m.nv:
@@ -829,7 +833,7 @@ def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverConte
dim=(d.nworld, m.ntree),
inputs=[],
outputs=[
d.island_dofadr,
d.island_idofadr,
d.island_nv,
d.island_nefc,
d.island_ne,
@@ -901,16 +905,17 @@ def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverConte
_island_scan_sizes,
dim=d.nworld,
inputs=[d.nisland],
outputs=[d.island_dofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
outputs=[d.island_idofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
)
# 4. Map DOFs
unconstrained_cnt = wp.zeros((d.nworld, 1), dtype=int)
d.island_dofadr.fill_(m.nv)
wp.launch(
_island_map_dofs,
dim=(d.nworld, m.nv),
inputs=[m.nv, d.dof_island, d.island_dofadr, d.nidof],
outputs=[d.island_nv, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
inputs=[m.nv, d.dof_island, d.island_idofadr, d.nidof],
outputs=[d.island_nv, d.island_dofadr, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
)
# 5. Map Constraints
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+183 -35
View File
@@ -40,6 +40,21 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
# Default value for mat_shininess in MuJoCo is 0.5
# With an 8 bit image format, the maximum value is 255.0
# So max shininess value for the Phong lighting model is 128.0
MAX_SHININESS = 128.0
# The exponent value for mat_shininess is 0.5 times the max shininess value
DEFAULT_MAT_SHININESS_EXPONENT = 0.5 * MAX_SHININESS
# Default value for mat_specular in MuJoCo is 0.5
DEFAULT_MAT_SPECULAR = 0.5
# Default value for mat_emission in MuJoCo is 0.0
DEFAULT_MAT_EMISSION = 0.0
NO_LIGHT_AMBIENT_FALLBACK = 0.3
@wp.func
def sample_texture(
@@ -416,37 +431,54 @@ def _make_compute_lighting(cast_ray_first_hit: wp.Function) -> wp.Function:
lightcastshadow: bool,
lightpos: wp.vec3,
lightdir: wp.vec3,
lightattenuation: wp.vec3,
lightcutoff_rad: float,
lightexp: float,
lightdiff: wp.vec3,
lightspec: wp.vec3,
normal: wp.vec3,
hitpoint: wp.vec3,
view_dir: wp.vec3,
mat_spec: float,
mat_shin_exp: float,
cull_backfaces: bool,
) -> float:
light_contribution = float(0.0)
enable_specular: bool,
default_attenuation: bool,
has_spot: bool,
) -> Tuple[wp.vec3, wp.vec3]:
diff_rgb = wp.vec3(0.0)
spec_rgb = wp.vec3(0.0)
# TODO: We should probably only be looping over active lights
# in the first place with a static loop of enabled light idx?
if not lightactive:
return light_contribution
return diff_rgb, spec_rgb
L = wp.vec3(0.0, 0.0, 0.0)
L = wp.vec3(0.0)
dist_to_light = float(MJ_MAXVAL)
attenuation = float(1.0)
attenuation = 1.0
if lighttype == 1: # directional light
L = wp.normalize(-lightdir)
# MuJoCo guarantees `lightdir` is unit length.
L = -lightdir
else:
L, dist_to_light = math.normalize_with_norm(lightpos - hitpoint)
attenuation = 1.0 / (1.0 + 0.02 * dist_to_light * dist_to_light)
if lighttype == 0: # spot light
spot_dir = wp.normalize(lightdir)
cos_theta = wp.dot(-L, spot_dir)
spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) * 10.0))
attenuation = attenuation * spot_factor
if not default_attenuation:
light_attenuation_factor = wp.vec3(1.0, dist_to_light, dist_to_light * dist_to_light)
attenuation = math.safe_div(1.0, wp.dot(light_attenuation_factor, lightattenuation))
if has_spot:
if lighttype == 0: # spot light
cos_theta = wp.dot(-L, lightdir)
cos_cutoff = wp.cos(lightcutoff_rad)
if cos_theta < cos_cutoff:
return diff_rgb, spec_rgb
attenuation = attenuation * wp.pow(wp.max(cos_theta, 0.0), lightexp)
ndotl = wp.max(0.0, wp.dot(normal, L))
if ndotl == 0.0:
return light_contribution
return diff_rgb, spec_rgb
visible = float(1.0)
visible = 1.0
if use_shadows and lightcastshadow:
# Nudge the origin slightly along the surface normal to avoid
@@ -486,9 +518,17 @@ def _make_compute_lighting(cast_ray_first_hit: wp.Function) -> wp.Function:
)
if shadow_geom_id != -1:
visible = 0.3
visible = NO_LIGHT_AMBIENT_FALLBACK
return ndotl * attenuation * visible
weight = attenuation * visible
diff_rgb = lightdiff * (ndotl * weight)
if enable_specular:
if mat_spec > 0.0 and mat_shin_exp > 0.0:
H = wp.normalize(L + view_dir)
ndoth = wp.max(0.0, wp.dot(normal, H))
spec_rgb = lightspec * (mat_spec * wp.pow(ndoth, mat_shin_exp) * weight)
return diff_rgb, spec_rgb
return compute_lighting
@@ -530,12 +570,21 @@ def render(m: Model, d: Data, rc: RenderContext):
light_type: wp.array2d[int],
light_castshadow: wp.array2d[bool],
light_active: wp.array2d[bool],
light_attenuation: wp.array2d[wp.vec3],
light_cutoff: wp.array2d[float],
light_exponent: wp.array2d[float],
light_ambient: wp.array2d[wp.vec3],
light_diffuse: wp.array2d[wp.vec3],
light_specular: wp.array2d[wp.vec3],
flex_vertadr: wp.array[int],
flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array[float],
mesh_faceadr: wp.array[int],
mat_texid: wp.array3d[int],
mat_texrepeat: wp.array2d[wp.vec2],
mat_emission: wp.array2d[float],
mat_specular: wp.array2d[float],
mat_shininess: wp.array2d[float],
mat_rgba: wp.array2d[wp.vec4],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
@@ -619,8 +668,9 @@ def render(m: Model, d: Data, rc: RenderContext):
wp.static(rc.znear),
)
ray_dir_world = cam_xmat_in[worldid, mujoco_cam_id] @ ray_dir_local_cam
ray_origin_world = cam_xpos_in[worldid, mujoco_cam_id]
cam_mat_world = cam_xmat_in[worldid, mujoco_cam_id]
ray_dir_world = cam_mat_world @ ray_dir_local_cam
geom_id, dist, normal, u, v, f, mesh_id = cast_ray(
geom_type,
@@ -694,7 +744,6 @@ def render(m: Model, d: Data, rc: RenderContext):
color = mat_rgba[worldid % mat_rgba.shape[0], geom_matid[worldid % geom_matid.shape[0], geom_id]]
base_color = wp.vec3(color[0], color[1], color[2])
hit_color = base_color
if wp.static(rc.use_textures):
if geom_id != -2:
@@ -721,18 +770,48 @@ def render(m: Model, d: Data, rc: RenderContext):
)
base_color = wp.cw_mul(base_color, tex_color)
result = wp.vec3(0.0, 0.0, 0.0)
if wp.static(rc.use_ambient_lighting):
len_n = wp.length(normal)
n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
n = wp.normalize(n)
hemispheric = 0.5 * (n[2] + 1.0)
ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
result = 0.5 * wp.cw_mul(base_color, ambient_color)
mat_spec = DEFAULT_MAT_SPECULAR
mat_shin_exp = DEFAULT_MAT_SHININESS_EXPONENT
mat_emis = DEFAULT_MAT_EMISSION
if wp.static(rc.enable_specular or rc.enable_emission):
if geom_id != -2:
mat_id_for_spec = geom_matid[worldid % geom_matid.shape[0], geom_id]
if mat_id_for_spec >= 0:
if wp.static(rc.enable_specular):
mat_spec = mat_specular[worldid % mat_specular.shape[0], mat_id_for_spec]
mat_shin_exp = mat_shininess[worldid % mat_shininess.shape[0], mat_id_for_spec] * MAX_SHININESS
if wp.static(rc.enable_emission):
mat_emis = mat_emission[worldid % mat_emission.shape[0], mat_id_for_spec]
# Apply lighting and shadows
for l in range(wp.static(m.nlight)):
light_contribution = compute_lighting(
result = wp.vec3(0.0)
if wp.static(rc.enable_emission):
result = base_color * mat_emis
if wp.static(rc.use_ambient_lighting):
if wp.static(rc.headlight_active):
result = result + wp.cw_mul(base_color, wp.static(rc.headlight_ambient))
elif wp.static(m.nlight == 0):
result = result + base_color * NO_LIGHT_AMBIENT_FALLBACK
if wp.static(rc.enable_per_light_ambient):
for light_index in range(wp.static(m.nlight)):
if light_active[worldid % light_active.shape[0], light_index]:
result = result + wp.cw_mul(base_color, light_ambient[worldid % light_ambient.shape[0], light_index])
view_dir = -ray_dir_world
light_cutoff_worldid = light_cutoff[worldid % light_cutoff.shape[0]]
light_active_worldid = light_active[worldid % light_active.shape[0]]
light_type_worldid = light_type[worldid % light_type.shape[0]]
light_castshadow_worldid = light_castshadow[worldid % light_castshadow.shape[0]]
light_xpos_in_worldid = light_xpos_in[worldid]
light_xdir_in_worldid = light_xdir_in[worldid]
light_attenuation_worldid = light_attenuation[worldid % light_attenuation.shape[0]]
light_exponent_worldid = light_exponent[worldid % light_exponent.shape[0]]
light_diffuse_worldid = light_diffuse[worldid % light_diffuse.shape[0]]
light_specular_worldid = light_specular[worldid % light_specular.shape[0]]
# Apply Lighting for each light
for light_index in range(wp.static(m.nlight)):
diff_rgb, spec_rgb = compute_lighting(
geom_type,
geom_dataid,
geom_size,
@@ -755,16 +834,76 @@ def render(m: Model, d: Data, rc: RenderContext):
flex_geom_edgeid,
flex_bvh_id,
flex_group_root,
light_active[worldid % light_active.shape[0], l],
light_type[worldid % light_type.shape[0], l],
light_castshadow[worldid % light_castshadow.shape[0], l],
light_xpos_in[worldid, l],
light_xdir_in[worldid, l],
light_active_worldid[light_index],
light_type_worldid[light_index],
light_castshadow_worldid[light_index],
light_xpos_in_worldid[light_index],
light_xdir_in_worldid[light_index],
light_attenuation_worldid[light_index],
light_cutoff_worldid[light_index] * wp.static(wp.pi / 180.0),
light_exponent_worldid[light_index],
light_diffuse_worldid[light_index],
light_specular_worldid[light_index],
normal,
hit_point,
view_dir,
mat_spec,
mat_shin_exp,
wp.static(rc.enable_backface_culling),
wp.static(rc.enable_specular),
wp.static(rc.light_attenuation_is_default),
wp.static(rc.has_spot_lights),
)
result = result + base_color * light_contribution
result = result + wp.cw_mul(base_color, diff_rgb) + spec_rgb
# Apply Headlight
if wp.static(rc.headlight_active):
cam_pos = ray_origin_world
cam_fwd = -cam_mat_world[:, 2]
hl_diff, hl_spec = compute_lighting(
geom_type,
geom_dataid,
geom_size,
flex_vertadr,
flex_edge,
flex_radius,
geom_xpos_in,
geom_xmat_in,
flexvert_xpos_in,
use_shadows,
bvh_id,
group_root[worldid],
bvh_ngeom,
bvh_nflexgeom,
enabled_geom_ids,
worldid,
mesh_bvh_id,
hfield_bvh_id,
flex_geom_flexid,
flex_geom_edgeid,
flex_bvh_id,
flex_group_root,
True,
1,
False,
cam_pos,
cam_fwd,
wp.vec3(1.0, 0.0, 0.0),
0.0,
0.0,
wp.static(rc.headlight_diffuse),
wp.static(rc.headlight_specular),
normal,
hit_point,
view_dir,
mat_spec,
mat_shin_exp,
wp.static(rc.enable_backface_culling),
wp.static(rc.enable_specular),
True,
False,
)
result = result + wp.cw_mul(base_color, hl_diff) + hl_spec
hit_color = wp.min(result, wp.vec3(1.0, 1.0, 1.0))
hit_color = wp.max(hit_color, wp.vec3(0.0, 0.0, 0.0))
@@ -792,12 +931,21 @@ def render(m: Model, d: Data, rc: RenderContext):
m.light_type,
m.light_castshadow,
m.light_active,
m.light_attenuation,
m.light_cutoff,
m.light_exponent,
m.light_ambient,
m.light_diffuse,
m.light_specular,
m.flex_vertadr,
m.flex_edge,
m.flex_radius,
m.mesh_faceadr,
m.mat_texid,
m.mat_texrepeat,
m.mat_emission,
m.mat_specular,
m.mat_shininess,
m.mat_rgba,
d.geom_xpos,
d.geom_xmat,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sleep.py
Vendored
+1032
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+837 -780
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+116 -9
View File
@@ -28,6 +28,7 @@ MJ_MINIMP = mujoco.mjMINIMP # minimum constraint impedance
MJ_MAXIMP = mujoco.mjMAXIMP # maximum constraint impedance
MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
MJ_MINMU = mujoco.mjMINMU # minimum friction
MJ_MINAWAKE = mujoco.mjMINAWAKE # minimum number of timesteps before sleeping
NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0.dev914519929")
TACTILE_DEPTH_SEMANTICS = check_version("mujoco>=3.9.0.dev921980899")
# maximum size (by number of edges) of an horizon in EPA algorithm
@@ -89,7 +90,7 @@ class BlockDim:
update_gradient_cholesky: int = 64
update_gradient_cholesky_blocked: int = 32
update_gradient_JTDAJ_sparse: int = 64
update_gradient_JTDAJ_dense: int = 96
update_gradient_JTDAJ_dense: int = 128
linesearch_iterative: int = 32
contact_jac_tiled: int = 32
# derivative
@@ -236,13 +237,44 @@ class EnableBit(enum.IntFlag):
Attributes:
ENERGY: energy computation
INVDISCRETE: discrete-time inverse dynamics
SLEEP: sleeping
"""
ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY
INVDISCRETE = mujoco.mjtEnableBit.mjENBL_INVDISCRETE
SLEEP = mujoco.mjtEnableBit.mjENBL_SLEEP
# unsupported: OVERRIDE, FWDINV, ISLAND
class SleepPolicy(enum.IntEnum):
"""Per-tree sleep policy.
Attributes:
AUTO: compiler chooses sleep policy
AUTO_NEVER: compiler sleep policy: never
AUTO_ALLOWED: compiler sleep policy: allowed
"""
AUTO = mujoco.mjtSleepPolicy.mjSLEEP_AUTO
AUTO_NEVER = mujoco.mjtSleepPolicy.mjSLEEP_AUTO_NEVER
AUTO_ALLOWED = mujoco.mjtSleepPolicy.mjSLEEP_AUTO_ALLOWED
# unsupported: NEVER, ALLOWED, INIT
class SleepState(enum.IntEnum):
"""Sleep state for bodies.
Attributes:
STATIC: body is static (world body or mocap)
ASLEEP: body is asleep
AWAKE: body is awake
"""
STATIC = mujoco.mjtSleepState.mjS_STATIC
ASLEEP = mujoco.mjtSleepState.mjS_ASLEEP
AWAKE = mujoco.mjtSleepState.mjS_AWAKE
class TrnType(enum.IntEnum):
"""Type of actuator transmission.
@@ -749,6 +781,7 @@ class Option:
tolerance: main solver tolerance
ls_tolerance: CG/Newton linesearch tolerance
ccd_tolerance: convex collision detection tolerance
sleep_tolerance: sleep velocity tolerance
gravity: gravitational acceleration
wind: wind (for lift, drag, and viscosity)
magnetic: global magnetic flux
@@ -783,6 +816,7 @@ class Option:
tolerance: array("*", float)
ls_tolerance: array("*", float)
ccd_tolerance: array("*", float)
sleep_tolerance: array("*", float)
gravity: array("*", wp.vec3)
wind: array("*", wp.vec3)
magnetic: array("*", wp.vec3)
@@ -978,9 +1012,11 @@ class Model:
dof_damping: damping coefficient (*, nv)
dof_dampingpoly: high-order damping coefficients (*, nv, 2)
dof_invweight0: diag. inverse inertia in qpos0 (*, nv)
dof_length: dof length for weighting velocity norm (nv,)
tree_bodynum: number of bodies in tree (incl. root) (ntree,)
tree_dofadr: start address of tree's dofs (ntree,)
tree_dofnum: number of dofs in tree (ntree,)
tree_sleep_policy: tree sleep policy (SleepPolicy) (ntree,)
geom_type: geometric type (GeomType) (ngeom,)
geom_contype: geom contact type (ngeom,)
geom_conaffinity: geom contact affinity (ngeom,)
@@ -1032,6 +1068,12 @@ class Model:
light_poscom0: global position rel. to sub-com in qpos0 (*, nlight, 3)
light_pos0: global position rel. to body in qpos0 (*, nlight, 3)
light_dir0: global direction in qpos0 (*, nlight, 3)
light_attenuation: OpenGL constant/linear/quadratic (*, nlight, 3)
light_cutoff: spotlight half-cone angle in degrees (*, nlight)
light_exponent: spotlight angular falloff exponent (*, nlight)
light_ambient: ambient RGB (*, nlight, 3)
light_diffuse: diffuse RGB (*, nlight, 3)
light_specular: specular RGB (*, nlight, 3)
flex_contype: flex contact type (nflex,)
flex_conaffinity: flex contact affinity (nflex,)
flex_condim: contact dimensionality (1, 3, 4, 6) (nflex,)
@@ -1100,6 +1142,9 @@ class Model:
hfield_data: elevation data (nhfielddata,)
mat_texid: texture id for rendering (*, nmat, mjNTEXROLE)
mat_texrepeat: texture repeat for rendering (*, nmat, 2)
mat_emission: emission scalar (self-illumination) (*, nmat)
mat_specular: specular reflection scalar (*, nmat)
mat_shininess: shininess in [0, 1], mapped to GL [0, 128](*, nmat)
mat_rgba: rgba (*, nmat, 4)
pair_dim: contact dimensionality (npair,)
pair_geom1: id of geom1 (npair,)
@@ -1402,9 +1447,11 @@ class Model:
dof_damping: array("*", "nv", float)
dof_dampingpoly: array("*", "nv", wp.vec2)
dof_invweight0: array("*", "nv", float)
dof_length: array("nv", float)
tree_bodynum: array("ntree", int)
tree_dofadr: array("ntree", int)
tree_dofnum: array("ntree", int)
tree_sleep_policy: array("ntree", int)
geom_type: array("ngeom", int)
geom_contype: array("ngeom", int)
geom_conaffinity: array("ngeom", int)
@@ -1456,6 +1503,12 @@ class Model:
light_poscom0: array("*", "nlight", wp.vec3)
light_pos0: array("*", "nlight", wp.vec3)
light_dir0: array("*", "nlight", wp.vec3)
light_attenuation: array("*", "nlight", wp.vec3)
light_cutoff: array("*", "nlight", float)
light_exponent: array("*", "nlight", float)
light_ambient: array("*", "nlight", wp.vec3)
light_diffuse: array("*", "nlight", wp.vec3)
light_specular: array("*", "nlight", wp.vec3)
flex_contype: array("nflex", int)
flex_conaffinity: array("nflex", int)
flex_condim: array("nflex", int)
@@ -1524,6 +1577,9 @@ class Model:
hfield_data: array("nhfielddata", float)
mat_texid: array("*", "nmat", 10, int)
mat_texrepeat: array("*", "nmat", wp.vec2)
mat_emission: array("*", "nmat", float)
mat_specular: array("*", "nmat", float)
mat_shininess: array("*", "nmat", float)
mat_rgba: array("*", "nmat", wp.vec4)
pair_dim: array("npair", int)
pair_geom1: array("npair", int)
@@ -1852,6 +1908,9 @@ class Data:
nefc: number of constraints (nworld,)
nisland: number of constraint islands (nworld,)
nidof: total DOFs in islands (nworld,)
ntree_awake: number of awake trees (nworld,)
nbody_awake: number of awake bodies (nworld,)
nv_awake: number of awake dofs (nworld,)
time: simulation time (nworld,)
energy: potential, kinetic energy (nworld, 2)
qpos: position (nworld, nq)
@@ -1868,6 +1927,7 @@ class Data:
qacc: acceleration (nworld, nv)
act_dot: time-derivative of actuator activation (nworld, na)
sensordata: sensor data array (nworld, nsensordata,)
tree_asleep: tree asleep counter; >=0: asleep cycle (nworld, ntree)
xpos: Cartesian position of body frame (nworld, nbody, 3)
xquat: Cartesian orientation of body frame (nworld, nbody, 4)
xmat: Cartesian orientation of body frame (nworld, nbody, 3, 3)
@@ -1906,6 +1966,10 @@ class Data:
qLD: upper Cholesky factorization (nworld, nv, nv) if dense
L'*D*L factorization of M (nworld, 1, nC) if sparse
qLDiagInv: 1/diag(D) (nworld, nv)
tree_awake: is tree awake; 0: asleep; 1: awake (nworld, ntree)
body_awake: body sleep state (SleepState) (nworld, nbody)
body_awake_ind: indices of awake/static bodies (nworld, nbody)
dof_awake_ind: indices of awake dofs (nworld, nv)
flexedge_velocity: flex edge velocities (nworld, nflexedge)
ten_velocity: tendon velocities (nworld, ntendon)
actuator_velocity: actuator velocities (nworld, nu)
@@ -1936,6 +2000,7 @@ class Data:
tree_island: island ID per tree (-1 if unconstrained) (nworld, ntree)
dof_island: island ID per DOF (-1 if unconstrained) (nworld, nv)
island_dofadr: island start address in dof vector (nworld, ntree)
island_idofadr: island start address in idof vector (nworld, ntree)
island_nv: DOFs per island (nworld, ntree)
island_nefc: constraints per island (nworld, ntree)
island_ne: equality constraints per island (nworld, ntree)
@@ -1970,6 +2035,9 @@ class Data:
nefc: array("nworld", int)
nisland: array("nworld", int)
nidof: array("nworld", int)
ntree_awake: array("nworld", int)
nbody_awake: array("nworld", int)
nv_awake: array("nworld", int)
time: array("nworld", float)
energy: array("nworld", wp.vec2)
qpos: array("nworld", "nq", float)
@@ -1986,6 +2054,7 @@ class Data:
qacc: array("nworld", "nv", float)
act_dot: array("nworld", "na", float)
sensordata: array("nworld", "nsensordata", float)
tree_asleep: array("nworld", "ntree", int)
xpos: array("nworld", "nbody", wp.vec3)
xquat: array("nworld", "nbody", wp.quat)
xmat: array("nworld", "nbody", wp.mat33)
@@ -2022,6 +2091,10 @@ class Data:
M: wp.array3d[float]
qLD: wp.array3d[float]
qLDiagInv: array("nworld", "nv", float)
tree_awake: array("nworld", "ntree", int)
body_awake: array("nworld", "nbody", int)
body_awake_ind: array("nworld", "nbody", int)
dof_awake_ind: array("nworld", "nv", int)
flexedge_velocity: array("nworld", "nflexedge", float)
ten_velocity: array("nworld", "ntendon", float)
actuator_velocity: array("nworld", "nu", float)
@@ -2050,6 +2123,7 @@ class Data:
tree_island: array("nworld", "ntree", int)
dof_island: array("nworld", "nv", int)
island_dofadr: array("nworld", "ntree", int)
island_idofadr: array("nworld", "ntree", int)
island_nv: array("nworld", "ntree", int)
island_nefc: array("nworld", "ntree", int)
island_ne: array("nworld", "ntree", int)
@@ -2083,9 +2157,6 @@ class InverseContext:
Jaref: wp.array2d[float]
search_dot: wp.array[float]
gauss: wp.array[float]
cost: wp.array[float]
prev_cost: wp.array[float]
done: wp.array[bool]
changed_efc_ids: wp.array2d[int]
changed_efc_count: wp.array[int]
@@ -2115,6 +2186,7 @@ class IslandSolverContext:
done: wp.array2d[bool] # per-island convergence
solver_niter: wp.array2d[int] # iterations per island
beta: wp.array2d[float]
beta_den: wp.array2d[float]
alpha: wp.array2d[float]
Ma: wp.array2d[float] # island-local Ma (nworld, nv)
@@ -2125,9 +2197,6 @@ class SolverContext:
Jaref: wp.array2d[float]
search_dot: wp.array[float]
gauss: wp.array[float]
cost: wp.array[float]
prev_cost: wp.array[float]
done: wp.array[bool]
grad: wp.array2d[float]
grad_dot: wp.array[float]
@@ -2138,9 +2207,11 @@ class SolverContext:
quad: wp.array2d[wp.vec3]
quad_gauss: wp.array[wp.vec3]
alpha: wp.array[float]
improvement: wp.array[float]
prev_grad: wp.array2d[float]
prev_Mgrad: wp.array2d[float]
beta: wp.array[float]
beta_den: wp.array[float]
h: wp.array3d[float]
hfactor: wp.array3d[float]
# Incremental Hessian update (Newton only)
@@ -2158,8 +2229,8 @@ class RenderContext:
cam_id_map: camera id map
use_textures: whether to use textures
use_shadows: whether to use shadows
use_ambient_lighting: whether to use ambient lighting
background_color: background color
use_ambient_lighting: top-level switch for ambient contributions
background_color: color used for missed rays when no skybox is rendered
use_precomputed_rays: whether to use precomputed rays
bvh_ngeom: number of geometries in the BVH
enabled_geom_ids: enabled geometry ids
@@ -2204,11 +2275,38 @@ class RenderContext:
render_skybox: whether to shade missed rays with the MuJoCo skybox texture
skybox_tex_id: index into textures of the skybox (MuJoCo tex_type == SKYBOX), -1 if none
skybox_face_width: pixel width of one skybox cube face (0 if no skybox)
headlight_active: whether to inject MuJoCo's vis.headlight as a synthetic
directional light at the active camera. Read from `mjm.vis.headlight.active`
at context creation; users disable the headlight by configuring it on the
MuJoCo model (e.g. `<visual><headlight active="0"/></visual>` in XML).
headlight_ambient: RGB ambient color of the headlight (from vis.headlight).
headlight_diffuse: RGB diffuse color of the headlight.
headlight_specular: RGB specular color of the headlight.
enable_backface_culling: drop primitive ray hits whose normal faces away
from the ray (i.e. the ray origin is inside the geom). Matches MuJoCo's
mesh-ray rule. When False, the renderer reports inner-surface hits, which
is faster but causes a camera placed inside a geom to render that geom's
back wall.
light_attenuation_is_default: True iff every light in the model has the
MuJoCo default `attenuation = (1, 0, 0)`. Computed once at context
creation; when True the kernel skips the per-light polynomial
attenuation evaluation (a divide + 3 multiplies + an add per
non-directional light per pixel) via `wp.static`.
has_spot_lights: True iff any light in the model has `type == SPOT`.
When False, the kernel skips the spot-cone branch (cos cutoff +
pow exponent) per non-directional light per pixel via `wp.static`.
enable_specular: when True, evaluate the Phong specular highlight per
light per pixel (uses `mat_specular` / `mat_shininess`). When False,
the entire specular branch is removed at compile time. Useful for
depth/segmentation-only workflows or when materials are matte.
enable_emission: when True, add `mat_emission * base_color` to each
shaded pixel. When False the term is dropped at compile time.
enable_per_light_ambient: when True and `use_ambient_lighting` is also
True, sum the per-light `light_ambient` colors into each shaded pixel
even when the surface normal is perpendicular to the light direction
or the pixel is shadowed. When False the second per-light loop for
ambient is removed at compile time. Headlight ambient and the no-light
fallback are controlled by `use_ambient_lighting`.
geom_ray_types: tuple of GeomType int values present in the scene, used to
statically eliminate unused intersection branches in the ray-cast kernels.
"""
@@ -2224,6 +2322,10 @@ class RenderContext:
render_skybox: bool
skybox_tex_id: int
skybox_face_width: int
headlight_active: bool
headlight_ambient: wp.vec3
headlight_diffuse: wp.vec3
headlight_specular: wp.vec3
bvh_ngeom: int
enabled_geom_ids: array("*", int)
mesh_registry: dict
@@ -2265,4 +2367,9 @@ class RenderContext:
znear: float
total_rays: int
enable_backface_culling: bool
enable_specular: bool
enable_emission: bool
enable_per_light_ambient: bool
light_attenuation_is_default: bool
has_spot_lights: bool
geom_ray_types: tuple = ()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/warp_util.py
Vendored
-25
View File
@@ -154,28 +154,3 @@ def check_toolkit_driver():
""",
stacklevel=2,
)
class scoped_mathdx_gemm_disabled:
"""Temporarily disable Warp MathDx GEMM kernels within this scope."""
def __init__(self, disable: bool = True):
self._disable = disable
self._config = None
self._prev = None
def __enter__(self):
if not self._disable:
return self
config = getattr(wp, "config", None)
if config is None or not hasattr(config, "enable_mathdx_gemm"):
return self
self._config = config
self._prev = config.enable_mathdx_gemm
self._config.enable_mathdx_gemm = False
return self
def __exit__(self, exc_type, exc, tb):
if self._config is not None:
self._config.enable_mathdx_gemm = self._prev
return False
mjx/mujoco/mjx/warp/collision_driver.py
+11
View File
@@ -23,6 +23,7 @@ import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
)
@@ -45,6 +46,7 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _collision_shim(
# Model
@@ -72,6 +74,7 @@ def _collision_shim(
flex_vertadr: wp.array[int],
flex_vertflexid: wp.array[int],
geom_aabb: wp.array3d[wp.vec3],
geom_bodyid: wp.array[int],
geom_conaffinity: wp.array[int],
geom_condim: wp.array[int],
geom_contype: wp.array[int],
@@ -139,11 +142,13 @@ def _collision_shim(
opt__ccd_iterations: int,
opt__ccd_tolerance: wp.array[float],
opt__disableflags: int,
opt__enableflags: int,
opt__sdf_initpoints: int,
opt__sdf_iterations: int,
# Data
naccdmax: int,
naconmax: int,
body_awake: wp.array2d[int],
flexvert_xpos: wp.array2d[wp.vec3],
geom_xmat: wp.array2d[wp.mat33],
geom_xpos: wp.array2d[wp.vec3],
@@ -194,6 +199,7 @@ def _collision_shim(
_m.flex_vertadr = flex_vertadr
_m.flex_vertflexid = flex_vertflexid
_m.geom_aabb = geom_aabb
_m.geom_bodyid = geom_bodyid
_m.geom_conaffinity = geom_conaffinity
_m.geom_condim = geom_condim
_m.geom_contype = geom_contype
@@ -252,6 +258,7 @@ def _collision_shim(
_m.opt.ccd_iterations = opt__ccd_iterations
_m.opt.ccd_tolerance = opt__ccd_tolerance
_m.opt.disableflags = opt__disableflags
_m.opt.enableflags = opt__enableflags
_m.opt.sdf_initpoints = opt__sdf_initpoints
_m.opt.sdf_iterations = opt__sdf_iterations
_m.pair_dim = pair_dim
@@ -263,6 +270,7 @@ def _collision_shim(
_m.pair_solreffriction = pair_solreffriction
_m.plugin = plugin
_m.plugin_attr = plugin_attr
_d.body_awake = body_awake
_d.contact.dim = contact__dim
_d.contact.dist = contact__dist
_d.contact.efc_address = contact__efc_address
@@ -385,6 +393,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
m.flex_vertadr,
m._impl.flex_vertflexid,
m.geom_aabb,
m.geom_bodyid,
m.geom_conaffinity,
m.geom_condim,
m.geom_contype,
@@ -452,10 +461,12 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
m.opt._impl.ccd_iterations,
m.opt._impl.ccd_tolerance,
m.opt.disableflags,
m.opt.enableflags,
m.opt._impl.sdf_initpoints,
m.opt._impl.sdf_iterations,
d._impl.naccdmax,
d._impl.naconmax,
d._impl.body_awake,
d._impl.flexvert_xpos,
d.geom_xmat,
d.geom_xpos,
mjx/mujoco/mjx/warp/forward.py
+405 -280
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/warp/render.py
+28
View File
@@ -67,9 +67,18 @@ def _render_shim(
geom_size: wp.array2d[wp.vec3],
geom_type: wp.array[int],
light_active: wp.array2d[bool],
light_ambient: wp.array2d[wp.vec3],
light_attenuation: wp.array2d[wp.vec3],
light_castshadow: wp.array2d[bool],
light_cutoff: wp.array2d[float],
light_diffuse: wp.array2d[wp.vec3],
light_exponent: wp.array2d[float],
light_specular: wp.array2d[wp.vec3],
light_type: wp.array2d[int],
mat_emission: wp.array2d[float],
mat_rgba: wp.array2d[wp.vec4],
mat_shininess: wp.array2d[float],
mat_specular: wp.array2d[float],
mat_texid: wp.array3d[int],
mat_texrepeat: wp.array2d[wp.vec2],
mesh_faceadr: wp.array[int],
@@ -107,9 +116,18 @@ def _render_shim(
_m.geom_size = geom_size
_m.geom_type = geom_type
_m.light_active = light_active
_m.light_ambient = light_ambient
_m.light_attenuation = light_attenuation
_m.light_castshadow = light_castshadow
_m.light_cutoff = light_cutoff
_m.light_diffuse = light_diffuse
_m.light_exponent = light_exponent
_m.light_specular = light_specular
_m.light_type = light_type
_m.mat_emission = mat_emission
_m.mat_rgba = mat_rgba
_m.mat_shininess = mat_shininess
_m.mat_specular = mat_specular
_m.mat_texid = mat_texid
_m.mat_texrepeat = mat_texrepeat
_m.mesh_faceadr = mesh_faceadr
@@ -154,6 +172,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
'geom_xpos',
'light_active',
'light_castshadow',
'light_cutoff',
'light_type',
'mat_rgba',
'mat_texid',
@@ -178,9 +197,18 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
m.geom_size,
m.geom_type,
m.light_active,
m._impl.light_ambient,
m._impl.light_attenuation,
m.light_castshadow,
m.light_cutoff,
m._impl.light_diffuse,
m._impl.light_exponent,
m._impl.light_specular,
m.light_type,
m._impl.mat_emission,
m.mat_rgba,
m._impl.mat_shininess,
m._impl.mat_specular,
m.mat_texid,
m._impl.mat_texrepeat,
m.mesh_faceadr,
mjx/mujoco/mjx/warp/smooth.py
+1
View File
@@ -23,6 +23,7 @@ import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
)
mjx/mujoco/mjx/warp/types.py
+64 -1
View File
@@ -23,7 +23,6 @@ from jax import tree_util
from jax.interpreters import batching
from mujoco.mjx._src import dataclasses as mjx_dataclasses
import numpy as np
if typing.TYPE_CHECKING:
GraphMode = int
@@ -145,6 +144,7 @@ class OptionWarp(PyTreeNode):
run_collision_detection: bool
sdf_initpoints: int
sdf_iterations: int
sleep_tolerance: jax.Array
class ModelWarp(PyTreeNode):
"""Derived fields from Model."""
@@ -178,6 +178,7 @@ class ModelWarp(PyTreeNode):
callback: Callback
cam_projection: np.ndarray
collision_sensor_adr: np.ndarray
dof_length: np.ndarray
dof_tri_col: np.ndarray
dof_tri_row: np.ndarray
eq_connect_adr: np.ndarray
@@ -231,11 +232,19 @@ class ModelWarp(PyTreeNode):
is_sparse: bool
jnt_limited_ball_adr: np.ndarray
jnt_limited_slide_hinge_adr: np.ndarray
light_ambient: jax.Array
light_attenuation: jax.Array
light_bodyid: np.ndarray
light_diffuse: jax.Array
light_exponent: jax.Array
light_specular: jax.Array
light_targetbodyid: np.ndarray
mapD2M: np.ndarray
mapM2D: np.ndarray
mapM2M: np.ndarray
mat_emission: jax.Array
mat_shininess: jax.Array
mat_specular: jax.Array
mat_texrepeat: jax.Array
max_ten_J_rownnz: int
mesh_polyadr: np.ndarray
@@ -321,6 +330,7 @@ class ModelWarp(PyTreeNode):
tree_bodynum: np.ndarray
tree_dofadr: np.ndarray
tree_dofnum: np.ndarray
tree_sleep_policy: np.ndarray
wrap_geom_adr: np.ndarray
wrap_jnt_adr: np.ndarray
wrap_pulley_scale: np.ndarray
@@ -332,6 +342,8 @@ class DataWarp(PyTreeNode):
M: jax.Array
actuator_moment: jax.Array
actuator_velocity: jax.Array
body_awake: jax.Array
body_awake_ind: jax.Array
cacc: jax.Array
cfrc_ext: jax.Array
cfrc_int: jax.Array
@@ -353,6 +365,7 @@ class DataWarp(PyTreeNode):
contact__vert: jax.Array
contact__worldid: jax.Array
crb: jax.Array
dof_awake_ind: jax.Array
dof_island: jax.Array
dof_islandid: jax.Array
efc__D: jax.Array
@@ -395,6 +408,7 @@ class DataWarp(PyTreeNode):
iqfrc_smooth: jax.Array
island_dofadr: jax.Array
island_efcadr: jax.Array
island_idofadr: jax.Array
island_ne: jax.Array
island_nefc: jax.Array
island_nf: jax.Array
@@ -411,6 +425,7 @@ class DataWarp(PyTreeNode):
naccdmax: int
nacon: jax.Array
naconmax: int
nbody_awake: jax.Array
ncollision: jax.Array
ne: jax.Array
nefc: jax.Array
@@ -421,6 +436,8 @@ class DataWarp(PyTreeNode):
njmax_nnz: int
njmax_pad: int
nl: jax.Array
ntree_awake: jax.Array
nv_awake: jax.Array
nworld: int
qLD: jax.Array
qLDiagInv: jax.Array
@@ -434,6 +451,8 @@ class DataWarp(PyTreeNode):
ten_velocity: jax.Array
ten_wrapadr: jax.Array
ten_wrapnum: jax.Array
tree_asleep: jax.Array
tree_awake: jax.Array
tree_island: jax.Array
wrap_obj: jax.Array
wrap_xpos: jax.Array
@@ -499,6 +518,8 @@ _NDIM = {
'actuator_length': 2,
'actuator_moment': 2,
'actuator_velocity': 2,
'body_awake': 2,
'body_awake_ind': 2,
'cacc': 3,
'cam_xmat': 4,
'cam_xpos': 3,
@@ -526,6 +547,7 @@ _NDIM = {
'crb': 3,
'ctrl': 2,
'cvel': 3,
'dof_awake_ind': 2,
'dof_island': 2,
'dof_islandid': 2,
'efc__D': 2,
@@ -572,6 +594,7 @@ _NDIM = {
'iqfrc_smooth': 2,
'island_dofadr': 2,
'island_efcadr': 2,
'island_idofadr': 2,
'island_ne': 2,
'island_nefc': 2,
'island_nf': 2,
@@ -590,6 +613,7 @@ _NDIM = {
'naccdmax': 0,
'nacon': 1,
'naconmax': 0,
'nbody_awake': 1,
'ncollision': 1,
'ne': 1,
'nefc': 1,
@@ -600,6 +624,8 @@ _NDIM = {
'njmax_nnz': 0,
'njmax_pad': 0,
'nl': 1,
'ntree_awake': 1,
'nv_awake': 1,
'nworld': 0,
'qLD': 3,
'qLDiagInv': 2,
@@ -633,6 +659,8 @@ _NDIM = {
'ten_wrapadr': 2,
'ten_wrapnum': 2,
'time': 1,
'tree_asleep': 2,
'tree_awake': 2,
'tree_island': 2,
'wrap_obj': 3,
'wrap_xpos': 3,
@@ -755,6 +783,7 @@ _NDIM = {
'dof_frictionloss': 2,
'dof_invweight0': 2,
'dof_jntid': 1,
'dof_length': 1,
'dof_parentid': 1,
'dof_solimp': 3,
'dof_solref': 3,
@@ -867,20 +896,29 @@ _NDIM = {
'jnt_stiffnesspoly': 3,
'jnt_type': 1,
'light_active': 2,
'light_ambient': 3,
'light_attenuation': 3,
'light_bodyid': 1,
'light_castshadow': 2,
'light_cutoff': 2,
'light_diffuse': 3,
'light_dir': 3,
'light_dir0': 3,
'light_exponent': 2,
'light_mode': 1,
'light_pos': 3,
'light_pos0': 3,
'light_poscom0': 3,
'light_specular': 3,
'light_targetbodyid': 1,
'light_type': 2,
'mapD2M': 1,
'mapM2D': 1,
'mapM2M': 1,
'mat_emission': 2,
'mat_rgba': 3,
'mat_shininess': 2,
'mat_specular': 2,
'mat_texid': 3,
'mat_texrepeat': 3,
'max_ten_J_rownnz': 0,
@@ -993,6 +1031,7 @@ _NDIM = {
'opt__run_collision_detection': 0,
'opt__sdf_initpoints': 0,
'opt__sdf_iterations': 0,
'opt__sleep_tolerance': 1,
'opt__solver': 0,
'opt__timestep': 1,
'opt__tolerance': 1,
@@ -1088,6 +1127,7 @@ _NDIM = {
'tree_bodynum': 1,
'tree_dofadr': 1,
'tree_dofnum': 1,
'tree_sleep_policy': 1,
'wrap_geom_adr': 1,
'wrap_jnt_adr': 1,
'wrap_objid': 1,
@@ -1118,6 +1158,7 @@ _NDIM = {
'run_collision_detection': 0,
'sdf_initpoints': 0,
'sdf_iterations': 0,
'sleep_tolerance': 1,
'solver': 0,
'timestep': 1,
'tolerance': 1,
@@ -1135,6 +1176,8 @@ _BATCH_DIM = {
'actuator_length': True,
'actuator_moment': True,
'actuator_velocity': True,
'body_awake': True,
'body_awake_ind': True,
'cacc': True,
'cam_xmat': True,
'cam_xpos': True,
@@ -1162,6 +1205,7 @@ _BATCH_DIM = {
'crb': True,
'ctrl': True,
'cvel': True,
'dof_awake_ind': True,
'dof_island': True,
'dof_islandid': True,
'efc__D': True,
@@ -1208,6 +1252,7 @@ _BATCH_DIM = {
'iqfrc_smooth': True,
'island_dofadr': True,
'island_efcadr': True,
'island_idofadr': True,
'island_ne': True,
'island_nefc': True,
'island_nf': True,
@@ -1226,6 +1271,7 @@ _BATCH_DIM = {
'naccdmax': False,
'nacon': False,
'naconmax': False,
'nbody_awake': True,
'ncollision': False,
'ne': True,
'nefc': True,
@@ -1236,6 +1282,8 @@ _BATCH_DIM = {
'njmax_nnz': False,
'njmax_pad': False,
'nl': True,
'ntree_awake': True,
'nv_awake': True,
'nworld': False,
'qLD': True,
'qLDiagInv': True,
@@ -1269,6 +1317,8 @@ _BATCH_DIM = {
'ten_wrapadr': True,
'ten_wrapnum': True,
'time': True,
'tree_asleep': True,
'tree_awake': True,
'tree_island': True,
'wrap_obj': True,
'wrap_xpos': True,
@@ -1391,6 +1441,7 @@ _BATCH_DIM = {
'dof_frictionloss': True,
'dof_invweight0': True,
'dof_jntid': False,
'dof_length': False,
'dof_parentid': False,
'dof_solimp': True,
'dof_solref': True,
@@ -1503,20 +1554,29 @@ _BATCH_DIM = {
'jnt_stiffnesspoly': True,
'jnt_type': False,
'light_active': True,
'light_ambient': True,
'light_attenuation': True,
'light_bodyid': False,
'light_castshadow': True,
'light_cutoff': True,
'light_diffuse': True,
'light_dir': True,
'light_dir0': True,
'light_exponent': True,
'light_mode': False,
'light_pos': True,
'light_pos0': True,
'light_poscom0': True,
'light_specular': True,
'light_targetbodyid': False,
'light_type': True,
'mapD2M': False,
'mapM2D': False,
'mapM2M': False,
'mat_emission': True,
'mat_rgba': True,
'mat_shininess': True,
'mat_specular': True,
'mat_texid': True,
'mat_texrepeat': True,
'max_ten_J_rownnz': False,
@@ -1629,6 +1689,7 @@ _BATCH_DIM = {
'opt__run_collision_detection': False,
'opt__sdf_initpoints': False,
'opt__sdf_iterations': False,
'opt__sleep_tolerance': True,
'opt__solver': False,
'opt__timestep': True,
'opt__tolerance': True,
@@ -1724,6 +1785,7 @@ _BATCH_DIM = {
'tree_bodynum': False,
'tree_dofadr': False,
'tree_dofnum': False,
'tree_sleep_policy': False,
'wrap_geom_adr': False,
'wrap_jnt_adr': False,
'wrap_objid': False,
@@ -1754,6 +1816,7 @@ _BATCH_DIM = {
'run_collision_detection': False,
'sdf_initpoints': False,
'sdf_iterations': False,
'sleep_tolerance': True,
'solver': False,
'timestep': True,
'tolerance': True,