Merge pull request #191 from thowell/inertia_box_fluid_model

inertia box fluid model

Author: thowell

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 mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_DENSE:
@@ -280,6 +277,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,
@@ -568,6 +568,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),
@@ -605,6 +606,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),
@@ -854,6 +856,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),
@@ -891,6 +894,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),
@@ -1071,12 +1075,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]

mujoco_warp/_src/io_test.py

@@ -222,20 +222,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(
+        """
       <mujoco>
-        <option wind="1 1 1" density="1" viscosity="1"/>
+        <option density="1"/>
         <worldbody>
-          <body>          
-            <geom type="sphere" size=".1"/>
+          <body>
+            <geom type="sphere" size=".1" fluidshape="ellipsoid"/>
             <freejoint/>
           </body>
-        </worldbody> 
-    </mujoco>
-    """)
-
-    with self.assertRaises(NotImplementedError):
+        </worldbody>
+      </mujoco>
+      """
+      )
       mjwarp.put_model(mjm)
 
   def test_jacobian_auto(self):

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,50 +126,187 @@ 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,
     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),
@@ -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],
+  )