diff --git a/.vscode/launch.json b/.vscode/launch.json deleted file mode 100644 index 4f26c7ce..00000000 --- a/.vscode/launch.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "version": "0.2.0", - "configurations": [ - { - "args": [ - "--extensionDevelopmentPath=${workspaceFolder}/contrib/kernel_analyzer" - ], - "name": "Launch Extension", - "outFiles": [ - "${workspaceFolder}/contrib/kernel_analyzer/out/**/*.js" - ], - "preLaunchTask": "${defaultBuildTask}", - "request": "launch", - "type": "extensionHost", - } - ] -} \ No newline at end of file diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 513cc1f0..0c9c28e3 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -47,6 +47,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: if mjm.tendon_frictionloss.any(): raise NotImplementedError("Tendon frictionloss is unsupported.") + if mjm.geom_fluid.any(): + raise NotImplementedError("Ellipsoid fluid model not implemented.") + # check options for opt, opt_types, msg in ( (mjm.opt.integrator, types.IntegratorType, "Integrator"), @@ -56,12 +59,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: if opt not in set(opt_types): raise NotImplementedError(f"{msg} {opt} is unsupported.") - if mjm.opt.wind.any(): - raise NotImplementedError("Wind is unsupported.") - - if mjm.opt.density > 0 or mjm.opt.viscosity > 0: - raise NotImplementedError("Fluid forces are unsupported.") - # TODO(team): remove after solver._update_gradient for Newton solver utilizes tile operations for islands nv_max = 60 if mjm.nv > nv_max and (not mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_SPARSE): @@ -273,6 +270,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: tolerance=mjm.opt.tolerance, ls_tolerance=mjm.opt.ls_tolerance, gravity=wp.vec3(mjm.opt.gravity), + wind=wp.vec3(mjm.opt.wind[0], mjm.opt.wind[1], mjm.opt.wind[2]), + density=mjm.opt.density, + viscosity=mjm.opt.viscosity, cone=mjm.opt.cone, solver=mjm.opt.solver, iterations=mjm.opt.iterations, @@ -552,6 +552,7 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in ctrl=wp.zeros((nworld, mjm.nu), dtype=float), qfrc_applied=wp.zeros((nworld, mjm.nv), dtype=float), xfrc_applied=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), + fluid_applied=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), eq_active=wp.array(np.tile(mjm.eq_active0, (nworld, 1)), dtype=bool), mocap_pos=wp.zeros((nworld, mjm.nmocap), dtype=wp.vec3), mocap_quat=wp.zeros((nworld, mjm.nmocap), dtype=wp.quat), @@ -589,6 +590,7 @@ def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: in qfrc_spring=wp.zeros((nworld, mjm.nv), dtype=float), qfrc_damper=wp.zeros((nworld, mjm.nv), dtype=float), qfrc_gravcomp=wp.zeros((nworld, mjm.nv), dtype=float), + qfrc_fluid=wp.zeros((nworld, mjm.nv), dtype=float), qfrc_passive=wp.zeros((nworld, mjm.nv), dtype=float), subtree_linvel=wp.zeros((nworld, mjm.nbody), dtype=wp.vec3), subtree_angmom=wp.zeros((nworld, mjm.nbody), dtype=wp.vec3), @@ -838,6 +840,7 @@ def padtile(x, length, dtype=None): ctrl=tile(mjd.ctrl), qfrc_applied=tile(mjd.qfrc_applied), xfrc_applied=tile(mjd.xfrc_applied, dtype=wp.spatial_vector), + fluid_applied=wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector), eq_active=tile(mjd.eq_active.astype(bool)), mocap_pos=tile(mjd.mocap_pos, dtype=wp.vec3), mocap_quat=tile(mjd.mocap_quat, dtype=wp.quat), @@ -875,6 +878,7 @@ def padtile(x, length, dtype=None): qfrc_spring=tile(mjd.qfrc_spring), qfrc_damper=tile(mjd.qfrc_damper), qfrc_gravcomp=tile(mjd.qfrc_gravcomp), + qfrc_fluid=tile(mjd.qfrc_fluid), qfrc_passive=tile(mjd.qfrc_passive), subtree_linvel=tile(mjd.subtree_linvel, dtype=wp.vec3), subtree_angmom=tile(mjd.subtree_angmom, dtype=wp.vec3), @@ -1055,12 +1059,14 @@ def get_data_into( result.cvel = d.cvel.numpy()[0] result.cdof_dot = d.cdof_dot.numpy()[0] result.qfrc_bias = d.qfrc_bias.numpy()[0] + result.qfrc_fluid = d.qfrc_fluid.numpy()[0] result.qfrc_passive = d.qfrc_passive.numpy()[0] result.subtree_linvel = d.subtree_linvel.numpy()[0] result.subtree_angmom = d.subtree_angmom.numpy()[0] result.qfrc_spring = d.qfrc_spring.numpy()[0] result.qfrc_damper = d.qfrc_damper.numpy()[0] result.qfrc_gravcomp = d.qfrc_gravcomp.numpy()[0] + result.qfrc_fluid = d.qfrc_fluid.numpy()[0] result.qfrc_actuator = d.qfrc_actuator.numpy()[0] result.qfrc_smooth = d.qfrc_smooth.numpy()[0] result.qfrc_constraint = d.qfrc_constraint.numpy()[0] diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index edb99f31..98c6f188 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -227,20 +227,21 @@ def test_get_data_into_m(self): np.testing.assert_allclose(mjd.qLD, mjd_ref.qLD) np.testing.assert_allclose(mjd.qM, mjd_ref.qM) - def test_option_physical_constants(self): - mjm = mujoco.MjModel.from_xml_string(""" + def test_ellipsoid_fluid_model(self): + with self.assertRaises(NotImplementedError): + mjm = mujoco.MjModel.from_xml_string( + """ - - """) - - with self.assertRaises(NotImplementedError): + + + """ + ) mjwarp.put_model(mjm) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index c382716d..80e58eed 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -17,10 +17,12 @@ from . import math from . import support +from .types import MJ_MINVAL from .types import Data from .types import DisableBit from .types import JointType from .types import Model +from .types import Option from .warp_util import event_scope @@ -94,25 +96,6 @@ def _spring_passive( qfrc_spring_out[worldid, dofid] = -stiffness * fdif -@wp.kernel -def _damper_passive( - # Model: - dof_damping: wp.array(dtype=float), - # Data in: - qvel_in: wp.array2d(dtype=float), - qfrc_spring_in: wp.array2d(dtype=float), - # Data out: - qfrc_damper_out: wp.array2d(dtype=float), - qfrc_passive_out: wp.array2d(dtype=float), -): - worldid, dofid = wp.tid() - - qfrc_damper = -dof_damping[dofid] * qvel_in[worldid, dofid] - - qfrc_damper_out[worldid, dofid] = qfrc_damper - qfrc_passive_out[worldid, dofid] = qfrc_damper + qfrc_spring_in[worldid, dofid] - - @wp.kernel def _gravity_force( # Model: @@ -143,35 +126,175 @@ def _gravity_force( @wp.kernel -def _qfrc_passive_gravcomp( +def _box_fluid( + # Model: + opt_wind: wp.vec3, + opt_density: float, + opt_viscosity: float, + body_rootid: wp.array(dtype=int), + body_mass: wp.array(dtype=float), + body_inertia: wp.array(dtype=wp.vec3), + # Data in: + xipos_in: wp.array2d(dtype=wp.vec3), + ximat_in: wp.array2d(dtype=wp.mat33), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cvel_in: wp.array2d(dtype=wp.spatial_vector), + # Data out: + fluid_applied_out: wp.array2d(dtype=wp.spatial_vector), +): + """Fluid forces based on inertia-box approximation.""" + + worldid, bodyid = wp.tid() + + # map from CoM-centered to local body-centered 6D velocity + + # body-inertial + pos = xipos_in[worldid, bodyid] + rot = ximat_in[worldid, bodyid] + rotT = wp.transpose(rot) + + # transform velocity + cvel = cvel_in[worldid, bodyid] + torque = wp.spatial_top(cvel) + force = wp.spatial_bottom(cvel) + subtree_com = subtree_com_in[worldid, body_rootid[bodyid]] + dif = pos - subtree_com + force -= wp.cross(dif, torque) + + lvel_torque = rotT @ torque + lvel_force = rotT @ force + + if opt_wind[0] or opt_wind[1] or opt_wind[2]: + # subtract translational component from body velocity + lvel_force -= rotT @ opt_wind + + lfrc_torque = wp.vec3(0.0) + lfrc_force = wp.vec3(0.0) + + viscosity = opt_viscosity > 0.0 + density = opt_density > 0.0 + + if viscosity or density: + inertia = body_inertia[bodyid] + mass = body_mass[bodyid] + scl = 6.0 / mass + box0 = wp.sqrt(wp.max(MJ_MINVAL, inertia[1] + inertia[2] - inertia[0]) * scl) + box1 = wp.sqrt(wp.max(MJ_MINVAL, inertia[0] + inertia[2] - inertia[1]) * scl) + box2 = wp.sqrt(wp.max(MJ_MINVAL, inertia[0] + inertia[1] - inertia[2]) * scl) + + if viscosity: + # diameter of sphere approximation + diam = (box0 + box1 + box2) / 3.0 + + # angular viscosity + lfrc_torque = -lvel_torque * wp.pow(diam, 3.0) * wp.pi * opt_viscosity + + # linear viscosity + lfrc_force = -3.0 * lvel_force * diam * wp.pi * opt_viscosity + + if density: + # force + lfrc_force -= wp.vec3( + 0.5 * opt_density * box1 * box2 * wp.abs(lvel_force[0]) * lvel_force[0], + 0.5 * opt_density * box0 * box2 * wp.abs(lvel_force[1]) * lvel_force[1], + 0.5 * opt_density * box0 * box1 * wp.abs(lvel_force[2]) * lvel_force[2], + ) + + # torque + scl = opt_density / 64.0 + box0_pow4 = wp.pow(box0, 4.0) + box1_pow4 = wp.pow(box1, 4.0) + box2_pow4 = wp.pow(box2, 4.0) + lfrc_torque -= wp.vec3( + box0 * (box1_pow4 + box2_pow4) * wp.abs(lvel_torque[0]) * lvel_torque[0] * scl, + box1 * (box0_pow4 + box2_pow4) * wp.abs(lvel_torque[1]) * lvel_torque[1] * scl, + box2 * (box0_pow4 + box1_pow4) * wp.abs(lvel_torque[2]) * lvel_torque[2] * scl, + ) + + # rotate to global orientation: lfrc -> bfrc + torque = rot @ lfrc_torque + force = rot @ lfrc_force + + fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force, torque) + + +def _fluid(m: Model, d: Data): + wp.launch( + _box_fluid, + dim=(d.nworld, m.nbody), + inputs=[ + m.opt.wind, + m.opt.density, + m.opt.viscosity, + m.body_rootid, + m.body_mass, + m.body_inertia, + d.xipos, + d.ximat, + d.subtree_com, + d.cvel, + ], + outputs=[ + d.fluid_applied, + ], + ) + + # TODO(team): ellipsoid fluid model + + support.apply_ft(m, d, d.fluid_applied, d.qfrc_fluid) + + +@wp.kernel +def _qfrc_passive( # Model: jnt_actgravcomp: wp.array(dtype=int), dof_jntid: wp.array(dtype=int), + dof_damping: wp.array(dtype=float), # Data in: + qvel_in: wp.array2d(dtype=float), + qfrc_spring_in: wp.array2d(dtype=float), qfrc_gravcomp_in: wp.array2d(dtype=float), + qfrc_fluid_in: wp.array2d(dtype=float), + # In: + gravcomp: bool, + fluid: bool, # Data out: + qfrc_damper_out: wp.array2d(dtype=float), qfrc_passive_out: wp.array2d(dtype=float), ): worldid, dofid = wp.tid() + # spring + qfrc_passive = qfrc_spring_in[worldid, dofid] + + # damper + qfrc_damper = -dof_damping[dofid] * qvel_in[worldid, dofid] + qfrc_damper_out[worldid, dofid] = qfrc_damper + + qfrc_passive += qfrc_damper + # add gravcomp unless added by actuators - if jnt_actgravcomp[dof_jntid[dofid]]: - return + if gravcomp and not jnt_actgravcomp[dof_jntid[dofid]]: + qfrc_passive += qfrc_gravcomp_in[worldid, dofid] - qfrc_passive_out[worldid, dofid] += qfrc_gravcomp_in[worldid, dofid] + # add fluid force + if fluid: + qfrc_passive += qfrc_fluid_in[worldid, dofid] + + qfrc_passive_out[worldid, dofid] = qfrc_passive @event_scope def passive(m: Model, d: Data): """Adds all passive forces.""" if m.opt.disableflags & DisableBit.PASSIVE: - d.qfrc_passive.zero_() + d.qfrc_spring.zero_() + d.qfrc_damper.zero_() d.qfrc_gravcomp.zero_() + d.qfrc_fluid.zero_() + d.qfrc_passive.zero_() return - # TODO(team): mj_ellipsoidFluidModel - # TODO(team): mj_inertiaBoxFluidModell - d.qfrc_spring.zero_() wp.launch( _spring_passive, @@ -179,14 +302,11 @@ def passive(m: Model, d: Data): inputs=[m.qpos_spring, m.jnt_type, m.jnt_qposadr, m.jnt_dofadr, m.jnt_stiffness, d.qpos], outputs=[d.qfrc_spring], ) - wp.launch( - _damper_passive, - dim=(d.nworld, m.nv), - inputs=[m.dof_damping, d.qvel, d.qfrc_spring], - outputs=[d.qfrc_damper, d.qfrc_passive], - ) - if m.ngravcomp and not (m.opt.disableflags & DisableBit.GRAVITY): - d.qfrc_gravcomp.zero_() + + gravcomp = m.ngravcomp and not (m.opt.disableflags & DisableBit.GRAVITY) + d.qfrc_gravcomp.zero_() + + if gravcomp: wp.launch( _gravity_force, dim=(d.nworld, m.nbody - 1, m.nv), @@ -203,9 +323,24 @@ def passive(m: Model, d: Data): ], outputs=[d.qfrc_gravcomp], ) - wp.launch( - _qfrc_passive_gravcomp, - dim=(d.nworld, m.nv), - inputs=[m.jnt_actgravcomp, m.dof_jntid, d.qfrc_gravcomp], - outputs=[d.qfrc_passive], - ) + + fluid = m.opt.density or m.opt.viscosity or m.opt.wind[0] or m.opt.wind[1] or m.opt.wind[2] + if fluid: + _fluid(m, d) + + wp.launch( + _qfrc_passive, + dim=(d.nworld, m.nv), + inputs=[ + m.jnt_actgravcomp, + m.dof_jntid, + m.dof_damping, + d.qvel, + d.qfrc_spring, + d.qfrc_gravcomp, + d.qfrc_fluid, + gravcomp, + fluid, + ], + outputs=[d.qfrc_damper, d.qfrc_passive], + ) diff --git a/mujoco_warp/_src/passive_test.py b/mujoco_warp/_src/passive_test.py index 0fe2f3b5..9378b320 100644 --- a/mujoco_warp/_src/passive_test.py +++ b/mujoco_warp/_src/passive_test.py @@ -53,6 +53,43 @@ def test_passive(self, gravity): # TODO(team): test DisableBit.PASSIVE + @parameterized.parameters( + (1, 0, 0, 0, 0), + (0, 1, 0, 0, 0), + (0, 0, 1, 0, 0), + (0, 0, 0, 1, 0), + (0, 0, 0, 0, 1), + (1, 1, 1, 1, 1), + ) + def test_fluid(self, density, viscosity, wind0, wind1, wind2): + """Tests fluid model.""" + + _, mjd, m, d = test_util.fixture( + xml=f""" + + + """, + keyframe=0, + ) + + for arr in (d.qfrc_passive, d.qfrc_fluid): + arr.zero_() + + mjwarp.passive(m, d) + + _assert_eq(d.qfrc_passive.numpy()[0], mjd.qfrc_passive, "qfrc_passive") + _assert_eq(d.qfrc_fluid.numpy()[0], mjd.qfrc_fluid, "qfrc_fluid") + @parameterized.parameters((True, True), (True, False), (False, True), (False, False)) def test_gravcomp(self, sparse, gravity): """Tests gravity compenstation.""" diff --git a/mujoco_warp/_src/support.py b/mujoco_warp/_src/support.py index 015cb94b..664c7720 100644 --- a/mujoco_warp/_src/support.py +++ b/mujoco_warp/_src/support.py @@ -198,25 +198,59 @@ def xfrc_accumulate_kernel( out[worldid, dofid] += accumul -@event_scope -def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)): +@wp.kernel +def _apply_ft( + # Model: + nbody: int, + body_parentid: wp.array(dtype=int), + body_rootid: wp.array(dtype=int), + dof_bodyid: wp.array(dtype=int), + # Data in: + xipos_in: wp.array2d(dtype=wp.vec3), + subtree_com_in: wp.array2d(dtype=wp.vec3), + cdof_in: wp.array2d(dtype=wp.spatial_vector), + # In: + ft_in: wp.array2d(dtype=wp.spatial_vector), + # Out: + qfrc_out: wp.array2d(dtype=float), +): + worldid, dofid = wp.tid() + cdof = cdof_in[worldid, dofid] + rotational_cdof = wp.vec3(cdof[0], cdof[1], cdof[2]) + jac = wp.spatial_vector(cdof[3], cdof[4], cdof[5], cdof[0], cdof[1], cdof[2]) + + dofbodyid = dof_bodyid[dofid] + accumul = float(0.0) + + for bodyid in range(dofbodyid, nbody): + # any body that is in the subtree of dofbodyid is part of the jacobian + parentid = bodyid + while parentid != 0 and parentid != dofbodyid: + parentid = body_parentid[parentid] + if parentid == 0: + continue # body is not part of the subtree + offset = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]] + cross_term = wp.cross(rotational_cdof, offset) + ft_body = ft_in[worldid, bodyid] + accumul += wp.dot(jac, ft_body) + wp.dot(cross_term, wp.spatial_top(ft_body)) + + qfrc_out[worldid, dofid] += accumul + + +def apply_ft(m: Model, d: Data, ft: wp.array2d(dtype=wp.spatial_vector), qfrc: wp.array2d(dtype=float)): wp.launch( - kernel=xfrc_accumulate_kernel, + kernel=_apply_ft, dim=(d.nworld, m.nv), - inputs=[ - m.nbody, - m.body_parentid, - m.body_rootid, - m.dof_bodyid, - d.xfrc_applied, - d.xipos, - d.subtree_com, - d.cdof, - ], + inputs=[m.nbody, m.body_parentid, m.body_rootid, m.dof_bodyid, d.xipos, d.subtree_com, d.cdof, ft], outputs=[qfrc], ) +@event_scope +def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)): + apply_ft(m, d, d.xfrc_applied, qfrc) + + @wp.func def bisection(x: wp.array(dtype=int), v: int, a_: int, b_: int) -> int: # Binary search for the largest index i such that x[i] <= v diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index f87ea93e..4c7fc361 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -392,6 +392,9 @@ class Option: epa_exact_neg_distance: flag for enabling the distance calculation for non-intersecting case in the convex narrowphase depth_extension: distance for which the closest point is not calculated for non-intersecting case in the convex narrowphase ls_parallel: evaluate engine solver step sizes in parallel + wind: wind (for lift, drag, and viscosity) + density: density of medium + viscosity: viscosity of medium """ timestep: float @@ -411,6 +414,9 @@ class Option: epa_exact_neg_distance: bool # warp only depth_extension: float # warp only ls_parallel: bool + wind: wp.vec3 + density: float + viscosity: float @dataclasses.dataclass @@ -1047,6 +1053,7 @@ class Data: ctrl: control (nworld, nu) qfrc_applied: applied generalized force (nworld, nv) xfrc_applied: applied Cartesian force/torque (nworld, nbody, 6) + fluid_applied: applied fluid force/torque (nworld, nbody, 6) eq_active: enable/disable constraints (nworld, neq) mocap_pos: position of mocap bodies (nworld, nmocap, 3) mocap_quat: orientation of mocap bodies (nworld, nmocap, 4) @@ -1084,6 +1091,7 @@ class Data: qfrc_spring: passive spring force (nworld, nv) qfrc_damper: passive damper force (nworld, nv) qfrc_gravcomp: passive gravity compensation force (nworld, nv) + qfrc_fluid: passive fluid force (nworld, nv) qfrc_passive: total passive force (nworld, nv) subtree_linvel: linear velocity of subtree com (nworld, nbody, 3) subtree_angmom: angular momentum about subtree com (nworld, nbody, 3) @@ -1153,6 +1161,7 @@ class Data: ctrl: wp.array2d(dtype=float) qfrc_applied: wp.array2d(dtype=float) xfrc_applied: wp.array2d(dtype=wp.spatial_vector) + fluid_applied: wp.array2d(dtype=wp.spatial_vector) # warp only eq_active: wp.array2d(dtype=bool) mocap_pos: wp.array2d(dtype=wp.vec3) mocap_quat: wp.array2d(dtype=wp.quat) @@ -1190,6 +1199,7 @@ class Data: qfrc_spring: wp.array2d(dtype=float) qfrc_damper: wp.array2d(dtype=float) qfrc_gravcomp: wp.array2d(dtype=float) + qfrc_fluid: wp.array2d(dtype=float) qfrc_passive: wp.array2d(dtype=float) subtree_linvel: wp.array2d(dtype=wp.vec3) subtree_angmom: wp.array2d(dtype=wp.vec3)