chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement midpoint integrator for free bodies.

Browse Source 
PiperOrigin-RevId: 899043541
Change-Id: I0bb38f6ad94e189b45ab16777a04ad6fefc6adf7
This commit is contained in:
Yuval Tassa
2026-04-13 09:37:51 -07:00
committed by Copybara-Service
parent d9b5d8babb
commit 0c337799bd
12 changed files with 913 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/engine/engine_derivative_test.cc
+78 -2
View File
@@ -320,17 +320,19 @@ TEST_F(DerivativeTest, PassiveDvel) {
mj_forward(model, data);
// get analytic derivatives
mju_zero(data->qDeriv, model->nD);
mjd_passive_vel(model, data);
mju_copy(qDerivAnalytic, data->qDeriv, nD);
// clear qDeriv, get finite-difference derivatives
mju_zero(data->qDeriv, nD);
mju_zero(qDerivFD, nD);
mjtNum eps = MjTol(1e-6, 1e-3);
mjtNum eps = MjTol(1e-6, 1e-4);
mjd_passive_velFD(model, data, eps);
// expect FD and analytic derivatives to be similar to tol precision
EXPECT_THAT(AsVector(data->qDeriv, nD),
Pointwise(MjNear(1e-4, 1e-3), AsVector(qDerivAnalytic, nD)));
Pointwise(MjNear(1e-6, 1e-4), AsVector(qDerivAnalytic, nD)));
}
mju_free(qDerivFD);
@@ -1733,5 +1735,79 @@ TEST_F(DerivativeTest, FlexInterpDerivativesDeformed) {
mj_deleteModel(model);
}
TEST_F(DerivativeTest, MidpointFluidAccuracy) {
const std::string xml_path =
GetTestDataFilePath(kTumblingThinObjectEllipsoidPath);
char error[1024];
mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
ASSERT_THAT(m, NotNull()) << error;
mjtNum dt_small = 1e-4;
mjtNum dt_large = m->opt.timestep; // 2e-3, the default
mjtNum duration = 0.5;
mjData* d_ref = mj_makeData(m);
mjData* d_midpoint = mj_makeData(m);
mjData* d_nomidpoint = mj_makeData(m);
// give initial angular velocity for tumbling
mj_resetData(m, d_ref);
mj_resetData(m, d_midpoint);
mj_resetData(m, d_nomidpoint);
d_ref->qvel[3] = 5;
d_ref->qvel[4] = 3;
d_ref->qvel[5] = 1;
d_midpoint->qvel[3] = 5;
d_midpoint->qvel[4] = 3;
d_midpoint->qvel[5] = 1;
d_nomidpoint->qvel[3] = 5;
d_nomidpoint->qvel[4] = 3;
d_nomidpoint->qvel[5] = 1;
int nsteps_large = static_cast<int>(duration / dt_large);
int substeps = static_cast<int>(dt_large / dt_small);
mjtNum error_midpoint = 0;
mjtNum error_nomidpoint = 0;
for (int i = 0; i < nsteps_large; i++) {
// reference: RK4 at small timestep
m->opt.integrator = mjINT_RK4;
m->opt.timestep = dt_small;
m->opt.enableflags &= ~mjENBL_INVDISCRETE;
for (int j = 0; j < substeps; j++) {
mj_step(m, d_ref);
}
// implicit with midpoint (default)
m->opt.integrator = mjINT_IMPLICIT;
m->opt.timestep = dt_large;
m->opt.enableflags &= ~mjENBL_INVDISCRETE;
mj_step(m, d_midpoint);
// implicit without midpoint
m->opt.enableflags |= mjENBL_INVDISCRETE;
mj_step(m, d_nomidpoint);
// accumulate position errors
for (int k = 0; k < 7; k++) {
mjtNum diff_mid = d_ref->qpos[k] - d_midpoint->qpos[k];
mjtNum diff_nomid = d_ref->qpos[k] - d_nomidpoint->qpos[k];
error_midpoint += diff_mid * diff_mid;
error_nomidpoint += diff_nomid * diff_nomid;
}
}
// expect midpoint to be more accurate
EXPECT_LT(error_midpoint, error_nomidpoint)
<< "implicit midpoint should be more accurate than implicit without "
<< "midpoint for a free body with fluid forces";
mj_deleteData(d_nomidpoint);
mj_deleteData(d_midpoint);
mj_deleteData(d_ref);
mj_deleteModel(m);
}
} // namespace
} // namespace mujoco
test/engine/engine_forward_test.cc
+305
View File
@@ -17,6 +17,7 @@
#include "src/engine/engine_forward.h"
#include "src/engine/engine_derivative.h"
#include <algorithm>
#include <array>
#include <cmath>
#include <cstdlib>
@@ -478,6 +479,310 @@ TEST_F(ImplicitIntegratorTest, EnergyConservation) {
mj_deleteModel(model);
}
// Energy and angmom conservation for free body with implicitfast (IMR)
TEST_F(ImplicitIntegratorTest, ConservationMidpoint) {
// aligned: CoM at joint origin
static constexpr char xml1[] = R"(
<mujoco>
<option integrator="implicitfast" timestep="0.01">
<flag energy="enable" gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30"/>
</body>
</worldbody>
</mujoco>
)";
// auto-aligned: CoM at joint origin
static constexpr char xml2[] = R"(
<mujoco>
<option integrator="implicitfast" timestep="0.01">
<flag energy="enable" gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint align="true"/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30" pos=".03 .02 .01"/>
</body>
</worldbody>
</mujoco>
)";
// non-aligned: CoM offset from joint origin
static constexpr char xml3[] = R"(
<mujoco>
<option integrator="implicitfast" timestep="0.01">
<flag energy="enable" gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30" pos=".03 .02 .01"/>
</body>
</worldbody>
</mujoco>
)";
int xml_idx = 1;
for (auto xml : {xml1, xml2, xml3}) {
SCOPED_TRACE(testing::Message() << "XML case " << xml_idx++);
char error[1024];
mjModel* model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
mjData* data = mj_makeData(model);
const int nstep = 500;
mjtNum energy_drift[2], angmom_drift[2]; // [0]=midpoint, [1]=rk4
for (int integrator : {mjINT_IMPLICITFAST, mjINT_RK4}) {
int idx = (integrator == mjINT_IMPLICITFAST) ? 0 : 1;
model->opt.integrator = integrator;
// reset
mj_resetData(model, data);
data->qvel[3] = 1.0;
data->qvel[4] = 2.0;
data->qvel[5] = 3.0;
mj_forward(model, data);
mjtNum initial_energy = data->energy[1];
mjtNum initial_angmom[3];
mj_subtreeVel(model, data);
mju_copy3(initial_angmom, data->subtree_angmom);
for (int i=0; i < nstep; i++) {
mj_step(model, data);
}
energy_drift[idx] = fabs(data->energy[1] - initial_energy);
mj_subtreeVel(model, data);
mjtNum angmom_err[3];
mju_sub3(angmom_err, data->subtree_angmom, initial_angmom);
angmom_drift[idx] = mju_norm3(angmom_err);
}
// midpoint should conserve energy better than RK4 (double only)
#ifndef mjUSESINGLE
EXPECT_LT(energy_drift[0], energy_drift[1]);
#endif
// both should conserve angular momentum well
EXPECT_LT(angmom_drift[0], MjTol(1e-3, 1e-2));
EXPECT_LT(angmom_drift[1], MjTol(1e-3, 1e-2));
mj_deleteData(data);
mj_deleteModel(model);
}
}
// verify second-order convergence of midpoint integration
TEST_F(ImplicitIntegratorTest, MidpointConvergenceOrder) {
// aligned: CoM at joint origin
static constexpr char xml1[] = R"(
<mujoco>
<option integrator="implicitfast">
<flag gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30"/>
</body>
</worldbody>
</mujoco>
)";
// non-aligned: CoM offset from joint origin
static constexpr char xml2[] = R"(
<mujoco>
<option integrator="implicitfast">
<flag gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30"
pos=".05 .03 .02"/>
</body>
</worldbody>
</mujoco>
)";
int xml_idx = 1;
for (auto xml : {xml1, xml2}) {
SCOPED_TRACE(testing::Message() << "XML case " << xml_idx++);
char error[1024];
mjModel* model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
mjtNum T = 1.0;
mjtNum h_coarse = 0.02;
mjtNum quat_coarse[4], quat_fine[4], quat_ref[4];
auto run = [&](mjtNum h, mjtNum quat_out[4]) {
model->opt.timestep = h;
mjData* data = mj_makeData(model);
data->qvel[3] = 1.0;
data->qvel[4] = 2.0;
data->qvel[5] = 3.0;
int nstep = (int)(T / h + 0.5);
for (int i = 0; i < nstep; i++) {
mj_step(model, data);
}
mju_copy4(quat_out, data->qpos + 3);
mj_deleteData(data);
};
run(h_coarse, quat_coarse);
run(h_coarse / 2, quat_fine);
run(h_coarse / 16, quat_ref);
// quaternion distance: ||quat - quat_ref|| (handles sign ambiguity)
auto quat_dist = [](const mjtNum a[4], const mjtNum b[4]) -> mjtNum {
mjtNum pos = 0, neg = 0;
for (int i = 0; i < 4; i++) {
pos += (a[i] - b[i]) * (a[i] - b[i]);
neg += (a[i] + b[i]) * (a[i] + b[i]);
}
return mju_sqrt(mju_min(pos, neg));
};
mjtNum err_coarse = quat_dist(quat_coarse, quat_ref);
mjtNum err_fine = quat_dist(quat_fine, quat_ref);
// second-order: error ratio should be ~4 when halving timestep
mjtNum ratio = err_coarse / err_fine;
EXPECT_GT(ratio, 3.5);
EXPECT_LT(ratio, 4.5);
mj_deleteModel(model);
}
}
// verify that Newton iteration in mj_midpoint converges quickly (aligned case)
TEST_F(ImplicitIntegratorTest, MidpointNewtonConvergence) {
// inertia ratios: symmetric, mildly asymmetric, extremely asymmetric
mjtNum inertias[][3] = {
{1.0, 1.0, 1.0},
{1.0, 2.0, 3.0},
{0.01, 1.0, 100.0},
{1.0, 1.0, 1000.0},
};
mjtNum timesteps[] = {0.001, 0.01, 0.1};
mjtNum velocities[][3] = {
{1.0, 2.0, 3.0},
{100.0, 0.0, 0.0},
{10.0, 10.0, 10.0},
{0.01, 0.01, 100.0},
};
mjtNum q_identity[4] = {1, 0, 0, 0};
mjtNum torques[][3] = {
{0, 0, 0},
{10.0, 20.0, 30.0},
{100.0, 0.0, 0.0},
{0.0, 0.0, 100.0},
};
int max_iter = 0;
int total_iter = 0;
int ncases = 0;
for (auto& I : inertias) {
for (mjtNum h : timesteps) {
for (auto& w : velocities) {
for (auto& tau : torques) {
mjtNum vel[6] = {0, 0, 0, w[0], w[1], w[2]};
mjtNum tau_ext[6] = {0, 0, 0, tau[0], tau[1], tau[2]};
mjtNum v_new[6];
mjtNum ipos[3] = {0, 0, 0};
int niter = mj_midpoint(1.0, I, ipos, q_identity, q_identity, vel,
tau_ext, NULL, h, v_new);
EXPECT_LT(niter, 10)
<< "Failed for I=(" << I[0] << "," << I[1] << "," << I[2] << ")"
<< " h=" << h
<< " w=(" << w[0] << "," << w[1] << "," << w[2] << ")"
<< " tau=(" << tau[0] << "," << tau[1] << "," << tau[2] << ")";
max_iter = std::max(max_iter, niter);
total_iter += niter;
ncases++;
}
}
}
}
EXPECT_LE(max_iter, 4);
EXPECT_LT((mjtNum)total_iter / ncases, 2.0);
}
// verify that Newton iteration in mj_midpoint converges quickly (non-aligned)
TEST_F(ImplicitIntegratorTest, MidpointFullNewtonConvergence) {
mjtNum masses[] = {0.1, 1.0, 10.0};
mjtNum inertias[][3] = {
{1.0, 1.0, 1.0},
{1.0, 2.0, 3.0},
{0.01, 1.0, 100.0},
};
mjtNum offsets[][3] = {
{0.1, 0.0, 0.0},
{0.05, 0.03, 0.02},
{0.0, 0.0, 0.5},
};
mjtNum timesteps[] = {0.001, 0.01, 0.1};
mjtNum velocities[][6] = {
{1.0, 0.0, 0.0, 1.0, 2.0, 3.0},
{0.0, 0.0, 0.0, 10.0, 10.0, 10.0},
{5.0, 5.0, 5.0, 0.01, 0.01, 100.0},
};
mjtNum q_identity[4] = {1, 0, 0, 0};
mjtNum forces[][6] = {
{0, 0, 0, 0, 0, 0},
{10.0, 20.0, 30.0, 1.0, 2.0, 3.0},
};
int max_iter = 0;
int total_iter = 0;
int ncases = 0;
for (mjtNum mass : masses) {
for (auto& I : inertias) {
for (auto& r : offsets) {
for (mjtNum h : timesteps) {
for (auto& vel : velocities) {
for (auto& frc : forces) {
mjtNum v_new[6];
int niter = mj_midpoint(mass, I, r, q_identity, q_identity,
vel, frc, NULL, h, v_new);
EXPECT_LT(niter, 10)
<< "Failed for mass=" << mass
<< " I=(" << I[0] << "," << I[1] << "," << I[2] << ")"
<< " r=(" << r[0] << "," << r[1] << "," << r[2] << ")"
<< " h=" << h;
max_iter = std::max(max_iter, niter);
total_iter += niter;
ncases++;
}
}
}
}
}
}
EXPECT_LE(max_iter, 6);
EXPECT_LT((mjtNum)total_iter / ncases, 3.0);
}
TEST_F(ForwardTest, ControlClamping) {
static constexpr char xml[] = R"(
<mujoco>
test/engine/engine_inverse_test.cc
+8 -9
View File
@@ -73,9 +73,16 @@ TEST_F(InverseTest, DiscreteInverseMatch) {
mjtNum* qvel_next = (mjtNum*)mju_malloc(nv * sizeof(mjtNum));
mjtNum* qacc_fd = (mjtNum*)mju_malloc(nv * sizeof(mjtNum));
for (auto integrator : {mjINT_EULER, mjINT_IMPLICIT, mjINT_IMPLICITFAST}) {
for (auto integrator : {mjINT_EULER, mjINT_IMPLICIT}) {
model->opt.integrator = integrator;
for (bool invdiscrete : {false, true}) {
// set/unset mjENBL_INVDISCRETE flag (affects both forward and inverse)
if (invdiscrete) {
model->opt.enableflags |= mjENBL_INVDISCRETE;
} else {
model->opt.enableflags &= ~mjENBL_INVDISCRETE;
}
// simulate
mj_resetData(model, data);
for (int i = 0; i < kSteps; ++i) {
@@ -98,17 +105,9 @@ TEST_F(InverseTest, DiscreteInverseMatch) {
mj_forward(model, data);
mju_copy(data->qacc, qacc_fd, nv);
// set/unset mjENBL_INVDISCRETE flag
if (invdiscrete) {
model->opt.enableflags |= mjENBL_INVDISCRETE;
} else {
model->opt.enableflags &= ~mjENBL_INVDISCRETE;
}
// call built-in testing function
mj_compareFwdInv(model, data);
// depending on mjENBL_INVDISCRETE flag, expect mismatch to be small/large
if (invdiscrete) {
mjtNum epsilon = MjTol(1e-9, 0.05);
EXPECT_LT(data->solver_fwdinv[0], epsilon);
test/engine/engine_sleep_test.cc
+53
View File
@@ -533,6 +533,59 @@ TEST_F(SleepTest, Equality) {
mj_deleteModel(m);
}
// Test that the midpoint integrator doesn't break the sleep qvel=0 invariant.
// A standalone free body (eligible for midpoint) with high viscosity should
// eventually go to sleep, and after sleeping, qvel/qacc must be exactly zero.
TEST_F(SleepTest, MidpointSleepZeroVelocity) {
static constexpr char xml[] = R"(
<mujoco>
<option integrator="implicitfast" viscosity="10"
sleep_tolerance="0.01">
<flag sleep="enable" gravity="disable" constraint="disable"
contact="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30"
pos=".03 .02 .01"/>
</body>
</worldbody>
</mujoco>
)";
char error[1024];
mjModel* m = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(m, NotNull()) << error;
mjData* d = mj_makeData(m);
// give initial velocity (both translational and angular)
d->qvel[0] = 0.5;
d->qvel[1] = 0.5;
d->qvel[2] = 0.5;
d->qvel[3] = 1.0;
d->qvel[4] = 2.0;
d->qvel[5] = 3.0;
// step until body goes to sleep
for (int step = 0; step < 1000; step++) {
mj_step(m, d);
if (d->ntree_awake == 0) break;
}
// body should have gone to sleep
ASSERT_EQ(d->ntree_awake, 0) << "body did not go to sleep";
// qvel and qacc must be exactly zero for sleeping body
for (int i = 0; i < 6; i++) {
EXPECT_EQ(d->qvel[i], 0.0) << "qvel[" << i << "] not zero after sleep";
EXPECT_EQ(d->qacc[i], 0.0) << "qacc[" << i << "] not zero after sleep";
}
mj_deleteData(d);
mj_deleteModel(m);
}
static const char* const kInitIslandFailModel =
"engine/testdata/sleep/init_island_fail.xml";