Mujoco_WASM/test/engine/engine_inverse_test.cc

// Copyright 2023 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Tests for engine/engine_inverse.c.

#include "src/engine/engine_inverse.h"

#include <string>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mujoco.h>
#include "test/fixture.h"

namespace mujoco {
namespace {

using ::testing::NotNull;
using InverseTest = MujocoTest;

const int kSteps = 70;
static const char* const kModelPath = "testdata/model.xml";

// test standard continuous-time inverse dynamics
TEST_F(InverseTest, ForwardInverseMatch) {
  const std::string xml_path = GetTestDataFilePath(kModelPath);
  char error[1024];
  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
  ASSERT_THAT(model, NotNull()) << error;
  mjData* data = mj_makeData(model);

  // set small tolerance and enough iterations for all solvers to converge
  model->opt.iterations = 500;
  model->opt.tolerance = 0;

  // solver names for diagnostics
  const char* solver_name[] = {"PGS", "CG", "Newton"};

  for (int diagexact = 0; diagexact < 2; diagexact++) {
    if (diagexact) {
      model->opt.enableflags |= mjENBL_DIAGEXACT;
    } else {
      model->opt.enableflags &= ~mjENBL_DIAGEXACT;
    }

    for (mjtSolver solver : {mjSOL_PGS, mjSOL_CG, mjSOL_NEWTON}) {
      model->opt.solver = solver;
      mj_resetData(model, data);

      // simulate, call mj_forward
      for (int i = 0; i < kSteps; ++i) {
        mj_step(model, data);
      }
      mj_forward(model, data);

      // call built-in testing function
      mj_compareFwdInv(model, data);

      // per-solver tolerances
      mjtNum epsilon;
      switch (solver) {
      case mjSOL_PGS:    epsilon = MjTol(1e-6, 1e-2);   break;
      case mjSOL_CG:     epsilon = MjTol(1e-9, 1e-1);   break;
      case mjSOL_NEWTON: epsilon = MjTol(1e-10, 1e-2);  break;
      }
      EXPECT_LT(data->solver_fwdinv[0], epsilon)
        << solver_name[solver] << " diagexact=" << diagexact;
      EXPECT_LT(data->solver_fwdinv[1], epsilon)
        << solver_name[solver] << " diagexact=" << diagexact;
    }
  }

  mj_deleteData(data);
  mj_deleteModel(model);
}

// test discrete-time inverse dynamics
TEST_F(InverseTest, DiscreteInverseMatch) {
  // load and allocate
  const std::string xml_path = GetTestDataFilePath(kModelPath);
  char error[1024];
  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
  ASSERT_THAT(model, NotNull()) << error;
  int nv = model->nv;
  mjData* data = mj_makeData(model);
  int nstate = mj_stateSize(model, mjSTATE_INTEGRATION);
  mjtNum* state = (mjtNum*)mju_malloc(nstate * sizeof(mjtNum));
  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}) {
    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) {
        mj_step(model, data);
      }

      // save state
      mj_getState(model, data, state, mjSTATE_INTEGRATION);

      // call step, save new qvel
      mj_step(model, data);
      mju_copy(qvel_next, data->qvel, nv);

      // reset the state, compute discrete-time (finite-differenced) qacc
      mj_setState(model, data, state, mjSTATE_INTEGRATION);
      mju_sub(qacc_fd, qvel_next, data->qvel, nv);
      mju_scl(qacc_fd, qacc_fd, 1/model->opt.timestep, nv);

      // call mj_forward, overwrite qacc with qacc_fd
      mj_forward(model, data);
      mju_copy(data->qacc, qacc_fd, nv);

      // call built-in testing function
      mj_compareFwdInv(model, data);

      if (invdiscrete) {
        mjtNum epsilon = MjTol(1e-9, 0.05);
        EXPECT_LT(data->solver_fwdinv[0], epsilon);
        EXPECT_LT(data->solver_fwdinv[1], epsilon);
      } else {
        EXPECT_GT(data->solver_fwdinv[0], 1.0);
        EXPECT_GT(data->solver_fwdinv[1], 1.0);
      }
    }
  }

  // deallocate
  mju_free(qacc_fd);
  mju_free(qvel_next);
  mju_free(state);
  mj_deleteData(data);
  mj_deleteModel(model);
}

}  // namespace
}  // namespace mujoco