Skip to content

inertia box fluid model #191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

inertia box fluid model #191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
fb7a1de
inertia box fluid model
thowell Apr 27, 2025
ccfbdf0
fix passive.py imports
thowell Apr 27, 2025
76e237d
Merge remote-tracking branch 'upstream/main' into inertia_box_fluid_m…
thowell May 8, 2025
6893c77
update apply_ft
thowell May 8, 2025
ecc2224
update passive.py
thowell May 8, 2025
6c6960c
move _inertia_box_fluid_model functionality ino _box_fluid kernel
thowell May 8, 2025
197341c
wp.pow
thowell May 8, 2025
34b542c
wp.pow and wp.static
thowell May 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 0 additions & 17 deletions .vscode/launch.json
View file
Open in desktop

This file was deleted.

15 changes: 9 additions & 6 deletions mujoco_warp/_src/io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,12 +56,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):
Expand Down Expand Up @@ -273,6 +267,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,
Expand Down Expand Up @@ -552,6 +549,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),
Expand Down Expand Up @@ -589,6 +587,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),
Expand Down Expand Up @@ -838,6 +837,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),
Expand Down Expand Up @@ -875,6 +875,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),
Expand Down Expand Up @@ -1055,12 +1056,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]
Expand Down
16 changes: 0 additions & 16 deletions mujoco_warp/_src/io_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,22 +227,6 @@ 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):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe we should still check for ellipsoid fluid and fail?

mjm = mujoco.MjModel.from_xml_string("""
<mujoco>
<option wind="1 1 1" density="1" viscosity="1"/>
<worldbody>
<body>
<geom type="sphere" size=".1"/>
<freejoint/>
</body>
</worldbody>
</mujoco>
""")

with self.assertRaises(NotImplementedError):
mjwarp.put_model(mjm)


if __name__ == "__main__":
wp.init()
Expand Down
221 changes: 180 additions & 41 deletions mujoco_warp/_src/passive.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,20 @@
# limitations under the License.
# ==============================================================================

from typing import Tuple

import warp as wp

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
from .warp_util import kernel as nested_kernel


@wp.kernel
Expand Down Expand Up @@ -94,25 +99,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:
Expand Down Expand Up @@ -142,51 +128,189 @@ def _gravity_force(
wp.atomic_add(qfrc_gravcomp_out[worldid], dofid, wp.dot(jac, force))


def _box_fluid(opt: Option):
opt_wind = opt.wind
opt_viscosity = opt.viscosity
opt_density = opt.density

@nested_kernel
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This gets a bit subjective, but I think the bar might be higher for using a nested kernel. For tiles we have no choice, but for here, we can pass the opts in as params.

If it led to a significant performance difference I would understand, but this kernel looks pretty cheap.

In particular, my understanding is that the performance hit comes from thread divergence. e.g. some threads have viscosity and others don't. I don't think that's the case here, no?

def box_fluid(
# Model:
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 wp.static(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 = wp.static(opt_viscosity > 0.0)
density = wp.static(opt_density > 0.0)

if wp.static(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) * wp.static(1.0 / 3.0)

# angular viscosity
lfrc_torque = -lvel_torque * wp.pow(diam, 3.0) * wp.static(wp.pi * opt_viscosity)

# linear viscosity
lfrc_force = -lvel_force * diam * wp.static(3.0 * wp.pi * opt_viscosity)

if density:
# force
lfrc_force -= wp.vec3(
wp.static(0.5 * opt_density) * box1 * box2 * wp.abs(lvel_force[0]) * lvel_force[0],
wp.static(0.5 * opt_density) * box0 * box2 * wp.abs(lvel_force[1]) * lvel_force[1],
wp.static(0.5 * opt_density) * box0 * box1 * wp.abs(lvel_force[2]) * lvel_force[2],
)

# torque
scl = wp.static(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)

return box_fluid


def _fluid(m: Model, d: Data):
wp.launch(
_box_fluid(m.opt),
dim=(d.nworld, m.nbody),
inputs=[
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_gravcomp(
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,
dim=(d.nworld, m.njnt),
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),
Expand All @@ -203,9 +327,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],
)
Loading
Loading