Skip to content

Sample the environment layer field in the fluid models - #3461

Draft
teerthsharma wants to merge 2 commits into
google-deepmind:mainfrom
teerthsharma:feat/env-sampler
Draft

Sample the environment layer field in the fluid models#3461
teerthsharma wants to merge 2 commits into
google-deepmind:mainfrom
teerthsharma:feat/env-sampler

Conversation

@teerthsharma

@teerthsharma teerthsharma commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Second of two steps toward #3452. #3460 adds the layer stack; this adds the reader and points the fluid models at it.

Relates to #3452. Depends on #3458 for sanitizer builds.

Stacked on #3460. Review that first; this branch contains only the commit on top of it.

Summary

#3460 adds the layer stack to mjModel but nothing reads it. This adds the reader and points the fluid models at it.

mj_envLayer returns the layer containing a height, and mj_envSample fills an mjEnv with the gravity, wind, density and viscosity in force at a world point. mj_ellipsoidFluidModel and mj_inertiaBoxFluidModel now sample at each geom's own position instead of reading mjOption, and mjd_ellipsoidFluid and mjd_inertiaBoxFluid sample at the same points so force and derivative stay consistent.

When nlayer <= 1, mj_envSample copies mjOption verbatim, so single-medium models take the same values they always did.

Gravity is deliberately not converted. It does not enter as a force at a point: mj_rne sets the world frame's fictitious acceleration to -g and propagates it down the tree, so each body picks it up through cinert * cacc, an integral over its mass distribution with no sample point to name. Three further reads are locked to that one — mj_gravcomp must cancel it exactly, mj_energyPos must match it, and the mj_fwdActuation gate — so sampling any one at xipos would desynchronise the others for every body spanning a boundary. Making gravity a field needs the integral, not a sample, and is left out.

mj_fluid and mjd_passive_vel each early-returned on mjOption alone, so a stratified model with opt.density == 0 would have received no fluid force at all. Both now gate on mj_hasFluid, which also consults the layers. At nlayer == 0 the layer loop does not execute and the test is identical to the old one.

Determinism

mj_envLayer performs no floating-point arithmetic: it reads heights, compares, and shifts integers. There is no add or multiply on mjtNum, so nothing is exposed to FMA contraction, x87 excess precision or vectorized reassociation, and the result is bit-identical across compilers, platforms and optimization levels. Index arithmetic is done in mjtSize throughout, so nothing is narrowed. The loop runs exactly ceil(log2(nlayer-1)) times, with no data-dependent iteration count, no global state and no warm start. A height exactly on a boundary belongs to the upper layer; +inf selects the top layer; -inf and NaN select layer 0, because every comparison against NaN is false and the search then always takes the upper branch.

Validation

check result
CTest, Linux, gcc 15.2.0, Release 1991/1991, 0 failed
EnvironmentTest under ASan + UBSan + LeakSanitizer, clang 21.1.8 14/14, 0 diagnostics

The sanitizer run is only meaningful with a positive control, so: the test binary carries 224 __asan/__ubsan symbols, and a deliberately planted heap-buffer-overflow compiled with the same flags is caught and reported. Building it required #3458 applied locally, -DCMAKE_INTERPROCEDURAL_OPTIMIZATION=OFF to avoid a ThinLTO link failure, and CMAKE_*_FLAGS_INIT rather than CMAKE_*_FLAGS so CMake's own defaults survive.

The 14 tests cover the boundary tie convention at ULP resolution, binary search against a linear scan over 64 models with 1 to 64 layers and 39 queries each, bit-identical mjOption fallback at nlayer 0 and 1 with poison values in the layer arrays, force and derivative each following the field, and seven adversarial cases: non-finite query heights, index always in range, monotonicity on adjacent points, statelessness across interleaved calls, -0.0 and denormal and 1-ulp-apart boundaries, single-medium purity, and idempotent sampling.

Test quality was measured by mutation rather than asserted. Eight mutations of the implementation were applied and reverted:

mutation caught by
layer_height[mid] <= z to < z 4 tests
delete the nlayer <= 1 early return 2 tests
hi = nbound to hi = nbound + 1 6 tests
NaN branch flip, !(z < h[mid]) 1 test, uniquely
memoize (z, layer) across calls 1 test, uniquely
drop env->viscosity on the layered path 5 tests

No test survives every mutation. Two mutations of an earlier revision survived, and investigating why showed the code they targeted was unreachable: a clamp on nlayer made the midpoint overflow it guarded against impossible, while itself converting a narrowing conversion into an out-of-bounds read. The clamp is gone and the search runs in mjtSize.

Cost

Measured by calling mj_envSample over 256 real geom positions, 2000 rounds per batch, best of 11, pinned at real-time priority:

nlayer ns per sample marginal search iterations
0 2.862 0, early return
2 4.875 +2.013 1
8 6.633 +3.771 3

+0.88 ns per search iteration, matching O(log nlayer). One sample per geom per step under Euler, two under implicit and implicitfast. On a 256-geom fluid scene that predicts +0.6% per step, which is below what step-level timing can resolve on this host; the step-level measurement of the same scenes indeed returns values inside its own noise, including physically impossible negative ones.

Limits

Gravity remains global, as above. mjOption values are still read by everything outside the two fluid models. The stratification is horizontal slabs along +z only.

The single-medium claim is exercised by the shipped corpus, where every model has nlayer == 0, and by the bit-identical trajectory checksums reported on #3460; the layered path is exercised only by the new tests and by synthetic scenes.

Sanitizer coverage is the EnvironmentTest binary, not the full suite, which did not complete under sanitizers on this host within a reasonable time. MJX and Warp are not extended.

mjOption carries one gravity, wind, density and viscosity that hold everywhere
in the model. This adds an ordered stack of horizontal medium slabs that let
those four quantities vary with height along +z.

mjModel gains nlayer and the layer_height, layer_gravity, layer_density,
layer_viscosity and layer_wind arrays. mjSpec gains mjsLayer together with
mjs_addLayer, mjs_getLayer and mjs_deleteLayers. <option> accepts any number of
<layer> children, each inheriting its unspecified attributes from the global
medium, and the compiler rejects heights that are not strictly ascending.
Layers are unnamed global option data rather than model elements: they have no
mjtObj type, no name and no id, hence no mjsElement, and the stack is cleared
at once rather than deleted element by element. The compiler holds them in a
deque so that pointers returned by mjs_addLayer survive later adds.

Nothing reads the new fields yet. A model with no <layer> element compiles to
nlayer == 0 and is unchanged.

Signed-off-by: teerthsharma <teerths57@gmail.com>
mjModel carries the layer stack but nothing reads it. mj_inertiaBoxFluidModel
and mj_ellipsoidFluidModel take wind, density and viscosity straight from
mjOption, so a layered model still behaves as one global medium. This adds the
reader and wires it into the two fluid models and their derivatives.

mj_envLayer binary-searches the nlayer-1 interior boundaries for the layer
containing a height, with a height exactly on a boundary opening the upper
layer. mj_envSample fills an mjEnv with that layer's gravity, wind, density and
viscosity. Both fluid models sample at the point the force is applied to, xipos
for the inertia box and geom_xpos for the ellipsoid, and mjd_inertiaBoxFluid
and mjd_ellipsoidFluid sample at the same points, so the derivative continues
to match the force. mj_fluid previously early-returned on mjOption density and
viscosity alone, which a layered model that leaves the global medium empty
would wrongly take; mj_hasFluid asks the layers as well.

The stack is stratified along +z rather than a general cell arrangement because
a piecewise-constant gravity field has a potential only when every value jump is
parallel to the normal of the face it crosses. Parallel layers satisfy that for
any assignment of values; a vertical split does not, and a closed circuit across
one gains energy without bound.

A model with nlayer == 0 takes the single-medium path, in which mj_envSample
copies mjOption bit for bit, so existing models are unaffected.

Signed-off-by: teerthsharma <teerths57@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant