chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
b04a18a00a55e6d4a23f56e5f5dd2d3a23d404bb
Mujoco_WASM/python/mujoco/functions.cc
T
Yuval Tassa 072e963fa0 Add SO3 transmission and native orientation actuator. 
https://youtu.be/17XpwnqyCXs

New transmission type mjTRN_SO3: a relative orientation, targeting a ball
joint or a site+refsite pair. It is the first transmission with more than
one force output: its length is the norm of the expmap vector of the
relative rotation and its moment axes are the 3 rows of the
relative rotational Jacobian, without projecting onto per-actuator gears.

New force law mjGAIN_SO3/mjBIAS_SO3: a geodesic PD servo, force =
kp * log(q_current^-1 * q_target) - kv * velocity, exact for arbitrary axis
combinations with a unique equilibrium at every commanded orientation.
Error, moment rows and velocity all live in the child frame (joint or
site): the right-difference error is the gradient of the geodesic
potential in that frame. The parent-frame (left) error is not: driving
child-frame torques with it pumps energy at large angles, settling into
steady-spinning limit cycles (the SO3LargeAngleConvergence test). The
integrator variant stores the 3D orientation setpoint in act (actnum = 3,
re-anchored to a bounded representative at integration time). Exposed in
MJCF as <orientation joint=|site=+refsite= kp kv|dampratio>, or via
<general gaintype="so3" biastype="so3">.

The setpoint input has two charts: an expmap target (3 controls, default)
or a quaternion target (4 controls) -- <orientation input="quat">, the
first actuator with different input and output widths. The signature is
recorded in a new per-actuator field actuator_ctrlspec (mjtCtrlChart),
whose meaning is scoped by the gain type the way gain/bias parameters are;
ctrlnum is derived from it at compile time and remains the layout
authority. An explicit field rather than width inference or a prm slot:
width-as-chart cannot express same-width signatures (upcoming servo input
subsets), and prm slots are the input_mode pattern this stack retires.
The force law normalizes the commanded quaternion, making it scale- and
antipodally-invariant. The all-zero ctrl still maps to the identity via
mju_normalize4, but it is a degenerate point (a nudge of any component
commands a half-turn), so quat inputs reset to the identity quaternion:
new mj_resetCtrl sets neutral ctrl values (zero, except qw = 1), called
by mj_resetData and the viewers' Clear All. The quat chart is
restricted to dyntype 'none': integrating a quaternion setpoint linearly
is not meaningful on the manifold. New mjsActuator.ctrlspec field carries
the signature through the spec and XML round-trip.

Actuator sensors (actuatorpos/vel/frc) now report one value per force
output; dim = 3 on an SO3 actuator.

As the first actuator with nu != nactuator, this commit also makes the
viewers multi-input aware: the control sliders in simulate and studio,
which indexed per-actuator arrays by control index (out of bounds on
this model class), are generated per control and labeled with the
actuator name plus an input suffix ("orient/qw"), via the new
introspection helper mj_actuatorInputName -- the single source of truth
for input names, extended by each new multi-input type (quaternion
components are w-first: qw, qx, qy, qz). Slider ranges now honor a
defined ctrlrange even when ctrllimited is false: range is the UI hint,
limited is the clamp -- wrapped and expmap setpoints are unbounded but
still want finite sliders, while quat components are truly bounded.

The rotational demo model is orientation.xml under
test/engine/testdata/actuation/, upgraded to a three-way contrast:
per-axis wrapped servos vs an expmap-commanded vs a quat-commanded
orientation actuator, on identical checker-textured boxes. It is loaded
by the mixed-axis contrast and input-name tests, and doubles as the
viewer test model (slider groups of 3 independent, 3 grouped, 4 grouped).

PiperOrigin-RevId: 951607063
Change-Id: If235dba8e2f2ca72672e7c62531a27e967c6a373
2026-07-21 11:36:13 -07:00

1822 lines
73 KiB
C++
Raw Blame History

