Add implementation of mjd_ellipsoidFluid derivative by Aaryan-549 · Pull Request #897 · google-deepmind/mujoco_warp

Aaryan-549 · 2025-12-10T08:20:43Z

Fixes #750
Implements derivative of ellipsoid-based fluid forces with respect to velocities. Includes all component derivatives:

Added mass forces
Viscous torque and drag
Kutta lift
Magnus force

The implementation adds the fluid force derivatives to qDeriv, following the pattern from the reference C implementation.

thowell · 2025-12-12T14:40:46Z

+    slender_drag_coef = geom_fluid[geomid, 2]
+    ang_drag_coef = geom_fluid[geomid, 3]
+
+    if density > 0.0 and magnus_coef != 0.0:


lets structure this with one if density > 0.0 check and checks for magnus_coef, kutta_coef and viscosity in the density check scope.

thowell · 2025-12-12T14:41:56Z



+@wp.func
+def _pow2_deriv(val: float) -> float:


this function is unused, lets remove it

thowell · 2025-12-12T14:46:29Z

+  lin_vel = wp.spatial_bottom(local_vels)
+  ang_vel = wp.spatial_top(local_vels)
+
+  virtual_lin_mom = wp.vec3(


fluid_density * wp.cw_mul(virtual_mass, lin_vel)

thowell · 2025-12-12T14:46:57Z

+    fluid_density * virtual_mass[1] * lin_vel[1],
+    fluid_density * virtual_mass[2] * lin_vel[2]
+  )
+  virtual_ang_mom = wp.vec3(


fluid_density * wp.cw_mul(virtual_inertia, ang_vel)

thowell · 2025-12-12T14:53:38Z

+    ang_drag_coef * II[2] + slender_drag_coef * (I_max - II[2])
+  )
+
+  mom_visc = wp.vec3(x * mom_coef[0], y * mom_coef[1], z * mom_coef[2])


wp.cw_mul(ang, mom_coef)

thowell · 2025-12-12T14:54:25Z

+
+  mom_visc = wp.vec3(x * mom_coef[0], y * mom_coef[1], z * mom_coef[2])
+  norm = wp.length(mom_visc)
+  density = fluid_density / wp.max(wp.static(1e-10), norm)


lets utilize MJ_MINVAL instead of wp.static(1e-10)

thowell · 2025-12-12T14:56:13Z

+  norm = wp.length(mom_visc)
+  density = fluid_density / wp.max(wp.static(1e-10), norm)
+
+  mom_sq = wp.vec3(


-density * wp.cw_mul(wp.cw_mul(ang, mom_coef), mom_coef)

thowell · 2025-12-12T14:58:20Z

+
+  proj_denom = aa * xx + bb * yy + cc * zz
+  proj_num = a * xx + b * yy + c * zz
+  dA_coef = wp.pi / wp.max(wp.static(1e-10), wp.sqrt(proj_num * proj_num * proj_num * proj_denom))


lets utilize MJ_MINVAL

thowell · 2025-12-12T14:58:56Z

+  proj_num = a * xx + b * yy + c * zz
+  dA_coef = wp.pi / wp.max(wp.static(1e-10), wp.sqrt(proj_num * proj_num * proj_num * proj_denom))
+
+  A_proj = wp.pi * wp.sqrt(proj_denom / wp.max(wp.static(1e-10), proj_num))


lets utilize MJ_MINVAL

thowell · 2025-12-12T14:59:14Z

+  A_proj = wp.pi * wp.sqrt(proj_denom / wp.max(wp.static(1e-10), proj_num))
+
+  norm = wp.sqrt(xx + yy + zz)
+  inv_norm = 1.0 / wp.max(wp.static(1e-10), norm)


lets utilize MJ_MINVAL

thowell · 2025-12-12T15:01:16Z

+    xz, yz, zz + inner
+  )
+
+  D = D * (-quad_coef * inv_norm)


thowell · 2025-12-12T15:05:50Z

