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

Add sensor matrices to mjd_transitionFD.

Browse Source 
- Also add missing `mj_jacSubtreeCom` to Python bindings.

PiperOrigin-RevId: 471610181
Change-Id: I31410d194527ce6a7bdbd01e176f3a41efa52d84
This commit is contained in:
Yuval Tassa
2022-09-01 13:00:02 -07:00
committed by Copybara-Service
parent 46da1285af
commit e3a82247c2
7 changed files with 157 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+1
View File
@@ -61,6 +61,7 @@ General
- Added a disable flag for sensors.
- :ref:`mju_mulQuat` and :ref:`mju_mulQuatAxis` support in place computation. For example
|br| ``mju_mulQuat(a, a, b);`` sets the quaternion ``a`` equal to the product of ``a`` and ``b``.
- Added sensor matrices to ``mjd_transitionFD`` (note this is an API change).
Deleted/deprecated features
^^^^^^^^^^^^^^^^^^^^^^^^^^^
include/mujoco/mujoco.h
+8 -3
View File
@@ -1091,12 +1091,17 @@ MJAPI mjtNum mju_sigmoid(mjtNum x);
//---------------------- Derivatives ---------------------------------------------------------------
// Finite differenced state-transition and control-transition matrices dx(t+h) = A*dx(t) + B*du(t).
// Finite differenced transition matrices (control theory notation)
// d(x_next) = A*dx + B*du
// d(sensor) = C*dx + D*du
// required output matrix dimensions:
// A: (2*nv+na x 2*nv+na)
// B: (2*nv+na x nu)
MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
mjtNum* A, mjtNum* B);
// D: (nsensordata x 2*nv+na)
// C: (nsensordata x nu)
void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
#if defined(__cplusplus)
introspect/functions.py
+13 -1
View File
@@ -6980,7 +6980,19 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='C',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='D',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
),
doc='Finite differenced state-transition and control-transition matrices dx(t+h) = A*dx(t) + B*du(t). required output matrix dimensions: A: (2*nv+na x 2*nv+na) B: (2*nv+na x nu)', # pylint: disable=line-too-long
doc='Finite differenced transition matrices (control theory notation) d(x_next) = A*dx + B*du d(sensor) = C*dx + D*du required output matrix dimensions: A: (2*nv+na x 2*nv+na) B: (2*nv+na x nu) D: (nsensordata x 2*nv+na) C: (nsensordata x nu)', # pylint: disable=line-too-long
)),
])
python/mujoco/functions.cc
+25 -2
View File
@@ -355,6 +355,17 @@ PYBIND11_MODULE(_functions, pymodule) {
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,
@@ -1033,7 +1044,9 @@ PYBIND11_MODULE(_functions, pymodule) {
pymodule,
[](const raw::MjModel* m, raw::MjData* d, mjtNum eps, mjtByte centered,
std::optional<Eigen::Ref<EigenArrayXX>> A,
std::optional<Eigen::Ref<EigenArrayXX>> B) {
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)");
@@ -1042,10 +1055,20 @@ PYBIND11_MODULE(_functions, pymodule) {
(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, centered,
A.has_value() ? A->data() : nullptr,
B.has_value() ? B->data() : nullptr);
B.has_value() ? B->data() : nullptr,
C.has_value() ? C->data() : nullptr,
D.has_value() ? D->data() : nullptr);
});
Def<traits::mju_Halton>(pymodule);
// Skipped: mju_strncpy (doesn't make sense in Python)
src/engine/engine_derivative.c
+34 -11
View File
@@ -1463,7 +1463,7 @@ void mjd_smooth_vel(const mjModel *m, mjData *d) {
// finite differenced Jacobian of (next_state, sensors) = mj_step(state, control)
// all outputs are optional
// output dimensions (transposed w.r.t common convention):
// output dimensions (transposed w.r.t Control Theory convention):
// DyDq: (nv x 2*nv+na)
// DyDv: (nv x 2*nv+na)
// DyDa: (na x 2*nv+na)
@@ -1712,30 +1712,53 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
// finite differenced state-transition and control-transition matrices dy = A*dx + B*du
// finite differenced transition matrices (control theory notation)
// d(x_next) = A*dx + B*du
// d(sensor) = C*dx + D*du
// required output matrix dimensions:
// A: (2*nv+na x 2*nv+na)
// B: (2*nv+na x nu)
// D: (nsensordata x 2*nv+na)
// C: (nsensordata x nu)
void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
mjtNum* A, mjtNum* B) {
int nv = m->nv, na = m->na, nu = m->nu;
int ndx = 2*nv+na; // row length of Jacobians
mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D) {
int nv = m->nv, na = m->na, nu = m->nu, ns = m->nsensordata;
int ndx = 2*nv+na; // row length of state Jacobians
// stepFD() offset pointers, initialised to NULL
mjtNum *DyDq, *DyDv, *DyDa, *DsDq, *DsDv, *DsDa;
DyDq = DyDv = DyDa = DsDq = DsDv = DsDa = NULL;
mjMARKSTACK;
// allocate transposed matrices
mjtNum *AT = mj_stackAlloc(d, ndx*ndx); // state-transition matrix (transposed)
mjtNum *AT = A ? mj_stackAlloc(d, ndx*ndx) : NULL; // state-transition matrix (transposed)
mjtNum *BT = B ? mj_stackAlloc(d, nu*ndx) : NULL; // control-transition matrix (transposed)
mjtNum *CT = C ? mj_stackAlloc(d, ndx*ns) : NULL; // state-observation matrix (transposed)
mjtNum *DT = D ? mj_stackAlloc(d, nu*ns) : NULL; // control-observation matrix (transposed)
// set offset pointers
if (A) {
DyDq = AT;
DyDv = AT+ndx*nv;
DyDa = AT+ndx*2*nv;
}
if (C) {
DsDq = CT;
DsDv = CT + ns*nv;
DsDa = CT + ns*2*nv;
}
// get Jacobians
if (A) {
mjd_stepFD(m, d, eps, centered, AT, AT+ndx*nv, AT+ndx*2*nv, BT, NULL, NULL, NULL, NULL);
} else {
mjd_stepFD(m, d, eps, centered, NULL, NULL, NULL, BT, NULL, NULL, NULL, NULL);
}
mjd_stepFD(m, d, eps, centered, DyDq, DyDv, DyDa, BT, DsDq, DsDv, DsDa, DT);
// transpose
if (A) mju_transpose(A, AT, ndx, ndx);
if (B) mju_transpose(B, BT, nu, ndx);
if (C) mju_transpose(C, CT, ndx, ns);
if (D) mju_transpose(D, DT, nu, ns);
mjFREESTACK;
}
src/engine/engine_derivative.h
+2 -2
View File
@@ -39,9 +39,9 @@ MJAPI void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps, mjtNum* Df
// advance simulation using control callback, skipstage is mjtStage
MJAPI void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor);
// finite differenced state-transition and control-transition matrices dy = A*dx + B*du
// finite differenced transition matrices (control theory notation)
MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
mjtNum* A, mjtNum* B);
mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
#ifdef __cplusplus
}
test/engine/engine_derivative_test.cc
+74 -12
View File
@@ -352,7 +352,8 @@ TEST_F(DerivativeTest, LinearSystem) {
mjtNum* AFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*2*nv);
mjtNum* BFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu);
mjd_transitionFD(model, data, eps, /*centered=*/0, AFD, BFD);
mjd_transitionFD(model, data, eps, /*centered=*/0,
AFD, BFD, nullptr, nullptr);
// uncomment for debugging:
// PrintMatrix(AFD, 2*nv, 2*nv);
@@ -365,13 +366,12 @@ TEST_F(DerivativeTest, LinearSystem) {
// central differenced A and B
mjtNum* AFDc = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*2*nv);
mjtNum* BFDc = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu);
mjd_transitionFD(model, data, eps, /*centered=*/1, AFDc, BFDc);
mjd_transitionFD(model, data, eps, /*centered=*/1,
AFDc, BFDc, nullptr, nullptr);
// expect central derivatives to be closer to analytic solution
EXPECT_LT(CompareMatrices(A, AFDc, 2*nv, 2*nv, eps),
CompareMatrices(A, AFD, 2*nv, 2*nv, eps));
EXPECT_LT(CompareMatrices(B, BFDc, 2*nv, nu, eps),
CompareMatrices(B, BFD, 2*nv, nu, eps));
// expect central derivatives to be equal to forward differences
CompareMatrices(AFD, AFDc, 2*nv, 2*nv, eps);
CompareMatrices(BFD, BFDc, 2*nv, nu, eps);
mju_free(BFDc);
mju_free(AFDc);
@@ -400,7 +400,7 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
// analytic B
mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu);
LinearSystem(model, data, NULL, B);
LinearSystem(model, data, nullptr, B);
// forward differenced A and B
mjtNum eps = 1e-6;
@@ -409,19 +409,22 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
// set ctrl to the limits, request forward differences
data->ctrl[0] = 1;
data->ctrl[1] = -1;
mjd_transitionFD(model, data, eps, /*centered=*/0, NULL, BFD);
mjd_transitionFD(model, data, eps, /*centered=*/0,
nullptr, BFD, nullptr, nullptr);
// expect FD and analytic derivatives to be similar to eps precision
CompareMatrices(B, BFD, 2*nv, nu, eps);
// ctrl remains at limits, request central differences
mjd_transitionFD(model, data, eps, /*centered=*/1, NULL, BFD);
mjd_transitionFD(model, data, eps, /*centered=*/1,
nullptr, BFD, nullptr, nullptr);
// expect FD and analytic derivatives to be similar to eps precision
CompareMatrices(B, BFD, 2*nv, nu, eps);
// set ctrl beyond limits, request forward differences
data->ctrl[0] = 2;
data->ctrl[1] = -2;
mjd_transitionFD(model, data, eps, /*centered=*/0, NULL, BFD);
mjd_transitionFD(model, data, eps, /*centered=*/0,
nullptr, BFD, nullptr, nullptr);
// expect derivatives to be 0
EXPECT_THAT(AsVector(BFD, 2*nv*nu), Each(Eq(0.0)));
@@ -430,7 +433,8 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
EXPECT_EQ(data->ctrl[1], -2.0);
// ctrl remains beyond limits, request centered differences
mjd_transitionFD(model, data, eps, /*centered=*/1, NULL, BFD);
mjd_transitionFD(model, data, eps, /*centered=*/1,
nullptr, BFD, nullptr, nullptr);
// expect derivatives to be 0
EXPECT_THAT(AsVector(BFD, 2*nv*nu), Each(Eq(0.0)));
@@ -440,5 +444,63 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
mj_deleteModel(model);
}
// compare FD sensor derivatives to analytic derivatives
TEST_F(DerivativeTest, SensorDerivatives) {
static constexpr char xml[] = R"(
<mujoco>
<worldbody>
<body>
<joint name="joint" type="slide"/>
<geom size=".1"/>
</body>
</worldbody>
<actuator>
<general name="actuator" joint="joint" gainprm="3"/>
</actuator>
<sensor>
<jointpos joint="joint"/>
<jointvel joint="joint"/>
<actuatorfrc actuator="actuator"/>
</sensor>
</mujoco>
)";
mjModel* model = LoadModelFromString(xml);
int nv = model->nv, nu = model->nu, ns = model->nsensordata;
mjData* data = mj_makeData(model);
// finite differenced C and D
mjtNum eps = 1e-6;
mjtNum* CFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*ns*2*nv);
mjtNum* DFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*ns*nu);
mjd_transitionFD(model, data, eps, /*centered=*/0,
nullptr, nullptr, CFD, DFD);
// expected analytic C and D
mjtNum C[6] = {
1, 0,
0, 1,
0, 0
};
mjtNum D[3] = {
0,
0,
3,
};
// compare expected and actual values
CompareMatrices(CFD, C, ns, 2*nv, eps);
CompareMatrices(DFD, D, ns, nu, eps);
mju_free(DFD);
mju_free(CFD);
mj_deleteData(data);
mj_deleteModel(model);
}
} // namespace
} // namespace mujoco