// Copyright 2022 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.
#include <stdalign.h>
#include <array>
#include <cstdint>
#include <cstdio>
#include <memory>
#include <string>
#include <optional>
#include <Eigen/Core>
#include <mujoco/mjxmacro.h>
#include <mujoco/mujoco.h>
#include "errors.h"
#include "function_traits.h"
#include "functions.h"
#include "private.h"
#include "raw.h"
#include "structs.h"
#include "util/func_wrap.h"
#include <pybind11/eigen.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
namespace mujoco::python {
namespace {
PYBIND11_MODULE(_functions, pymodule, pybind11::mod_gil_not_used()) {
namespace py = ::pybind11;
namespace traits = python_traits;
using EigenVectorI = Eigen::Vector<int, Eigen::Dynamic>;
using EigenVectorX = Eigen::Vector<mjtNum, Eigen::Dynamic>;
using EigenArrayXX = Eigen::Array<
mjtNum, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
// Import the _structs module so that pybind11 knows about Python bindings
// for MjWrapper types and therefore generates prettier docstrings.
py::module::import("mujoco._structs");
// Virtual file system
// Skipped entire section
// Asset cache
Def<traits::mj_getCacheSize>(pymodule, [](void* cache) {
return mj_getCacheSize(static_cast<mjCache*>(cache));
});
Def<traits::mj_getCacheCapacity>(pymodule, [](void* cache) {
return mj_getCacheCapacity(static_cast<mjCache*>(cache));
});
Def<traits::mj_setCacheCapacity>(pymodule, [](void* cache, std::size_t size) {
return mj_setCacheCapacity(static_cast<mjCache*>(cache), size);
});
Def<traits::mj_getCache>(pymodule, []() {
return static_cast<void*>(mj_getCache());
});
Def<traits::mj_clearCache>(pymodule, [](void* cache) {
return mj_clearCache(static_cast<mjCache*>(cache));
});
// Parse and compile
// Skipped: mj_loadXML (have MjModel.from_xml_string)
DEF_WITH_OMITTED_PY_ARGS(traits::mj_saveLastXML, "error", "error_sz")(
pymodule,
[](const char* filename, const mjModel* m) {
std::array<char, 1024> error;
int success = InterceptMjErrors(::mj_saveLastXML)(
filename, m, error.data(), error.size());
if (!success) {
throw FatalError(std::string(error.data()));
}
});
// Skipped: mj_freeLastXML
DEF_WITH_OMITTED_PY_ARGS(traits::mj_printSchema,
"filename", "buffer", "buffer_sz")(
pymodule, [](bool flg_html, bool flg_pad) {
constexpr int kBufferSize = 60000;
auto buffer = std::unique_ptr<char[]>(new char[kBufferSize]);
const int out_length = InterceptMjErrors(::mj_printSchema)(
nullptr, buffer.get(), kBufferSize, flg_html, flg_pad);
if (out_length >= kBufferSize) {
throw UnexpectedError("output buffer too small");
}
return std::string(buffer.get(), out_length);
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_getXMLDependencies,
"dependencies")(
pymodule, [](const char* filename){
mjStringVec dependencies;
InterceptMjErrors(::mju_getXMLDependencies)(filename, &dependencies);
return dependencies;
});
// Main simulation
pymodule.def(
"mj_step",
InterceptMjErrors(
[](const MjModelWrapper& m, MjDataWrapper& d, int nstep) {
const raw::MjModel* const m_ptr = m.get();
raw::MjData* const d_ptr = d.get();
for (int i = 0; i < nstep; ++i) {
::mj_step(m_ptr, d_ptr);
}
}),
py::arg("m"), py::arg("d"), py::arg_v("nstep", 1),
py::doc((std::string(traits::mj_step::doc) +
std::string(" Optionally, repeat nstep times.")).c_str()),
py::call_guard<py::gil_scoped_release>());
Def<traits::mj_step1>(pymodule);
Def<traits::mj_step2>(pymodule);
Def<traits::mj_forward>(pymodule);
Def<traits::mj_inverse>(pymodule);
Def<traits::mj_forwardSkip>(pymodule);
Def<traits::mj_inverseSkip>(pymodule);
// Initialization
Def<traits::mj_defaultLROpt>(pymodule);
Def<traits::mj_defaultSolRefImp>(pymodule);
Def<traits::mj_defaultOption>(pymodule);
Def<traits::mj_defaultVisual>(pymodule);
// Skipped: mj_copyModel (have MjModel.__copy__, memory managed by MjModel)
pymodule.def(
"mj_saveModel",
[](const MjModelWrapper& m,
const std::optional<std::string>& filename = std::nullopt,
std::optional<Eigen::Ref<Eigen::Vector<std::uint8_t, Eigen::Dynamic>>>
buffer = std::nullopt) {
void* buffer_ptr = nullptr;
int buffer_sz = 0;
if (buffer.has_value()) {
buffer_ptr = buffer->data();
buffer_sz = buffer->size();
}
return InterceptMjErrors(::mj_saveModel)(
m.get(), filename.has_value() ? filename->c_str() : nullptr,
buffer_ptr, buffer_sz);
},
py::arg("m"), py::arg_v("filename", std::nullopt),
py::arg_v("buffer", std::nullopt), py::doc(traits::mj_saveModel::doc),
py::call_guard<py::gil_scoped_release>());
// Skipped: mj_loadModel (have MjModel.from_binary_path)
// Skipped: mj_deleteModel (have MjModel.__del__)
Def<traits::mj_sizeModel>(pymodule);
// Skipped: mj_makeData (have MjData.__init__)
DEF_WITH_OMITTED_PY_ARGS(traits::mj_copyData)(
pymodule,
[](raw::MjData* dest, const raw::MjModel* m, const raw::MjData* src) {
InterceptMjErrors(::mj_copyData)(dest, m, src);
});
Def<traits::mj_resetData>(pymodule);
Def<traits::mj_resetCtrl>(pymodule);
Def<traits::mj_resetDataDebug>(pymodule);
Def<traits::mj_resetDataKeyframe>(pymodule);
// Skipped: mj_stackAllocByte (doesn't make sense in Python)
// Skipped: mj_deleteData (have MjData.__del__)
Def<traits::mj_resetCallbacks>(pymodule);
Def<traits::mj_setConst>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mj_setLengthRange, "error", "error_sz")(
pymodule,
[](raw::MjModel* m, raw::MjData* d, int index, const raw::MjLROpt* opt) {
std::array<char, 1024> error;
int success = InterceptMjErrors(::mj_setLengthRange)(
m, d, index, opt, error.data(), error.size());
if (!success) {
throw FatalError(std::string(error.data()));
}
});
// Printing
Def<traits::mj_printFormattedModel>(pymodule);
Def<traits::mj_printModel>(pymodule);
Def<traits::mj_printFormattedData>(pymodule);
Def<traits::mj_printData>(pymodule);
Def<traits::mj_printFormattedScene>(pymodule);
Def<traits::mj_printScene>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mju_printMat, "nr", "nc")(
pymodule,
[](Eigen::Ref<const EigenArrayXX> mat) {
return ::mju_printMat(mat.data(), mat.rows(), mat.cols());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_printMatSparse, "nr")(
pymodule,
[](Eigen::Ref<const EigenVectorX> mat,
Eigen::Ref<const Eigen::Vector<int, Eigen::Dynamic>> rownnz,
Eigen::Ref<const Eigen::Vector<int, Eigen::Dynamic>> rowadr,
Eigen::Ref<const Eigen::Vector<int, Eigen::Dynamic>> colind) {
if (rownnz.size() != rowadr.size()) {
throw py::type_error("size of rownnz should equal size of rowadr");
}
const int nnz = rowadr[rowadr.size() - 1] + rownnz[rownnz.size() - 1];
if (mat.size() != nnz) {
throw py::type_error(
"size of mat should equal rownnz[-1] + rowadr[-1]");
}
if (colind.size() != nnz) {
throw py::type_error(
"size of colind should equal rownnz[-1] + rowadr[-1]");
}
return InterceptMjErrors(::mju_printMatSparse)(
mat.data(), rowadr.size(), rownnz.data(),
rowadr.data(), colind.data());
});
// Components
Def<traits::mj_fwdKinematics>(pymodule);
Def<traits::mj_fwdPosition>(pymodule);
Def<traits::mj_fwdVelocity>(pymodule);
Def<traits::mj_fwdActuation>(pymodule);
Def<traits::mj_fwdAcceleration>(pymodule);
Def<traits::mj_fwdConstraint>(pymodule);
Def<traits::mj_Euler>(pymodule);
Def<traits::mj_RungeKutta>(pymodule);
Def<traits::mj_implicit>(pymodule);
Def<traits::mj_invPosition>(pymodule);
Def<traits::mj_invVelocity>(pymodule);
Def<traits::mj_invConstraint>(pymodule);
Def<traits::mj_compareFwdInv>(pymodule);
// Sub components
Def<traits::mj_sensorPos>(pymodule);
Def<traits::mj_sensorVel>(pymodule);
Def<traits::mj_sensorAcc>(pymodule);
Def<traits::mj_energyPos>(pymodule);
Def<traits::mj_energyVel>(pymodule);
Def<traits::mj_checkPos>(pymodule);
Def<traits::mj_checkVel>(pymodule);
Def<traits::mj_checkAcc>(pymodule);
Def<traits::mj_kinematics>(pymodule);
Def<traits::mj_comPos>(pymodule);
Def<traits::mj_camlight>(pymodule);
Def<traits::mj_flex>(pymodule);
Def<traits::mj_tendon>(pymodule);
Def<traits::mj_transmission>(pymodule);
Def<traits::mj_crb>(pymodule);
Def<traits::mj_makeM>(pymodule);
Def<traits::mj_factorM>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM, "n")(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, Eigen::Ref<EigenArrayXX> x,
Eigen::Ref<const EigenArrayXX> y) {
if (x.rows() != y.rows()) {
throw py::type_error(
"the first dimension of x and y should be of the same size");
}
if (x.cols() != m->nv) {
throw py::type_error(
"the last dimension of x should be of size nv");
}
if (y.cols() != m->nv) {
throw py::type_error(
"the last dimension of y should be of size nv");
}
return InterceptMjErrors(::mj_solveM)(
m, d, x.data(), y.data(), y.rows());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM2, "n")(
pymodule, [](const raw::MjModel* m, raw::MjData* d,
Eigen::Ref<EigenArrayXX> x, Eigen::Ref<const EigenArrayXX> y,
Eigen::Ref<const EigenArrayXX> sqrtInvD) {
if (x.rows() != y.rows()) {
throw py::type_error(
"the first dimension of x and y should be of the same size");
}
if (x.cols() != m->nv) {
throw py::type_error("the last dimension of x should be of size nv");
}
if (y.cols() != m->nv) {
throw py::type_error("the last dimension of y should be of size nv");
}
if (sqrtInvD.size() != m->nv) {
throw py::type_error("the size of sqrtInvD should be nv");
}
return InterceptMjErrors(::mj_solveM2)(m, d, x.data(), y.data(),
sqrtInvD.data(), y.rows());
});
Def<traits::mj_comVel>(pymodule);
Def<traits::mj_passive>(pymodule);
Def<traits::mj_subtreeVel>(pymodule);
Def<traits::mj_rne>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, int flg_acc,
Eigen::Ref<EigenVectorX> result) {
if (result.size() != m->nv) {
throw py::type_error("result should have length nv");
}
return InterceptMjErrors(::mj_rne)(
m, d, flg_acc, result.data());
});
Def<traits::mj_rnePostConstraint>(pymodule);
Def<traits::mj_maxContact>(pymodule);
Def<traits::mj_collision>(pymodule);
Def<traits::mj_makeConstraint>(pymodule);
Def<traits::mj_island>(pymodule);
Def<traits::mj_projectConstraint>(pymodule);
Def<traits::mj_referenceConstraint>(pymodule);
Def<traits::mj_constraintUpdate>(
pymodule, [](const raw::MjModel* m, raw::MjData* d,
Eigen::Ref<const EigenVectorX> jar,
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 1>>> cost,
int flg_coneHessian) {
if (jar.size() != d->nefc) {
throw py::type_error("size of jar should equal nefc");
}
return InterceptMjErrors(::mj_constraintUpdate)(
m, d, jar.data(), cost.has_value() ? cost->data() : nullptr,
flg_coneHessian);
});
// Support
Def<traits::mj_stateSize>(pymodule);
Def<traits::mj_getState>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d,
Eigen::Ref<EigenVectorX> state, unsigned int sig) {
if (state.size() != mj_stateSize(m, sig)) {
throw py::type_error("state size should equal mj_stateSize(m, sig)");
}
return InterceptMjErrors(::mj_getState)(m, d, state.data(), sig);
});
Def<traits::mj_extractState>(
pymodule,
[](const raw::MjModel* m,
Eigen::Ref<const EigenVectorX> src, unsigned int srcsig,
Eigen::Ref<EigenVectorX> dst, unsigned int dstsig) {
if (src.size() != mj_stateSize(m, srcsig)) {
throw py::type_error("src size should equal mj_stateSize(m, srcsig)");
}
if (dst.size() != mj_stateSize(m, dstsig)) {
throw py::type_error("dst size should equal mj_stateSize(m, dstsig)");
}
return InterceptMjErrors(::mj_extractState)(m, src.data(), srcsig,
dst.data(), dstsig);
});
Def<traits::mj_setState>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
Eigen::Ref<const EigenVectorX> state, unsigned int sig) {
if (state.size() != mj_stateSize(m, sig)) {
throw py::type_error("state size should equal mj_stateSize(m, sig)");
}
return InterceptMjErrors(::mj_setState)(m, d, state.data(), sig);
});
Def<traits::mj_copyState>(pymodule);
Def<traits::mj_readCtrl>(pymodule);
Def<traits::mj_readSensor>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d, int id, mjtNum time,
Eigen::Ref<EigenVectorX> result, int order) {
int dim = m->sensor_dim[id];
if (result.size() != dim) {
throw py::type_error("result should have length sensor_dim[id]");
}
const mjtNum* ptr = InterceptMjErrors(::mj_readSensor)(
m, d, id, time, result.data(), order);
if (ptr && ptr != result.data()) {
for (int i = 0; i < dim; ++i) {
result[i] = ptr[i];
}
}
return result;
});
Def<traits::mj_initCtrlHistory>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, int id,
std::optional<Eigen::Ref<const EigenVectorX>> times,
Eigen::Ref<const EigenVectorX> values) {
int nhistory = m->actuator_history[2*id];
if (times.has_value() && times->size() != nhistory) {
throw py::type_error(
"times should have length actuator_history[2*id]");
}
if (values.size() != nhistory) {
throw py::type_error(
"values should have length actuator_history[2*id]");
}
return InterceptMjErrors(::mj_initCtrlHistory)(
m, d, id,
times.has_value() ? times->data() : nullptr, values.data());
});
Def<traits::mj_initSensorHistory>(
pymodule, [](const raw::MjModel* m, raw::MjData* d, int id,
std::optional<Eigen::Ref<const EigenVectorX>> times,
Eigen::Ref<const EigenArrayXX> values, mjtNum phase) {
int nhistory = m->sensor_history[2 * id];
int dim = m->sensor_dim[id];
if (times.has_value() && times->size() != nhistory) {
throw py::type_error("times should have length sensor_history[2*id]");
}
if (values.rows() != nhistory || values.cols() != dim) {
throw py::type_error(
"values should have shape (sensor_history[2*id], "
"sensor_dim[id])");
}
return InterceptMjErrors(::mj_initSensorHistory)(
m, d, id, times.has_value() ? times->data() : nullptr,
values.data(), phase);
});
Def<traits::mj_setKeyframe>(pymodule);
Def<traits::mj_addContact>(pymodule);
Def<traits::mj_isPyramidal>(pymodule);
Def<traits::mj_isSparse>(pymodule);
Def<traits::mj_isDual>(pymodule);
Def<traits::mj_mulJacVec>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != d->nefc) {
throw py::type_error("res should be of length nefc");
}
if (vec.size() != m->nv) {
throw py::type_error("vec should be of length nv");
}
return InterceptMjErrors(::mj_mulJacVec)(m, d, res.data(), vec.data());
});
Def<traits::mj_mulJacTVec>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != m->nv) {
throw py::type_error("res should be of length nv");
}
if (vec.size() != d->nefc) {
throw py::type_error("vec should be of length nefc");
}
return InterceptMjErrors(::mj_mulJacTVec)(m, d, res.data(), vec.data());
});
Def<traits::mj_jac>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp,
std::optional<Eigen::Ref<EigenArrayXX>> jacr,
const mjtNum (*point)[3], int body) {
if (jacp.has_value() &&
(jacp->rows() != 3 || jacp->cols() != m->nv)) {
throw py::type_error("jacp should be of shape (3, nv)");
}
if (jacr.has_value() &&
(jacr->rows() != 3 || jacr->cols() != m->nv)) {
throw py::type_error("jacr should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_jac)(
m, d,
jacp.has_value() ? jacp->data() : nullptr,
jacr.has_value() ? jacr->data() : nullptr,
&(*point)[0], body);
});
Def<traits::mj_jacBody>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp,
std::optional<Eigen::Ref<EigenArrayXX>> jacr, int body) {
if (jacp.has_value() &&
(jacp->rows() != 3 || jacp->cols() != m->nv)) {
throw py::type_error("jacp should be of shape (3, nv)");
}
if (jacr.has_value() &&
(jacr->rows() != 3 || jacr->cols() != m->nv)) {
throw py::type_error("jacr should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_jacBody)(
m, d, jacp.has_value() ? jacp->data() : nullptr,
jacr.has_value() ? jacr->data() : nullptr, body);
});
Def<traits::mj_jacBodyCom>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp,
std::optional<Eigen::Ref<EigenArrayXX>> jacr, int body) {
if (jacp.has_value() &&
(jacp->rows() != 3 || jacp->cols() != m->nv)) {
throw py::type_error("jacp should be of shape (3, nv)");
}
if (jacr.has_value() &&
(jacr->rows() != 3 || jacr->cols() != m->nv)) {
throw py::type_error("jacr should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_jacBodyCom)(
m, d, jacp.has_value() ? jacp->data() : nullptr,
jacr.has_value() ? jacr->data() : nullptr, body);
});
Def<traits::mj_jacSubtreeCom>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp, int body) {
if (jacp.has_value() &&
(jacp->rows() != 3 || jacp->cols() != m->nv)) {
throw py::type_error("jacp should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_jacSubtreeCom)(
m, d, jacp.has_value() ? jacp->data() : nullptr, body);
});
Def<traits::mj_jacGeom>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp,
std::optional<Eigen::Ref<EigenArrayXX>> jacr, int geom) {
if (jacp.has_value() &&
(jacp->rows() != 3 || jacp->cols() != m->nv)) {
throw py::type_error("jacp should be of shape (3, nv)");
}
if (jacr.has_value() &&
(jacr->rows() != 3 || jacr->cols() != m->nv)) {
throw py::type_error("jacr should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_jacGeom)(
m, d, jacp.has_value() ? jacp->data() : nullptr,
jacr.has_value() ? jacr->data() : nullptr, geom);
});
Def<traits::mj_jacSite>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp,
std::optional<Eigen::Ref<EigenArrayXX>> jacr, int site) {
if (jacp.has_value() &&
(jacp->rows() != 3 || jacp->cols() != m->nv)) {
throw py::type_error("jacp should be of shape (3, nv)");
}
if (jacr.has_value() &&
(jacr->rows() != 3 || jacr->cols() != m->nv)) {
throw py::type_error("jacr should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_jacSite)(
m, d, jacp.has_value() ? jacp->data() : nullptr,
jacr.has_value() ? jacr->data() : nullptr, site);
});
Def<traits::mj_jacPointAxis>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp,
std::optional<Eigen::Ref<EigenArrayXX>> jacr,
const mjtNum (*point)[3], const mjtNum (*axis)[3], int body) {
if (jacp.has_value() &&
(jacp->rows() != 3 || jacp->cols() != m->nv)) {
throw py::type_error("jacp should be of shape (3, nv)");
}
if (jacr.has_value() &&
(jacr->rows() != 3 || jacr->cols() != m->nv)) {
throw py::type_error("jacr should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_jacPointAxis)(
m, d, jacp.has_value() ? jacp->data() : nullptr,
jacr.has_value() ? jacr->data() : nullptr,
&(*point)[0], &(*axis)[0], body);
});
Def<traits::mj_jacDot>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp,
std::optional<Eigen::Ref<EigenArrayXX>> jacr,
const mjtNum (*point)[3], int body) {
if (jacp.has_value() &&
(jacp->rows() != 3 || jacp->cols() != m->nv)) {
throw py::type_error("jacp should be of shape (3, nv)");
}
if (jacr.has_value() &&
(jacr->rows() != 3 || jacr->cols() != m->nv)) {
throw py::type_error("jacr should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_jacDot)(
m, d,
jacp.has_value() ? jacp->data() : nullptr,
jacr.has_value() ? jacr->data() : nullptr,
&(*point)[0], body);
});
Def<traits::mj_angmomMat>(
pymodule, [](const raw::MjModel* m, raw::MjData* d,
Eigen::Ref<EigenArrayXX> mat, int body) {
if (mat.rows() != 3 || mat.cols() != m->nv) {
throw py::type_error("mat should be of shape (3, nv)");
}
return InterceptMjErrors(::mj_angmomMat)(m, d, mat.data(), body);
});
Def<traits::mj_name2id>(pymodule);
Def<traits::mj_id2name>(pymodule);
Def<traits::mj_actuatorInputName>(pymodule);
Def<traits::mj_fullM>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d,
Eigen::Ref<EigenArrayXX> dst) {
if (dst.cols() != m->nv || dst.rows() != m->nv) {
throw py::type_error("dst should be of shape (nv, nv)");
}
return ::mj_fullM(m, d, dst.data());
});
Def<traits::mj_mulM>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d,
Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != m->nv) {
throw py::type_error("res should be of size nv");
}
if (vec.size() != m->nv) {
throw py::type_error("vec should be of size nv");
}
return InterceptMjErrors(::mj_mulM)(m, d, res.data(), vec.data());
});
Def<traits::mj_mulM2>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d,
Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != m->nv) {
throw py::type_error("res should be of size nv");
}
if (vec.size() != m->nv) {
throw py::type_error("vec should be of size nv");
}
return InterceptMjErrors(::mj_mulM2)(m, d, res.data(), vec.data());
});
Def<traits::mj_addM>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, Eigen::Ref<EigenVectorX> dst,
Eigen::Ref<Eigen::Vector<int, Eigen::Dynamic>> rownnz,
Eigen::Ref<Eigen::Vector<int, Eigen::Dynamic>> rowadr,
Eigen::Ref<Eigen::Vector<int, Eigen::Dynamic>> colind) {
if (dst.size() != m->nM) {
throw py::type_error("dst should be of size nM");
}
if (rownnz.size() != m->nv) {
throw py::type_error("rownnz should be of size nv");
}
if (rowadr.size() != m->nv) {
throw py::type_error("rowadr should be of size nv");
}
if (colind.size() != m->nM) {
throw py::type_error("colind should be of size nM");
}
return InterceptMjErrors(::mj_addM)(
m, d, dst.data(), rownnz.data(), rowadr.data(), colind.data());
});
Def<traits::mj_applyFT>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
const mjtNum (*force)[3], const mjtNum (*torque)[3],
const mjtNum (*point)[3], int body,
Eigen::Ref<EigenVectorX> qfrc_target) {
if (qfrc_target.size() != m->nv) {
throw py::type_error("qfrc_target should be of size nv");
}
return InterceptMjErrors(::mj_applyFT)(
m, d, &(*force)[0], &(*torque)[0], &(*point)[0],
body, qfrc_target.data());
});
Def<traits::mj_objectVelocity>(pymodule);
Def<traits::mj_objectAcceleration>(pymodule);
Def<traits::mj_contactForce>(pymodule);
Def<traits::mj_geomDistance>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
int geom1, int geom2, mjtNum distmax,
std::optional<Eigen::Ref<EigenArrayXX>> fromto) {
if (fromto.has_value() && fromto->size() != 6) {
throw py::type_error("fromto should be of size 6");
}
return InterceptMjErrors(::mj_geomDistance)(
m, d, geom1, geom2, distmax,
fromto.has_value() ? fromto->data() : nullptr);
});
Def<traits::mj_differentiatePos>(
pymodule,
[](const raw::MjModel* m, Eigen::Ref<EigenVectorX> qvel,
mjtNum dt, Eigen::Ref<const EigenVectorX> qpos1,
Eigen::Ref<const EigenVectorX> qpos2) {
if (qvel.size() != m->nv) {
throw py::type_error("qvel should be of size nv");
}
if (qpos1.size() != m->nq) {
throw py::type_error("qpos1 should be of size nq");
}
if (qpos2.size() != m->nq) {
throw py::type_error("qpos2 should be of size nq");
}
return InterceptMjErrors(::mj_differentiatePos)(
m, qvel.data(), dt, qpos1.data(), qpos2.data());
});
Def<traits::mj_integratePos>(
pymodule,
[](const raw::MjModel* m, Eigen::Ref<EigenVectorX> qpos,
Eigen::Ref<const EigenVectorX> qvel, mjtNum dt) {
if (qpos.size() != m->nq) {
throw py::type_error("qpos should be of size nq");
}
if (qvel.size() != m->nv) {
throw py::type_error("qvel should be of size nv");
}
return InterceptMjErrors(::mj_integratePos)(
m, qpos.data(), qvel.data(), dt);
});
Def<traits::mj_normalizeQuat>(
pymodule,
[](const raw::MjModel* m, Eigen::Ref<EigenVectorX> qpos) {
if (qpos.size() != m->nq) {
throw py::type_error("qpos should be of size nq");
}
return InterceptMjErrors(::mj_normalizeQuat)(m, qpos.data());
});
Def<traits::mj_local2Global>(pymodule);
Def<traits::mj_getTotalmass>(pymodule);
Def<traits::mj_setTotalmass>(pymodule);
Def<traits::mj_loadPluginLibrary>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mj_loadAllPluginLibraries, "callback")(
pymodule,
[](const std::string& directory) {
InterceptMjErrors(::mj_loadAllPluginLibraries)(
directory.c_str(), nullptr);
});
Def<traits::mj_version>(pymodule);
Def<traits::mj_versionString>(pymodule);
// Thread pool
Def<traits::mju_threadpool>(pymodule);
// Ray collision
Def<traits::mj_multiRay>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, const mjtNum(*pnt)[3],
Eigen::Ref<const EigenVectorX> vec,
std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
geomgroup,
mjtBool flg_static, int bodyexclude, Eigen::Ref<EigenVectorI> geomid,
Eigen::Ref<EigenVectorX> dist,
std::optional<Eigen::Ref<EigenVectorX>> normal,
int nray, mjtNum cutoff) {
if (dist.size() != nray || geomid.size() != nray) {
throw py::type_error("dist and geomid should be of size nray");
}
if (vec.size() != 3 * nray) {
throw py::type_error("vec should be of size 3*nray");
}
if (normal.has_value() && normal->size() != 3 * nray) {
throw py::type_error("normal should be of size 3*nray");
}
InterceptMjErrors(::mj_multiRay)(
m, d, &(*pnt)[0], vec.data(),
geomgroup.has_value() ? geomgroup->data() : nullptr, flg_static,
bodyexclude, geomid.data(), dist.data(),
normal.has_value() ? normal->data() : nullptr, nray, cutoff);
});
pymodule.def(
"mj_ray",
util::UnwrapArgs(
[](const raw::MjModel* m, const raw::MjData* d, const mjtNum(*pnt)[3],
const mjtNum(*vec)[3],
std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
geomgroup,
mjtBool flg_static, int bodyexclude,
std::optional<Eigen::Ref<Eigen::Vector<int, 1>>> geomid,
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mj_ray(m, d, &(*pnt)[0], &(*vec)[0],
geomgroup.has_value() ? geomgroup->data() : nullptr,
flg_static, bodyexclude,
geomid.has_value() ? geomid->data() : nullptr,
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("m"), py::arg("d"), py::arg("pnt"), py::arg("vec"),
py::arg("geomgroup"), py::arg("flg_static"), py::arg("bodyexclude"),
py::arg("geomid"), py::arg("normal") = std::nullopt,
py::doc(traits::mj_ray::doc),
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"mj_rayHfield",
util::UnwrapArgs(
[](const raw::MjModel* m, const raw::MjData* d, int geomid,
const mjtNum(*pnt)[3], const mjtNum(*vec)[3],
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mj_rayHfield(m, d, geomid, &(*pnt)[0], &(*vec)[0],
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("m"), py::arg("d"), py::arg("geomid"), py::arg("pnt"),
py::arg("vec"), py::arg("normal") = std::nullopt,
py::doc(traits::mj_rayHfield::doc),
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"mj_rayMesh",
util::UnwrapArgs(
[](const raw::MjModel* m, const raw::MjData* d, int geomid,
const mjtNum(*pnt)[3], const mjtNum(*vec)[3],
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mj_rayMesh(m, d, geomid, &(*pnt)[0], &(*vec)[0],
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("m"), py::arg("d"), py::arg("geomid"), py::arg("pnt"),
py::arg("vec"), py::arg("normal") = std::nullopt,
py::doc(traits::mj_rayMesh::doc),
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"mju_rayGeom",
util::UnwrapArgs(
[](const mjtNum(*pos)[3], const mjtNum(*mat)[9],
const mjtNum(*size)[3], const mjtNum(*pnt)[3],
const mjtNum(*vec)[3], int geomtype,
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mju_rayGeom(&(*pos)[0], &(*mat)[0], &(*size)[0], &(*pnt)[0],
&(*vec)[0], geomtype,
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("pos"), py::arg("mat"), py::arg("size"), py::arg("pnt"),
py::arg("vec"), py::arg("geomtype"), py::arg("normal") = std::nullopt,
py::doc(traits::mju_rayGeom::doc),
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"mj_rayFlex",
util::UnwrapArgs(
[](const raw::MjModel* m, const raw::MjData* d, int flex_layer,
mjtBool flg_vert, mjtBool flg_edge, mjtBool flg_face,
mjtBool flg_skin, int flexid, const mjtNum(*pnt)[3],
const mjtNum(*vec)[3],
std::optional<Eigen::Ref<Eigen::Vector<int, 1>>> vertid,
std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
return mj_rayFlex(m, d, flex_layer, flg_vert, flg_edge, flg_face,
flg_skin, flexid, &(*pnt)[0], &(*vec)[0],
vertid.has_value() ? vertid->data() : nullptr,
normal.has_value() ? normal->data() : nullptr);
}),
py::arg("m"), py::arg("d"), py::arg("flex_layer"), py::arg("flg_vert"),
py::arg("flg_edge"), py::arg("flg_face"), py::arg("flg_skin"),
py::arg("flexid"), py::arg("pnt"), py::arg("vec"),
py::arg("vertid") = std::nullopt, py::arg("normal") = std::nullopt,
py::doc(traits::mj_rayFlex::doc),
py::call_guard<py::gil_scoped_release>());
Def<traits::mju_raySkin>(pymodule);
// Interaction
Def<traits::mjv_defaultCamera>(pymodule);
Def<traits::mjv_defaultFreeCamera>(pymodule);
Def<traits::mjv_defaultPerturb>(pymodule);
Def<traits::mjv_room2model>(pymodule);
Def<traits::mjv_model2room>(pymodule);
Def<traits::mjv_cameraInModel>(pymodule);
Def<traits::mjv_cameraInRoom>(pymodule);
Def<traits::mjv_frustumHeight>(pymodule);
Def<traits::mjv_alignToCamera>(pymodule);
Def<traits::mjv_moveCamera>(pymodule);
Def<traits::mjv_movePerturb>(pymodule);
Def<traits::mjv_moveModel>(pymodule);
Def<traits::mjv_initPerturb>(pymodule);
Def<traits::mjv_applyPerturbPose>(pymodule);
Def<traits::mjv_applyPerturbForce>(pymodule);
// Skipped: mjv_averageCamera (defined in structs.cc due to the return type)
Def<traits::mjv_select>(pymodule);
// Visualization
Def<traits::mjv_defaultOption>(pymodule);
Def<traits::mjv_defaultFigure>(pymodule);
Def<traits::mjv_initGeom>(pymodule);
Def<traits::mjv_connector>(pymodule);
// Skipped: mjv_defaultScene (have MjvScene.__init__, memory managed by
// MjvScene).
// Skipped: mjv_makeScene (have MjvScene.__init__)
// Skipped: mjv_freeScene (have MjvScene.__del__)
Def<traits::mjv_updateScene>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, const raw::MjvOption* opt,
const std::optional<raw::MjvPerturb*> pert, raw::MjvCamera* cam,
int catmask, raw::MjvScene* scn) {
const raw::MjvPerturb* pert_ptr = pert.has_value() ? *pert : nullptr;
return mjv_updateScene(m, d, opt, pert_ptr, cam, catmask, scn);
});
Def<traits::mjv_addGeoms>(pymodule);
Def<traits::mjv_makeLights>(pymodule);
Def<traits::mjv_updateCamera>(pymodule);
Def<traits::mjv_updateSkin>(pymodule);
Def<traits::mjv_cameraFrame>(pymodule);
Def<traits::mjv_cameraFrustum>(pymodule);
// UI framework
// Skipped: entire section (can add this if there's demand)
// Error and memory
// Skipped: everything other than the function below (Python has exceptions)
Def<traits::mju_writeLog>(pymodule);
// Standard math
// This section consists only of preprocessor macros.
// Vector math
Def<traits::mju_zero3>(pymodule);
Def<traits::mju_copy3>(pymodule);
Def<traits::mju_scl3>(pymodule);
Def<traits::mju_add3>(pymodule);
Def<traits::mju_sub3>(pymodule);
Def<traits::mju_addTo3>(pymodule);
Def<traits::mju_subFrom3>(pymodule);
Def<traits::mju_addToScl3>(pymodule);
Def<traits::mju_addScl3>(pymodule);
Def<traits::mju_normalize3>(pymodule);
Def<traits::mju_norm3>(pymodule);
Def<traits::mju_dot3>(pymodule);
Def<traits::mju_dist3>(pymodule);
Def<traits::mju_mulMatVec3>(pymodule);
Def<traits::mju_mulMatTVec3>(pymodule);
Def<traits::mju_cross>(pymodule);
Def<traits::mju_zero4>(pymodule);
Def<traits::mju_unit4>(pymodule);
Def<traits::mju_copy4>(pymodule);
Def<traits::mju_normalize4>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mju_zero, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res) {
return InterceptMjErrors(::mju_zero)(res.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_fill, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res, mjtNum val) {
return InterceptMjErrors(::mju_fill)(res.data(), val, res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_copy, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> data) {
if (res.size() != data.size()) {
throw py::type_error("res and data should have the same size");
}
return InterceptMjErrors(::mju_copy)(
res.data(), data.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_sum, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> vec) {
return ::mju_sum(vec.data(), vec.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_L1, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> vec) {
return InterceptMjErrors(::mju_L1)(vec.data(), vec.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_scl, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec, mjtNum scl) {
if (res.size() != vec.size()) {
throw py::type_error("res and vec should have the same size");
}
return InterceptMjErrors(::mju_scl)(
res.data(), vec.data(), scl, res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_add, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec1,
Eigen::Ref<const EigenVectorX> vec2) {
if (res.size() != vec1.size()) {
throw py::type_error("res and vec1 should have the same size");
}
if (res.size() != vec2.size()) {
throw py::type_error("res and vec2 should have the same size");
}
return ::mju_add(res.data(), vec1.data(), vec2.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_sub, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec1,
Eigen::Ref<const EigenVectorX> vec2) {
if (res.size() != vec1.size()) {
throw py::type_error("res and vec1 should have the same size");
}
if (res.size() != vec2.size()) {
throw py::type_error("res and vec2 should have the same size");
}
return InterceptMjErrors(::mju_sub)(
res.data(), vec1.data(), vec2.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_addTo, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != vec.size()) {
throw py::type_error("res and vec should have the same size");
}
return InterceptMjErrors(::mju_addTo)(
res.data(), vec.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_subFrom, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != vec.size()) {
throw py::type_error("res and vec should have the same size");
}
return InterceptMjErrors(::mju_subFrom)(
res.data(), vec.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_addToScl, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec, mjtNum scl) {
if (res.size() != vec.size()) {
throw py::type_error("res and vec should have the same size");
}
return ::mju_addToScl(res.data(), vec.data(), scl, res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_addScl, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenVectorX> vec1,
Eigen::Ref<const EigenVectorX> vec2, mjtNum scl) {
if (res.size() != vec1.size()) {
throw py::type_error("res and vec1 should have the same size");
}
if (res.size() != vec2.size()) {
throw py::type_error("res and vec2 should have the same size");
}
return InterceptMjErrors(::mju_addScl)(
res.data(), vec1.data(), vec2.data(), scl, res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_normalize, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> vec) {
return InterceptMjErrors(::mju_normalize)(vec.data(), vec.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_norm, "n")(
pymodule,
[](Eigen::Ref<const EigenVectorX> vec) {
return InterceptMjErrors(::mju_norm)(vec.data(), vec.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_dot, "n")(
pymodule,
[](Eigen::Ref<const EigenVectorX> vec1,
Eigen::Ref<const EigenVectorX> vec2) {
if (vec1.size() != vec2.size()) {
throw py::type_error("vec1 and vec2 should have the same size");
}
return InterceptMjErrors(::mju_dot)(
vec1.data(), vec2.data(), vec1.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulMatVec, "nr", "nc")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != mat.rows()) {
throw py::type_error(
"size of res should equal the number of rows in mat");
}
if (vec.size() != mat.cols()) {
throw py::type_error(
"size of vec should equal the number of columns in mat");
}
return InterceptMjErrors(::mju_mulMatVec)(
res.data(), mat.data(), vec.data(), mat.rows(), mat.cols());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulMatTVec, "nr", "nc")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != mat.cols()) {
throw py::type_error(
"size of res should equal the number of columns in mat");
}
if (vec.size() != mat.rows()) {
throw py::type_error(
"size of vec should equal the number of rows in mat");
}
return InterceptMjErrors(::mju_mulMatTVec)(
res.data(), mat.data(), vec.data(), mat.rows(), mat.cols());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulVecMatVec, "n")(
pymodule,
[](Eigen::Ref<const EigenVectorX> vec1,
Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<const EigenVectorX> vec2) {
if (vec1.size() != vec2.size()) {
throw py::type_error(
"size of vec1 should equal the size of vec2");
}
if (vec1.size() != mat.cols()) {
throw py::type_error(
"size of vectors should equal the number of columns in mat");
}
if (vec1.size() != mat.rows()) {
throw py::type_error(
"size of vectors should equal the number of rows in mat");
}
return InterceptMjErrors(::mju_mulVecMatVec)(
vec1.data(), mat.data(), vec2.data(), vec1.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_transpose, "nr", "nc")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenArrayXX> mat) {
if (res.cols() != mat.rows()) {
throw py::type_error("#columns in res should equal #rows in mat");
}
if (res.rows() != mat.cols()) {
throw py::type_error("#rows in res should equal #columns in mat");
}
return InterceptMjErrors(::mju_transpose)(
res.data(), mat.data(), mat.rows(), mat.cols());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_symmetrize, "n")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenArrayXX> mat) {
if (mat.cols() != mat.rows()) {
throw py::type_error("mat should be square");
}
if (res.cols() != mat.cols() || res.rows() != mat.rows()) {
throw py::type_error("res and mat should have the same shape");
}
return InterceptMjErrors(::mju_symmetrize)(
res.data(), mat.data(), mat.rows());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_eye, "n")(
pymodule,
[](Eigen::Ref<EigenArrayXX> mat) {
if (mat.cols() != mat.rows()) {
throw py::type_error("mat should be square");
}
return InterceptMjErrors(::mju_eye)(mat.data(), mat.rows());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulMatMat, "r1", "c1", "c2")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenArrayXX> mat1,
Eigen::Ref<const EigenArrayXX> mat2) {
if (res.rows() != mat1.rows()) {
throw py::type_error("#rows in res should equal #rows in mat1");
}
if (res.cols() != mat2.cols()) {
throw py::type_error(
"#columns in res should equal #columns in mat2");
}
if (mat1.cols() != mat2.rows()) {
throw py::type_error("#columns in mat1 should equal #rows in mat2");
}
return InterceptMjErrors(::mju_mulMatMat)(
res.data(), mat1.data(), mat2.data(),
mat1.rows(), mat1.cols(), mat2.cols());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulMatMatT, "r1", "c1", "r2")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenArrayXX> mat1,
Eigen::Ref<const EigenArrayXX> mat2) {
if (res.rows() != mat1.rows()) {
throw py::type_error("#rows in res should equal #rows in mat1");
}
if (res.cols() != mat2.rows()) {
throw py::type_error("#columns in res should equal #rows in mat2");
}
if (mat1.cols() != mat2.cols()) {
throw py::type_error(
"#columns in mat1 should equal #columns in mat2");
}
return InterceptMjErrors(::mju_mulMatMatT)(
res.data(), mat1.data(), mat2.data(),
mat1.rows(), mat1.cols(), mat2.rows());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulMatTMat, "r1", "c1", "c2")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenArrayXX> mat1,
Eigen::Ref<const EigenArrayXX> mat2) {
if (res.rows() != mat1.cols()) {
throw py::type_error("#rows in res should equal #columns in mat1");
}
if (res.cols() != mat2.cols()) {
throw py::type_error(
"#columns in res should equal #columns in mat2");
}
if (mat1.rows() != mat2.rows()) {
throw py::type_error("#rows in mat1 should equal #rows in mat2");
}
return ::mju_mulMatTMat(res.data(), mat1.data(), mat2.data(),
mat1.rows(), mat1.cols(), mat2.cols());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_sqrMatTD, "nr", "nc")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenArrayXX> mat,
std::optional<Eigen::Ref<EigenVectorX>> diag) {
if (res.rows() != mat.cols()) {
throw py::type_error("#rows in res should equal #columns in mat");
}
if (res.cols() != mat.cols()) {
throw py::type_error("#rows in res should equal #columns in mat");
}
if (diag.has_value() && diag->size() != mat.rows()) {
throw py::type_error(
"size of diag should equal the number of rows in mat");
}
return InterceptMjErrors(::mju_sqrMatTD)(
res.data(), mat.data(),
diag.has_value() ? diag->data() : nullptr,
mat.rows(), mat.cols());
});
Def<traits::mju_transformSpatial>(pymodule);
// Sparse math
DEF_WITH_OMITTED_PY_ARGS(traits::mju_dense2sparse, "nr", "nc", "nnz")(
pymodule,
[](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<EigenVectorI> rownnz, Eigen::Ref<EigenVectorI> rowadr,
Eigen::Ref<EigenVectorI> colind) {
if (mat.rows() != rownnz.size()) {
throw py::type_error("#rows in mat should equal size of rownnz");
}
if (mat.rows() != rowadr.size()) {
throw py::type_error("#rows in mat should equal size of rowadr");
}
if (res.size() != colind.size()) {
throw py::type_error("#size of res should equal size of colind");
}
return ::mju_dense2sparse(res.data(), mat.data(), mat.rows(),
mat.cols(), rownnz.data(), rowadr.data(),
colind.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_sparse2dense, "nr", "nc")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenVectorX> mat,
Eigen::Ref<const EigenVectorI> rownnz,
Eigen::Ref<const EigenVectorI> rowadr,
Eigen::Ref<const EigenVectorI> colind) {
if (res.rows() != rownnz.size()) {
throw py::type_error("#rows in res should equal size of rownnz");
}
if (res.rows() != rowadr.size()) {
throw py::type_error("#rows in res should equal size of rowadr");
}
return ::mju_sparse2dense(res.data(), mat.data(), res.rows(),
res.cols(), rownnz.data(), rowadr.data(),
colind.data());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_sym2dense, "n")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenVectorX> mat,
Eigen::Ref<const EigenVectorI> rownnz,
Eigen::Ref<const EigenVectorI> rowadr,
Eigen::Ref<const EigenVectorI> colind) {
if (res.rows() != res.cols()) {
throw py::type_error("res should be a square matrix");
}
if (res.rows() != rownnz.size()) {
throw py::type_error("#rows in res should equal size of rownnz");
}
if (res.rows() != rowadr.size()) {
throw py::type_error("#rows in res should equal size of rowadr");
}
if (res.rows() > 0) {
int nnz = rowadr.array().tail(1)[0] + rownnz.array().tail(1)[0];
if (mat.size() < nnz) {
throw py::type_error("mat size is too small for the given sparse "
"structure");
}
if (colind.size() < nnz) {
throw py::type_error("colind size is too small for the given "
"sparse structure");
}
}
return ::mju_sym2dense(res.data(), mat.data(), res.rows(),
rownnz.data(), rowadr.data(), colind.data());
});
// Quaternions
Def<traits::mju_rotVecQuat>(pymodule);
Def<traits::mju_negQuat>(pymodule);
Def<traits::mju_mulQuat>(pymodule);
Def<traits::mju_mulQuatAxis>(pymodule);
Def<traits::mju_axisAngle2Quat>(pymodule);
Def<traits::mju_quat2Vel>(pymodule);
Def<traits::mju_subQuat>(pymodule);
Def<traits::mju_quat2Mat>(pymodule);
Def<traits::mju_mat2Quat>(pymodule);
Def<traits::mju_derivQuat>(pymodule);
Def<traits::mju_quatIntegrate>(pymodule);
Def<traits::mju_quatZ2Vec>(pymodule);
Def<traits::mju_mat2Rot>(pymodule);
Def<traits::mju_euler2Quat>(pymodule);
// Poses
Def<traits::mju_mulPose>(pymodule);
Def<traits::mju_negPose>(pymodule);
Def<traits::mju_trnVecPose>(pymodule);
// Decompositions
DEF_WITH_OMITTED_PY_ARGS(traits::mju_cholFactor, "n")(
pymodule,
[](Eigen::Ref<EigenArrayXX> mat, mjtNum mindiag) {
if (mat.rows() != mat.cols()) {
throw py::type_error("mat should be a square matrix");
}
return InterceptMjErrors(::mju_cholFactor)(
mat.data(), mat.rows(), mindiag);
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_cholSolve, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<const EigenVectorX> vec) {
if (mat.rows() != mat.cols()) {
throw py::type_error("mat should be a square matrix");
}
if (res.size() != mat.rows()) {
throw py::type_error(
"size of res should equal the number of rows in mat");
}
if (vec.size() != mat.cols()) {
throw py::type_error(
"size of vec should equal the number of rows in mat");
}
return InterceptMjErrors(::mju_cholSolve)(
res.data(), mat.data(), vec.data(), mat.rows());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_cholUpdate, "n")(
pymodule,
[](Eigen::Ref<EigenArrayXX> mat, Eigen::Ref<EigenVectorX> x,
int flg_plus) {
if (mat.rows() != mat.cols()) {
throw py::type_error("mat should be a square matrix");
}
if (x.size() != mat.rows()) {
throw py::type_error(
"size of x should equal the number of rows in mat");
}
return InterceptMjErrors(::mju_cholUpdate)(
mat.data(), x.data(), mat.rows(), flg_plus);
});
Def<traits::mju_cholFactorBand>(
pymodule, [](Eigen::Ref<EigenVectorX> mat, int ntotal, int nband,
int ndense, mjtNum diagadd, mjtNum diagmul) {
int nMat = (ntotal - ndense) * nband + ndense * ntotal;
if (mat.size() != nMat) {
throw py::type_error(
"mat must have size (ntotal-ndense)*nband + ndense*ntotal");
}
return InterceptMjErrors(::mju_cholFactorBand)(
mat.data(), ntotal, nband, ndense, diagadd, diagmul);
});
Def<traits::mju_cholSolveBand>(
pymodule,
[](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenVectorX> mat,
Eigen::Ref<const EigenVectorX> vec, int ntotal, int nband,
int ndense) {
int nMat = (ntotal - ndense) * nband + ndense * ntotal;
if (mat.size() != nMat) {
throw py::type_error(
"mat must have (ntotal-ndense)*nband + "
"ndense*ntotal elements");
}
if (res.size() != ntotal) {
throw py::type_error("size of res should equal ntotal");
}
if (vec.size() != ntotal) {
throw py::type_error("size of vec should equal ntotal");
}
return InterceptMjErrors(::mju_cholSolveBand)(
res.data(), mat.data(), vec.data(), ntotal, nband, ndense);
});
Def<traits::mju_band2Dense>(
pymodule,
[](Eigen::Ref<EigenArrayXX> res, Eigen::Ref<const EigenVectorX> mat,
int ntotal, int nband, int ndense, mjtBool flg_sym) {
int nMat = (ntotal - ndense) * nband + ndense * ntotal;
if (mat.size() != nMat) {
throw py::type_error(
"mat must have size (ntotal-ndense)*nband + ndense*ntotal");
}
if (res.rows() != ntotal) {
throw py::type_error("res should have ntotal rows");
}
if (res.cols() != ntotal) {
throw py::type_error("res should have ntotal columns");
}
return InterceptMjErrors(::mju_band2Dense)(
res.data(), mat.data(), ntotal, nband, ndense, flg_sym);
});
Def<traits::mju_dense2Band>(pymodule, [](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const EigenArrayXX> mat,
int ntotal, int nband, int ndense) {
int nRes = (ntotal - ndense) * nband + ndense * ntotal;
if (res.size() != nRes) {
throw py::type_error(
"res must have size (ntotal-ndense)*nband + ndense*ntotal");
}
if (mat.rows() != ntotal) {
throw py::type_error("mat should have ntotal rows");
}
if (mat.cols() != ntotal) {
throw py::type_error("mat should have ntotal columns");
}
return InterceptMjErrors(::mju_dense2Band)(res.data(), mat.data(), ntotal,
nband, ndense);
});
Def<traits::mju_bandMulMatVec>(
pymodule,
[](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<const EigenArrayXX> vec, int ntotal, int nband, int ndense,
int nVec, mjtBool flg_sym) {
int nMat = (ntotal - ndense) * nband + ndense * ntotal;
if (mat.size() != nMat) {
throw py::type_error(
"mat must have size (ntotal-ndense)*nband + ndense*ntotal");
}
if (res.rows() != ntotal) {
throw py::type_error("res should have ntotal rows");
}
if (res.cols() != nVec) {
throw py::type_error("res should have nVec columns");
}
if (vec.rows() != ntotal) {
throw py::type_error("vec should have ntotal rows");
}
if (vec.cols() != nVec) {
throw py::type_error("vec should have nVec columns");
}
return InterceptMjErrors(::mju_bandMulMatVec)(res.data(), mat.data(),
vec.data(), ntotal, nband,
ndense, nVec, flg_sym);
});
Def<traits::mju_bandDiag>(pymodule);
Def<traits::mju_eig3>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mju_boxQP, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<EigenArrayXX> R,
std::optional<Eigen::Ref<Eigen::Vector<int, Eigen::Dynamic>>> index,
Eigen::Ref<const EigenArrayXX> H,
Eigen::Ref<const EigenVectorX> g,
std::optional<Eigen::Ref<const EigenVectorX>> lower,
std::optional<Eigen::Ref<const EigenVectorX>> upper) {
int n = res.size();
if (R.size() != n*(n+7)) {
throw py::type_error("size of R should be n*(n+7)");
}
if (index.has_value() && (index->size() != n)) {
throw py::type_error("size of index should equal n");
}
if (H.rows() != n || H.cols() != n) {
throw py::type_error("H should be of shape (n, n)");
}
if (g.size() != n) {
throw py::type_error("size of g should equal n");
}
if (lower.has_value() && (lower->size() != n)) {
throw py::type_error("size of lower should equal n");
}
if (upper.has_value() && (upper->size() != n)) {
throw py::type_error("size of upper should equal n");
}
return InterceptMjErrors(::mju_boxQP)(
res.data(), R.data(),
index.has_value() ? index->data() : nullptr,
H.data(), g.data(), n,
lower.has_value() ? lower->data() : nullptr,
upper.has_value() ? upper->data() : nullptr);
});
// Miscellaneous
Def<traits::mju_muscleGain>(pymodule);
Def<traits::mju_muscleBias>(pymodule);
Def<traits::mju_muscleDynamics>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mju_encodePyramid, "dim")(
pymodule,
[](Eigen::Ref<EigenVectorX> pyramid, Eigen::Ref<const EigenVectorX> force,
Eigen::Ref<const EigenVectorX> mu) {
if (pyramid.size() != 2*mu.size()) {
throw py::type_error(
"size of pyramid should be twice as large as size of mu");
}
if (force.size() != mu.size() + 1) {
throw py::type_error(
"size of force should be exactly one larger than size of mu");
}
return InterceptMjErrors(::mju_encodePyramid)(
pyramid.data(), force.data(), mu.data(), mu.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_decodePyramid, "dim")(
pymodule,
[](Eigen::Ref<EigenVectorX> force, Eigen::Ref<const EigenVectorX> pyramid,
Eigen::Ref<const EigenVectorX> mu) {
if (pyramid.size() != 2*mu.size()) {
throw py::type_error(
"size of pyramid should be twice as large as size of mu");
}
if (force.size() != mu.size() + 1) {
throw py::type_error(
"size of force should be exactly one larger than size of mu");
}
return InterceptMjErrors(::mju_decodePyramid)(
force.data(), pyramid.data(), mu.data(), mu.size());
});
Def<traits::mju_springDamper>(pymodule);
Def<traits::mju_min>(pymodule);
Def<traits::mju_max>(pymodule);
Def<traits::mju_clip>(pymodule);
Def<traits::mju_sign>(pymodule);
Def<traits::mju_round>(pymodule);
Def<traits::mju_type2Str>(pymodule);
Def<traits::mju_str2Type>(pymodule);
Def<traits::mju_writeNumBytes>(pymodule);
Def<traits::mju_warningText>(pymodule);
Def<traits::mju_isBad>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mju_isZero, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> vec) {
return InterceptMjErrors(::mju_isZero)(vec.data(), vec.size());
});
Def<traits::mju_standardNormal>(
pymodule,
[](std::optional<mjtNum> num2) {
return InterceptMjErrors(::mju_standardNormal)(
num2.has_value() ? &*num2 : nullptr);
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_f2n, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const Eigen::Vector<float, Eigen::Dynamic>> vec) {
if (res.size() != vec.size()) {
throw py::type_error("res and vec should have the same size");
}
return InterceptMjErrors(::mju_f2n)(res.data(), vec.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_n2f, "n")(
pymodule,
[](Eigen::Ref<Eigen::Vector<float, Eigen::Dynamic>> res,
Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != vec.size()) {
throw py::type_error("res and vec should have the same size");
}
return InterceptMjErrors(::mju_n2f)(res.data(), vec.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_d2n, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res,
Eigen::Ref<const Eigen::Vector<double, Eigen::Dynamic>> vec) {
if (res.size() != vec.size()) {
throw py::type_error("res and vec should have the same size");
}
return InterceptMjErrors(::mju_d2n)(res.data(), vec.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_n2d, "n")(
pymodule,
[](Eigen::Ref<Eigen::Vector<double, Eigen::Dynamic>> res,
Eigen::Ref<const EigenVectorX> vec) {
if (res.size() != vec.size()) {
throw py::type_error("res and vec should have the same size");
}
return InterceptMjErrors(::mju_n2d)(res.data(), vec.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_insertionSort, "n")(
pymodule,
[](Eigen::Ref<EigenVectorX> res) {
return InterceptMjErrors(::mju_insertionSort)(res.data(), res.size());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_insertionSortInt, "n")(
pymodule,
[](Eigen::Ref<Eigen::Vector<int, Eigen::Dynamic>> res) {
return InterceptMjErrors(::mju_insertionSortInt)(
res.data(), res.size());
});
Def<traits::mju_Halton>(pymodule);
// Skipped: mju_strncpy (doesn't make sense in Python)
Def<traits::mju_sigmoid>(pymodule);
// Derivatives
Def<traits::mjd_transitionFD>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
mjtNum eps, mjtBool flg_centered,
std::optional<Eigen::Ref<EigenArrayXX>> A,
std::optional<Eigen::Ref<EigenArrayXX>> B,
std::optional<Eigen::Ref<EigenArrayXX>> C,
std::optional<Eigen::Ref<EigenArrayXX>> D) {
if (A.has_value() &&
(A->rows() != 2*m->nv+m->na || A->cols() != 2*m->nv+m->na)) {
throw py::type_error("A should be of shape (2*nv+na, 2*nv+na)");
}
if (B.has_value() &&
(B->rows() != 2*m->nv+m->na || B->cols() != m->nu)) {
throw py::type_error("B should be of shape (2*nv+na, nu)");
}
if (C.has_value() &&
(C->rows() != m->nsensordata || C->cols() != 2*m->nv+m->na)) {
throw py::type_error("C should be of shape (nsensordata, 2*nv+na)");
}
if (D.has_value() &&
(D->rows() != m->nsensordata || D->cols() != m->nu)) {
throw py::type_error("D should be of shape (nsensordata, nu)");
}
return InterceptMjErrors(::mjd_transitionFD)(
m, d, eps, flg_centered,
A.has_value() ? A->data() : nullptr,
B.has_value() ? B->data() : nullptr,
C.has_value() ? C->data() : nullptr,
D.has_value() ? D->data() : nullptr);
});
Def<traits::mjd_inverseFD>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
mjtNum eps, mjtBool flg_actuation,
std::optional<Eigen::Ref<EigenArrayXX>> DfDq,
std::optional<Eigen::Ref<EigenArrayXX>> DfDv,
std::optional<Eigen::Ref<EigenArrayXX>> DfDa,
std::optional<Eigen::Ref<EigenArrayXX>> DsDq,
std::optional<Eigen::Ref<EigenArrayXX>> DsDv,
std::optional<Eigen::Ref<EigenArrayXX>> DsDa,
std::optional<Eigen::Ref<EigenArrayXX>> DmDq) {
if (DfDq.has_value() &&
(DfDq->rows() != m->nv || DfDq->cols() != m->nv)) {
throw py::type_error("DfDq should be of shape (nv, nv)");
}
if (DfDv.has_value() &&
(DfDv->rows() != m->nv || DfDv->cols() != m->nv)) {
throw py::type_error("DfDv should be of shape (nv, nv)");
}
if (DfDa.has_value() &&
(DfDa->rows() != m->nv || DfDa->cols() != m->nv)) {
throw py::type_error("DfDa should be of shape (nv, nv)");
}
if (DsDq.has_value() &&
(DsDq->rows() != m->nv || DsDq->cols() != m->nsensordata)) {
throw py::type_error("DsDq should be of shape (nv, nsensordata)");
}
if (DsDv.has_value() &&
(DsDv->rows() != m->nv || DsDv->cols() != m->nsensordata)) {
throw py::type_error("DsDv should be of shape (nv, nsensordata)");
}
if (DsDa.has_value() &&
(DsDa->rows() != m->nv || DsDa->cols() != m->nsensordata)) {
throw py::type_error("DsDa should be of shape (nv, nsensordata)");
}
if (DmDq.has_value() &&
(DmDq->rows() != m->nv || DmDq->cols() != m->nC)) {
throw py::type_error("DmDq should be of shape (nv, nC)");
}
return InterceptMjErrors(::mjd_inverseFD)(
m, d, eps, flg_actuation,
DfDq.has_value() ? DfDq->data() : nullptr,
DfDv.has_value() ? DfDv->data() : nullptr,
DfDa.has_value() ? DfDa->data() : nullptr,
DsDq.has_value() ? DsDq->data() : nullptr,
DsDv.has_value() ? DsDv->data() : nullptr,
DsDa.has_value() ? DsDa->data() : nullptr,
DmDq.has_value() ? DmDq->data() : nullptr);
});
Def<traits::mjd_subQuat>(
pymodule,
[](Eigen::Ref<const EigenVectorX> qa, Eigen::Ref<const EigenVectorX> qb,
std::optional<Eigen::Ref<EigenArrayXX>> Da,
std::optional<Eigen::Ref<EigenArrayXX>> Db) {
if (qa.size() != 4) {
throw py::type_error("qa must have size 4");
}
if (qb.size() != 4) {
throw py::type_error("qb must have size 4");
}
if (Da.has_value() && Da->size() != 9) {
throw py::type_error("Da must have size 9");
}
if (Db.has_value() && Db->size() != 9) {
throw py::type_error("Db must have size 9");
}
return InterceptMjErrors(::mjd_subQuat)(
qa.data(), qb.data(),
Da.has_value() ? Da->data() : nullptr,
Db.has_value() ? Db->data() : nullptr);
});
Def<traits::mjd_quatIntegrate>(pymodule);
pymodule.def(
"_realloc_con_efc",
[](MjDataWrapper& d, int ncon, int nefc, int nJ) {
raw::MjData* data = d.get();
auto cleanup = [](raw::MjData* data, int nJ) {
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION(
static_cast<char*>(data->arena),
data->narena - data->pstack);
#endif
data->parena = 0;
data->ncon = 0;
data->nefc = 0;
if (nJ > -1) data->nJ = 0;
data->contact = static_cast<raw::MjContact*>(data->arena);
#define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_SOLVER
MJDATA_ARENA_POINTERS_DUAL
MJDATA_ARENA_POINTERS_ISLAND
MJDATA_ARENA_POINTERS_EFM
#undef X
};
char error_msg[128];
error_msg[0] = '\0';
const char* error_msg_fmt =
"Insufficient arena memory, currently allocated memory=\"%s\". "
"Increase using <size memory=\"X\"/>.";
cleanup(data, nJ);
data->ncon = ncon;
data->nefc = nefc;
if (nJ > -1) data->nJ = nJ;
data->contact =
static_cast<raw::MjContact*>(InterceptMjErrors(::mj_arenaAllocByte)(
data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact)));
if (!data->contact) {
cleanup(data, nJ);
std::snprintf(error_msg, sizeof(error_msg), error_msg_fmt,
mju_writeNumBytes(data->narena));
throw FatalError(error_msg);
}
#undef MJ_M
#define MJ_M(x) d.model().get()->x
#undef MJ_D
#define MJ_D(x) data->x
#define X(type, name, nr, nc) \
data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
data, sizeof(type) * (nr) * (nc), alignof(type))); \
if (!data->name) { \
cleanup(data, nJ); \
std::snprintf(error_msg, sizeof(error_msg), error_msg_fmt, \
mju_writeNumBytes(data->narena)); \
throw FatalError(error_msg); \
}
MJDATA_ARENA_POINTERS_SOLVER
if (mj_isDual(d.model().get())) {
MJDATA_ARENA_POINTERS_DUAL
}
#undef X
#undef MJ_D
#define MJ_D(x) x
#undef MJ_M
#define MJ_M(x) x
},
py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("nJ") = -1,
py::call_guard<py::gil_scoped_release>());
pymodule.def(
"_realloc_island",
[](MjDataWrapper& d, int nisland, int nidof) {
raw::MjData* data = d.get();
size_t parena_start = data->parena;
// Find island block start in arena to reclaim memory on re-allocation.
char* min_ptr = nullptr;
#define X(type, name, nr, nc) \
if (data->name && \
(!min_ptr || reinterpret_cast<char*>(data->name) < min_ptr)) { \
min_ptr = reinterpret_cast<char*>(data->name); \
}
MJDATA_ARENA_POINTERS_ISLAND
#undef X
if (min_ptr && data->arena) {
parena_start = min_ptr - static_cast<char*>(data->arena);
}
auto cleanup = [](raw::MjData* data, size_t target_parena) {
#define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_ISLAND
#undef X
data->nisland = 0;
data->nidof = 0;
data->parena = target_parena;
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION(
static_cast<char*>(data->arena) + target_parena,
data->narena - data->pstack - target_parena);
#endif
};
cleanup(data, parena_start);
char error_msg[128];
error_msg[0] = '\0';
const char* error_msg_fmt =
"Insufficient arena memory, currently allocated memory=\"%s\". "
"Increase using <size memory=\"X\"/>.";
data->nisland = nisland;
data->nidof = nidof;
#undef MJ_M
#define MJ_M(x) d.model().get()->x
#undef MJ_D
#define MJ_D(x) data->x
#define X(type, name, nr, nc) \
data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
data, sizeof(type) * (nr) * (nc), alignof(type))); \
if (!data->name) { \
cleanup(data, parena_start); \
std::snprintf(error_msg, sizeof(error_msg), error_msg_fmt, \
mju_writeNumBytes(data->narena)); \
throw FatalError(error_msg); \
}
MJDATA_ARENA_POINTERS_ISLAND
#undef X
#undef MJ_D
#define MJ_D(x) x
#undef MJ_M
#define MJ_M(x) x
},
py::arg("d"), py::arg("nisland"), py::arg("nidof"),
py::call_guard<py::gil_scoped_release>());
} // PYBIND11_MODULE NOLINT(readability/fn_size)
} // namespace
} // namespace mujoco::python
Reference in New Issue View Git Blame Copy Permalink