+  proj_num = a * xx + b * yy + c * zz
+  norm2 = xx + yy + zz
+  df_denom = wp.pi * kutta_lift_coef * fluid_density / wp.max(
+    wp.static(1e-10), wp.sqrt(proj_denom * proj_num * norm2)


MJ_MINVAL

thowell · 2025-12-12T15:06:20Z

+  dfx_coef = yy * (a - b) + zz * (a - c)
+  dfy_coef = xx * (b - a) + zz * (b - c)
+  dfz_coef = xx * (c - a) + yy * (c - b)
+  proj_term = proj_num / wp.max(wp.static(1e-10), proj_denom)


MJ_MINVAL

thowell · 2025-12-12T15:06:27Z

+  dfy_coef = xx * (b - a) + zz * (b - c)
+  dfz_coef = xx * (c - a) + yy * (c - b)
+  proj_term = proj_num / wp.max(wp.static(1e-10), proj_denom)
+  cos_term = proj_num / wp.max(wp.static(1e-10), norm2)


MJ_MINVAL

thowell · 2025-12-12T15:11:52Z

+  proj_term = proj_num / wp.max(wp.static(1e-10), proj_denom)
+  cos_term = proj_num / wp.max(wp.static(1e-10), norm2)
+
+  D = wp.mat33(


a - a, b - b and c - c can be replaced with 0.0

thowell · 2025-12-12T15:12:21Z

+    a - b, b - b, c - b,
+    a - c, b - c, c - c
+  )
+  D = D * (wp.static(2.0) * proj_num)


*= and we can remove the parentheses

thowell · 2025-12-12T15:12:43Z

+  )
+  D = D * (wp.static(2.0) * proj_num)
+
+  inner_term = wp.vec3(


we can utilize wp.cw_mul here

thowell · 2025-12-12T15:15:47Z

+  lin_vel_scaled = lin_vel * magnus_coef
+  ang_vel_scaled = ang_vel * magnus_coef
+
+  D_ang = _cross_deriv_a(ang_vel_scaled, lin_vel)


should lin_vel -> lin_vel_scaled?

thowell · 2025-12-12T15:16:01Z

+  ang_vel_scaled = ang_vel * magnus_coef
+
+  D_ang = _cross_deriv_a(ang_vel_scaled, lin_vel)
+  D_lin = _cross_deriv_b(ang_vel, lin_vel_scaled)


should ang_vel -> ang_vel_scaled?

thowell · 2025-12-12T15:21:28Z

+
+
+@wp.func
+def _cross_deriv_a(a: wp.vec3, b: wp.vec3) -> wp.mat33:


lets consider combining _cross_deriv_a and _cross_deriv_b into one function that returns 2 wp.mat33 since it looks like we always compute derivatives wrt to both cross product inputs

thowell

@Aaryan-549 thank you for contributing this pr!

in addition to the comments, lets add a test. please let us know if you have any questions. thanks!

Aaryan-549 · 2026-01-02T14:28:02Z

Hey @thowell, really sorry for the delay. I have implemented all fixes that you asked for. Let me know If there are any other changes needed.

erikfrey · 2026-01-05T22:44:08Z

@Aaryan-549 thank you for the contribution. It looks like there are some ruff linter issues, kernel analyzer issues, and unit test failures.

You should be able to run these three tools locally to reproduce and fix the errors reported in CI. Please give that a go - once the issues are fixed we'll take another look at the PR.

ruff and pytest issues should be straightforward. To install and use the kernel analyzer, see the readme here.

Aaryan-549 · 2026-01-07T12:28:16Z

Thanks, apologies for the noise. I'll run the tools locally and fix the errors shortly.

Aaryan-549 · 2026-01-08T06:11:08Z

I've run all three tools locally and fixed the issues.

Ruff formatting and linting both pass. I've fixed 21 failing tests and now have 925/926 tests passing. The one remaining failure (test_collision19) is a pre-existing collision bug that was already failing before my commits.

The kernel analyzer shows 3 style warnings in my code - two missing "# In:" comments and one parameter ordering convention issue. These don't break anything but I can fix them if you'd like, or address them in a follow-up after review.

thowell · 2026-01-13T14:33:18Z

@Aaryan-549 thank you for the update!

all of the checks should pass including tests, formatting, and the linter. pre-commit will provide information about what should be fixed, please resolve any issues this tool reports.

thanks!

thowell · 2026-01-15T15:11:35Z

    mjw.step(m, d)
    self.assertGreater(d.time.numpy()[0], 0.0)

+  @parameterized.product(


since the implicit integrator is not implemented in this pr we can remove this test

thowell · 2026-01-15T15:13:04Z

    _assert_eq(mjw_out, mj_out, "qM - dt * qDeriv")

+  @parameterized.parameters(
+    (mujoco.mjtJacobian.mjJAC_DENSE, "ellipsoid"),


no need to parameterized 'fluidshape' since both cases are ellipsoid

thowell · 2026-01-15T15:13:16Z

+          <site name="site2" pos="0 0 1"/>
+        </body>
+      </worldbody>
+      <tendon>


do we need tendons for this test?

SUZ-tsinghua · 2026-01-15T15:53:57Z

Hi @thowell , I'd happy to review this pr. @Aaryan-549 For the derivative test, there is no need to include tendon and actuator. You can just copy my test.

Aaryan-549 · 2026-01-15T17:08:51Z

Thanks @thowell! I've addressed all the latest feedback in the recent push:

Reverted changes to io.py, io_test.py, and forward_test.py. I agree it's better to keep the implicit integrator support for a separate PR.
Simplified derivative_test.py by removing unnecessary tendons, actuators, and the redundant fluidshape parameterization.
Removed unnecessary tendons and actuators from the test XML as suggested.
Reverted the unrelated formatting changes in kernel_analyzer.

Checks are passing. Ready for final review!

thowell · 2026-01-16T13:30:05Z

please revert all changes to this file

thowell · 2026-01-16T13:31:08Z

we shouldn't need any changes to this file

thowell · 2026-01-16T13:31:21Z

we shouldn't need any changes to this file

Aaryan-549 · 2026-01-23T21:04:36Z

Hey @thowell ,

I am really sorry for taking a whole week to reply back but I've been working on implementing the mjd_ellipsoidFluid derivative for PR #897, and I've run into a persistent issue that I could use some guidance on.

The Problem:
The test_smooth_vel_fluid tests are failing with ~50% of matrix elements mismatched. The diagonal elements match perfectly, but the off-diagonal elements have incorrect values. I've traced through the code extensively and believe the issue may be related to coordinate frame conventions or matrix storage order differences between MuJoCo and Warp.

What I've Verified/Fixed:

Integrator parameter type (changed from wp.array(dtype=int) to int)
Integrator enum check (using == 2 for IMPLICITFAST)
Dense vs sparse iteration (created separate kernels to iterate over all nv×nv DOF pairs for dense mode)
Attempted to match MuJoCo's column-major B matrix storage by transposing D in _add_to_quadrant

Current State:
The B matrix derivative functions (_added_mass_forces_deriv, _viscous_drag_deriv, etc.) appear to match MuJoCo's implementations, and the Jacobian computation follows the same approach. However, something about how the pieces fit together is producing wrong off-diagonal values.

Side Note:
I also noticed that when I ran @SUZ-tsinghua's original branch locally, there were 32 failing tests out of the box. Were the tests passing on your end when you reviewed? Just want to make sure I'm not chasing issues that existed before my changes.

Any pointers on what might be causing the mismatch would be really appreciated. Happy to hop on a call if that would be easier.

Thanks!

SUZ-tsinghua · 2026-01-24T03:50:42Z

Hi @Aaryan-549 , on my side, I only encountered 3 failed io tests on my branch. I have fixed these errors by merging the original branch into my branch.

Aaryan-549 · 2026-01-24T05:54:47Z

Hey @SUZ-tsinghua even I am facing 3 failed tests on my end(for my code). Do you have any idea of what could be causing these?(we only need to take care of 2 as 1 of the test was failing in main too, so I don't think we caused it).

Implements derivative of ellipsoid-based fluid forces with respect to velocities. Includes all component derivatives: - Added mass forces - Viscous torque and drag - Kutta lift - Magnus force The implementation adds the fluid force derivatives to qDeriv, following the pattern from the reference C implementation.

…J_MINVAL, restructure density checks, combine cross derivatives

…tion

…alues need debugging)

thowell · 2026-04-29T20:42:25Z

@Aaryan-549 updated the pr branch. please take a look at the changes. thanks!

Aaryan-549 · 2026-04-29T20:43:51Z

@thowell Thanks for the changes. I'll take a look at them shortly.

Aaryan-549 · 2026-05-01T17:54:38Z

Looks good to me, Thanks a lot for the assist!

thowell self-requested a review December 11, 2025 16:29