Add private functions and associated tests related to constraint island discovery.
PiperOrigin-RevId: 553156672 Change-Id: Ifd73e3df443f7955359c5243d5f7756069335c05
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
efca1a76aa
commit
a063e2f0ef
@@ -23,6 +23,7 @@
|
||||
#include <mujoco/mjmodel.h>
|
||||
#include <mujoco/mjxmacro.h>
|
||||
#include "engine/engine_array_safety.h"
|
||||
#include "engine/engine_crossplatform.h"
|
||||
#include "engine/engine_core_smooth.h"
|
||||
#include "engine/engine_io.h"
|
||||
#include "engine/engine_support.h"
|
||||
@@ -1486,7 +1487,7 @@ static inline int mj_nl(const mjModel* m, const mjData* d, int *nnz) {
|
||||
|
||||
margin = m->jnt_margin[i];
|
||||
|
||||
// slider and hinge joint limits can be bilateral, check both side
|
||||
// slider and hinge joint limits can be bilateral, check both sides
|
||||
if (m->jnt_type[i] == mjJNT_SLIDE || m->jnt_type[i] == mjJNT_HINGE) {
|
||||
value = d->qpos[m->jnt_qposadr[i]];
|
||||
for (side=-1; side <= 1; side+=2) {
|
||||
@@ -2095,3 +2096,129 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
|
||||
*cost = s;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
//---------------------------- constraint islands --------------------------------------------------
|
||||
|
||||
// comparison function for lexicographic edge sorting
|
||||
quicksortfunc(edgecompare, context, edge0, edge1) {
|
||||
int* e0 = (int*)edge0;
|
||||
int* e1 = (int*)edge1;
|
||||
int v00 = e0[0];
|
||||
int v10 = e1[0];
|
||||
|
||||
if (v00 < v10) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (v00 == v10) {
|
||||
int v01 = e0[1];
|
||||
int v11 = e1[1];
|
||||
|
||||
if (v01 < v11) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (v01 == v11) {
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
|
||||
|
||||
// construct sparse matrix from unsorted edge array, return number of nonzeros
|
||||
int mj_edge2Sparse(int* rownnz, int* rowadr, int* colind, int* edge, int ne, int nr) {
|
||||
if (!ne) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// sort edges
|
||||
mjQUICKSORT(edge, ne, 2*sizeof(int), edgecompare, NULL);
|
||||
|
||||
// construct sparse
|
||||
int nnz = 0; // number of nonzeros
|
||||
int e = 0; // current edge
|
||||
for (int r=0; r < nr; r++) {
|
||||
// init row
|
||||
rownnz[r] = 0;
|
||||
rowadr[r] = nnz;
|
||||
|
||||
// copy values while making unique and checking indices
|
||||
while (e < ne && edge[2*e] == r) {
|
||||
int v0 = edge[2*e];
|
||||
int v1 = edge[2*e + 1];
|
||||
|
||||
// skip if duplicate
|
||||
if (rownnz[r] && v0 == edge[2*e - 2] && v1 == edge[2*e - 1]) {
|
||||
e++;
|
||||
continue;
|
||||
}
|
||||
|
||||
// check for invalid indices
|
||||
if (v0 < 0 || v0 >= nr) mju_error("invalid row index %d in edge %d", v0, e);
|
||||
if (v1 < 0 || v1 >= nr) mju_error("invalid column index %d in edge %d", v1, e);
|
||||
|
||||
// copy column index, increment nnz, e, rownnz
|
||||
colind[nnz++] = edge[2*(e++) + 1];
|
||||
rownnz[r]++;
|
||||
}
|
||||
}
|
||||
|
||||
return nnz;
|
||||
}
|
||||
|
||||
|
||||
// find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix
|
||||
// arguments:
|
||||
// island (nr) - island index assigned to vertex, -1 if vertex has no edges
|
||||
// nr - number of rows/columns of adjacency matrix
|
||||
// rownnz (nr) - matrix row nonzeros
|
||||
// rowadr (nr) - matrix row addresses
|
||||
// colind (nnz) - matrix column indices
|
||||
// stack (nnz) - stack space
|
||||
// returns number of islands
|
||||
int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind,
|
||||
int* stack) {
|
||||
// initialize island count, set ids to -1
|
||||
int nisland = 0;
|
||||
for (int i=0; i < nr; i++) island[i] = -1;
|
||||
|
||||
// iterate over vertices, discover islands
|
||||
for (int i=0; i < nr; i++) {
|
||||
// vertex already in island or singleton with no edges: skip
|
||||
if (island[i] != -1 || !rownnz[i]) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// push i onto stack
|
||||
int nstack = 0;
|
||||
stack[nstack++] = i;
|
||||
|
||||
// DFS traversal of island
|
||||
while (nstack) {
|
||||
// pop v from stack
|
||||
int v = stack[--nstack];
|
||||
|
||||
// if v is already assigned, continue
|
||||
if (island[v] != -1) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// assign v to current island
|
||||
island[v] = nisland;
|
||||
|
||||
// push adjacent vertices onto stack
|
||||
memcpy(stack + nstack, colind + rowadr[v], rownnz[v]*sizeof(int));
|
||||
nstack += rownnz[v];
|
||||
}
|
||||
|
||||
// island is filled: increment nisland
|
||||
nisland++;
|
||||
}
|
||||
|
||||
return nisland;
|
||||
}
|
||||
|
||||
@@ -114,6 +114,13 @@ MJAPI void mj_referenceConstraint(const mjModel* m, mjData* d);
|
||||
// optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians
|
||||
MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
|
||||
mjtNum cost[1], int flg_coneHessian);
|
||||
|
||||
// construct sparse matrix from unsorted edge array, return number of nonzeros
|
||||
MJAPI int mj_edge2Sparse(int* rownnz, int* rowadr, int* colind, int* edge, int ne, int nr);
|
||||
|
||||
MJAPI int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind,
|
||||
int* scratch);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -35,8 +35,8 @@ mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
|
||||
const int nnz2, const int* ind2);
|
||||
|
||||
// convert matrix from dense to sparse
|
||||
void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
|
||||
int* rownnz, int* rowadr, int* colind);
|
||||
MJAPI void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
|
||||
int* rownnz, int* rowadr, int* colind);
|
||||
|
||||
// convert matrix from sparse to dense
|
||||
void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
|
||||
|
||||
@@ -24,12 +24,15 @@
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "src/engine/engine_core_constraint.h"
|
||||
#include "src/engine/engine_support.h"
|
||||
#include "src/engine/engine_util_sparse.h"
|
||||
#include "test/fixture.h"
|
||||
|
||||
namespace mujoco {
|
||||
namespace {
|
||||
|
||||
using ::testing::DoubleNear;
|
||||
using ::testing::ElementsAre;
|
||||
using ::testing::ElementsAreArray;
|
||||
using ::testing::Pointwise;
|
||||
using CoreConstraintTest = MujocoTest;
|
||||
|
||||
@@ -83,7 +86,7 @@ TEST_F(CoreConstraintTest, WeldRotJacobian) {
|
||||
ASSERT_THAT(model, testing::NotNull());
|
||||
ASSERT_EQ(model->nq, 7);
|
||||
ASSERT_EQ(model->nv, 6);
|
||||
static const int nv = 6; // for increased readabilty
|
||||
static const int nv = 6; // for increased readability
|
||||
mjData* data = mj_makeData(model);
|
||||
|
||||
// arbitrary initial values for the ball and hinge joints
|
||||
@@ -269,5 +272,233 @@ TEST_F(CoreConstraintTest, CombineSparseCount) {
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(CoreConstraintTest, EdgeToSparse4) {
|
||||
// unsorted edges, with duplication
|
||||
constexpr int ne = 6;
|
||||
constexpr int nr = 5;
|
||||
int edge[2*ne] = {
|
||||
1, 1,
|
||||
0, 0,
|
||||
0, 1,
|
||||
3, 2,
|
||||
1, 1,
|
||||
0, 0
|
||||
};
|
||||
|
||||
int rownnz[nr];
|
||||
int rowadr[nr];
|
||||
int colind[ne];
|
||||
|
||||
int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, ne, nr);
|
||||
|
||||
constexpr int expected_nnz = 4;
|
||||
|
||||
EXPECT_EQ(nnz, expected_nnz);
|
||||
|
||||
EXPECT_THAT(rownnz, ElementsAre(2, 1, 0, 1, 0));
|
||||
EXPECT_THAT(rowadr, ElementsAre(0, 2, 3, 3, 4));
|
||||
|
||||
int expected_colind[expected_nnz] = {0, 1, 1, 2};
|
||||
EXPECT_THAT(expected_colind, ElementsAreArray(colind, expected_nnz));
|
||||
}
|
||||
|
||||
TEST_F(CoreConstraintTest, EdgeToSparse2) {
|
||||
// unsorted edges, with duplication
|
||||
constexpr int ne = 4;
|
||||
constexpr int nr = 5;
|
||||
int edge[2*ne] = {
|
||||
3, 4,
|
||||
1, 1,
|
||||
3, 4,
|
||||
1, 1
|
||||
};
|
||||
|
||||
int rownnz[nr];
|
||||
int rowadr[nr];
|
||||
int colind[ne];
|
||||
|
||||
int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, ne, nr);
|
||||
|
||||
constexpr int expected_nnz = 2;
|
||||
|
||||
EXPECT_EQ(nnz, expected_nnz);
|
||||
|
||||
EXPECT_THAT(rownnz, ElementsAre(0, 1, 0, 1, 0));
|
||||
EXPECT_THAT(rowadr, ElementsAre(0, 0, 1, 1, 2));
|
||||
|
||||
int expected_colind[expected_nnz] = {1, 4};
|
||||
EXPECT_THAT(expected_colind, ElementsAreArray(colind, expected_nnz));
|
||||
}
|
||||
|
||||
TEST_F(CoreConstraintTest, EdgeToSparse3) {
|
||||
// unsorted edges, with duplication
|
||||
constexpr int ne = 3;
|
||||
constexpr int nr = 1;
|
||||
int edge[2*ne] = {
|
||||
0, 0,
|
||||
0, 0,
|
||||
0, 0
|
||||
};
|
||||
|
||||
int rownnz[nr];
|
||||
int rowadr[nr];
|
||||
int colind[ne];
|
||||
|
||||
int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, ne, nr);
|
||||
|
||||
constexpr int expected_nnz = 1;
|
||||
|
||||
EXPECT_EQ(nnz, expected_nnz);
|
||||
|
||||
EXPECT_THAT(rownnz, ElementsAre(1));
|
||||
EXPECT_THAT(rowadr, ElementsAre(0));
|
||||
|
||||
int expected_colind[expected_nnz] = {0};
|
||||
EXPECT_THAT(expected_colind, ElementsAreArray(colind, expected_nnz));
|
||||
}
|
||||
|
||||
|
||||
TEST_F(CoreConstraintTest, FloodFillSingleton) {
|
||||
// adjacency matrix for the graph 0 1 2
|
||||
// U U
|
||||
// (3 singletons, 0 and 2 have self-edges)
|
||||
mjtNum mat[9] = {
|
||||
1, 0, 0,
|
||||
0, 0, 0,
|
||||
0, 0, 1
|
||||
};
|
||||
constexpr int nr = 3;
|
||||
constexpr int nnz = 2;
|
||||
int rownnz[nr];
|
||||
int rowadr[nr];
|
||||
int colind[nnz];
|
||||
mjtNum res[nnz]; // unused
|
||||
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
|
||||
|
||||
// outputs / scratch
|
||||
int island[nr];
|
||||
int scratch[2*nr];
|
||||
|
||||
// flood fill
|
||||
int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, scratch);
|
||||
|
||||
EXPECT_EQ(nisland, 2);
|
||||
EXPECT_THAT(island, ElementsAre(0, -1, 1));
|
||||
}
|
||||
|
||||
TEST_F(CoreConstraintTest, FloodFill1) {
|
||||
// adjacency matrix for the graph 0 - 1 - 2
|
||||
mjtNum mat[9] = {
|
||||
0, 1, 0,
|
||||
1, 0, 1,
|
||||
0, 1, 0
|
||||
};
|
||||
constexpr int nr = 3;
|
||||
constexpr int nnz = 4;
|
||||
int rownnz[nr];
|
||||
int rowadr[nr];
|
||||
int colind[nnz];
|
||||
mjtNum res[nnz]; // unused
|
||||
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
|
||||
|
||||
// outputs / stack
|
||||
int island[nr];
|
||||
int stack[nnz];
|
||||
|
||||
int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, stack);
|
||||
|
||||
EXPECT_EQ(nisland, 1);
|
||||
EXPECT_THAT(island, ElementsAre(0, 0, 0));
|
||||
}
|
||||
|
||||
TEST_F(CoreConstraintTest, FloodFill2) {
|
||||
// adjacency matrix for the graph 6 – 1 – 4 0 – 3 – 5 – 2
|
||||
mjtNum mat[49] = {
|
||||
0, 0, 0, 1, 0, 0, 0,
|
||||
0, 0, 0, 0, 1, 0, 1,
|
||||
0, 0, 0, 0, 0, 1, 0,
|
||||
1, 0, 0, 0, 0, 1, 0,
|
||||
0, 1, 0, 0, 0, 0, 0,
|
||||
0, 0, 1, 1, 0, 0, 0,
|
||||
0, 1, 0, 0, 0, 0, 0,
|
||||
};
|
||||
constexpr int nr = 7;
|
||||
constexpr int nnz = 10;
|
||||
int rownnz[nr];
|
||||
int rowadr[nr];
|
||||
int colind[nnz];
|
||||
mjtNum res[nnz]; // unused
|
||||
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
|
||||
|
||||
// outputs / stack
|
||||
int island[nr];
|
||||
int stack[nnz];
|
||||
|
||||
int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, stack);
|
||||
|
||||
EXPECT_EQ(nisland, 2);
|
||||
EXPECT_THAT(island, ElementsAre(0, 1, 0, 0, 1, 0, 1));
|
||||
}
|
||||
|
||||
TEST_F(CoreConstraintTest, FloodFill3a) {
|
||||
// adjacency matrix for the graph 0 2 1 – 3
|
||||
// U
|
||||
mjtNum mat[16] = {
|
||||
0, 0, 0, 0,
|
||||
0, 0, 0, 1,
|
||||
0, 0, 1, 0,
|
||||
0, 1, 0, 0,
|
||||
};
|
||||
constexpr int nr = 4;
|
||||
constexpr int nnz = 3;
|
||||
int rownnz[nr];
|
||||
int rowadr[nr];
|
||||
int colind[nnz];
|
||||
mjtNum res[nnz]; // unused
|
||||
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
|
||||
|
||||
// outputs / stack
|
||||
int island[nr];
|
||||
int stack[nnz];
|
||||
|
||||
int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, stack);
|
||||
|
||||
EXPECT_EQ(nisland, 2);
|
||||
EXPECT_THAT(island, ElementsAre(-1, 0, 1, 0));
|
||||
}
|
||||
|
||||
TEST_F(CoreConstraintTest, FloodFill3b) {
|
||||
/*
|
||||
adjacency matrix for the graph 1 – 2 3 4 – 5
|
||||
U | \ |
|
||||
0 – 6
|
||||
*/
|
||||
mjtNum mat[49] = {
|
||||
0, 0, 0, 0, 1, 0, 1,
|
||||
0, 1, 1, 0, 0, 0, 0,
|
||||
0, 1, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0,
|
||||
1, 0, 0, 0, 0, 1, 1,
|
||||
0, 0, 0, 0, 1, 0, 1,
|
||||
1, 0, 0, 0, 1, 1, 0,
|
||||
};
|
||||
constexpr int nr = 7;
|
||||
constexpr int nnz = 13;
|
||||
int rownnz[nr];
|
||||
int rowadr[nr];
|
||||
int colind[nnz];
|
||||
mjtNum res[nnz]; // unused
|
||||
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
|
||||
|
||||
// outputs / stack
|
||||
int island[nr];
|
||||
int stack[nnz];
|
||||
|
||||
int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, stack);
|
||||
|
||||
EXPECT_EQ(nisland, 2);
|
||||
EXPECT_THAT(island, ElementsAre(0, 1, 1, -1, 0, 0, 0));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
|
||||
Reference in New Issue
Block a user