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

Use sparse (uncompressed) actuator_moment in mj_transmission.

Browse Source 
PiperOrigin-RevId: 692179704
Change-Id: Ic30ac5a98dc13de2028e378df65dc88ba3912bf5
This commit is contained in:
Taylor Howell
2024-11-01 08:05:14 -07:00
committed by Copybara-Service
parent 1d58576d28
commit a51f346059
13 changed files with 230 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+81 -47
View File
@@ -857,6 +857,9 @@ void mj_transmission(const mjModel* m, mjData* d) {
// outputs
mjtNum* length = d->actuator_length;
mjtNum* moment = d->actuator_moment;
int *rownnz = d->moment_rownnz;
int *rowadr = d->moment_rowadr;
int *colind = d->moment_colind;
// allocate Jacbians
mj_markStack(d);
@@ -875,6 +878,10 @@ void mj_transmission(const mjModel* m, mjData* d) {
// compute lengths and moments
for (int i=0; i < nu; i++) {
rownnz[i] = 0;
rowadr[i] = i == 0 ? 0 : rowadr[i-1] + rownnz[i-1];
int adr = rowadr[i];
// extract info
int id = m->actuator_trnid[2*i];
mjtNum* gear = m->actuator_gear+6*i;
@@ -885,18 +892,19 @@ void mj_transmission(const mjModel* m, mjData* d) {
case mjTRN_JOINTINPARENT: // joint, force in parent frame
// slide and hinge joint: scalar gear
if (m->jnt_type[id] == mjJNT_SLIDE || m->jnt_type[id] == mjJNT_HINGE) {
// sparsity
rownnz[i]++;
colind[adr] = m->jnt_dofadr[id];
length[i] = d->qpos[m->jnt_qposadr[id]]*gear[0];
moment[i*nv + m->jnt_dofadr[id]] = gear[0];
moment[adr] = gear[0];
}
// ball joint: 3D wrench gear
else if (m->jnt_type[id] == mjJNT_BALL) {
// j: qpos start address
int j = m->jnt_qposadr[id];
// axis: expmap representation of quaternion
mjtNum axis[3], quat[4];
mju_copy4(quat, d->qpos+j);
mju_copy4(quat, d->qpos+m->jnt_qposadr[id]);
mju_normalize4(quat);
mju_quat2Vel(axis, quat, 1);
@@ -912,11 +920,17 @@ void mj_transmission(const mjModel* m, mjData* d) {
// length: axis*gearAxis
length[i] = mju_dot3(axis, gearAxis);
// j: dof start address
j = m->jnt_dofadr[id];
// dof start address
int jnt_dofadr = m->jnt_dofadr[id];
// sparsity
for (int j = 0; j < 3; j++) {
colind[adr+j] = jnt_dofadr + j;
}
rownnz[i] += 3;
// moment: gearAxis
mju_copy3(moment+i*nv+j, gearAxis);
mju_copy3(moment+adr, gearAxis);
}
// free joint: 6D wrench gear
@@ -924,35 +938,30 @@ void mj_transmission(const mjModel* m, mjData* d) {
// cannot compute meaningful length, set to 0
length[i] = 0;
// j: qpos start address
int j = m->jnt_qposadr[id];
// vec: translational components
mjtNum vec[3];
mju_copy3(vec, d->qpos+j);
// axis: expmap representation of quaternion
mjtNum axis[3], quat[4];
mju_quat2Vel(axis, d->qpos+j+3, 1);
mju_copy4(quat, d->qpos+j+3);
mju_normalize4(quat);
mju_quat2Vel(axis, quat, 1);
// gearAxis: rotate to world frame if necessary
mjtNum gearAxis[3];
if (m->actuator_trntype[i] == mjTRN_JOINT) {
mju_copy3(gearAxis, gear+3);
} else {
mjtNum quat[4];
mju_copy4(quat, d->qpos+m->jnt_qposadr[id]+3);
mju_normalize4(quat);
mju_negQuat(quat, quat);
mju_rotVecQuat(gearAxis, gear+3, quat);
}
// j: dof start address
j = m->jnt_dofadr[id];
// dof start address
int jnt_dofadr = m->jnt_dofadr[id];
// sparsity
for (int j = 0; j < 6; j++) {
colind[adr+j] = jnt_dofadr + j;
}
rownnz[i] += 6;
// moment: gear(tran), gearAxis
mju_copy3(moment+i*nv+j, gear);
mju_copy3(moment+i*nv+j+3, gearAxis);
mju_copy3(moment+adr, gear);
mju_copy3(moment+adr+3, gearAxis);
}
break;
@@ -1000,20 +1009,26 @@ void mj_transmission(const mjModel* m, mjData* d) {
mj_jacSite(m, d, jac, 0, id);
mju_subFrom(jac, jacS, 3*nv);
// sparsity
for (int j = 0; j < nv; j++) {
colind[adr+j] = j;
}
rownnz[i] += nv;
// clear moment
mju_zero(moment+i*nv, nv);
mju_zero(moment + adr, nv);
// apply chain rule
for (int j=0; j < nv; j++) {
for (int k=0; k < 3; k++) {
moment[i*nv+j] += dlda[k]*jacA[k*nv+j] + dldv[k]*jac[k*nv+j];
moment[adr+j] += dlda[k]*jacA[k*nv+j] + dldv[k]*jac[k*nv+j];
}
}
// scale by gear ratio
length[i] *= gear[0];
for (int j = 0; j < nv; j++) {
moment[i*nv + j] *= gear[0];
moment[adr+j] *= gear[0];
}
}
break;
@@ -1022,20 +1037,32 @@ void mj_transmission(const mjModel* m, mjData* d) {
length[i] = d->ten_length[id]*gear[0];
// moment: sparse or dense
if (mj_isSparse(m)) {
// clear moment
mju_zero(moment+i*nv, nv);
if (issparse) {
// sparsity
int ten_J_rownnz = d->ten_J_rownnz[id];
int ten_J_rowadr = d->ten_J_rowadr[id];
rownnz[i] += ten_J_rownnz;
mju_copyInt(colind + adr, d->ten_J_colind + ten_J_rowadr, ten_J_rownnz);
int end = d->ten_J_rowadr[id] + d->ten_J_rownnz[id];
for (int j=d->ten_J_rowadr[id]; j < end; j++) {
moment[i*nv + d->ten_J_colind[j]] = d->ten_J[j] * gear[0];
}
mju_scl(moment + adr, d->ten_J + ten_J_rowadr, gear[0], ten_J_rownnz);
} else {
mju_scl(moment + i*nv, d->ten_J + id*nv, gear[0], nv);
// sparsity
for (int j = 0; j < nv; j++) {
colind[adr+j] = j;
}
rownnz[i] += nv;
mju_scl(moment+adr, d->ten_J + id*nv, gear[0], nv);
}
break;
case mjTRN_SITE: // site
// sparsity
for (int j = 0; j < nv; j++) {
colind[adr+j] = j;
}
rownnz[i] += nv;
// get site translation (jac) and rotation (jacS) Jacobians in global frame
mj_jacSite(m, d, jac, jacS, id);
@@ -1050,9 +1077,9 @@ void mj_transmission(const mjModel* m, mjData* d) {
mju_mulMatVec3(wrench+3, d->site_xmat+9*id, gear+3); // rotation
// moment: global Jacobian projected on wrench
mju_mulMatTVec(moment+i*nv, jac, wrench, 3, nv); // translation
mju_mulMatTVec(jac, jacS, wrench+3, 3, nv); // rotation
mju_addTo(moment+i*nv, jac, nv); // add the two
mju_mulMatTVec(moment+adr, jac, wrench, 3, nv); // translation
mju_mulMatTVec(jac, jacS, wrench+3, 3, nv); // rotation
mju_addTo(moment+adr, jac, nv); // add the two
}
// reference site defined
@@ -1089,7 +1116,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
}
// clear moment
mju_zero(moment+i*nv, nv);
mju_zero(moment+adr, nv);
// translational transmission
if (!mju_isZero(gear, 3)) {
@@ -1121,7 +1148,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
mju_mulMatVec3(wrench, d->site_xmat+9*refid, gear);
// moment: global Jacobian projected on wrench
mju_mulMatTVec(moment+i*nv, jac, wrench, 3, nv);
mju_mulMatTVec(moment+adr, jac, wrench, 3, nv);
}
// rotational transmission
@@ -1162,18 +1189,24 @@ void mj_transmission(const mjModel* m, mjData* d) {
// moment_tmp: global Jacobian projected on wrench, add to moment
if (!moment_tmp) moment_tmp = mj_stackAllocNum(d, nv);
mju_mulMatTVec(moment_tmp, jacS, wrench, 3, nv);
mju_addTo(moment+i*nv, moment_tmp, nv);
mju_addTo(moment+adr, moment_tmp, nv);
}
}
break;
case mjTRN_BODY: // body (adhesive contacts)
// sparsity
for (int j = 0; j < nv; j++) {
colind[adr+j] = j;
}
rownnz[i] += nv;
// cannot compute meaningful length, set to 0
length[i] = 0;
// clear moment
mju_zero(moment+i*nv, nv);
mju_zero(moment+adr, nv);
// moment is average of all contact normal Jacobians
{
@@ -1257,15 +1290,16 @@ void mj_transmission(const mjModel* m, mjData* d) {
// moment is average over contact normal Jacobians, make negative for adhesion
if (counter) {
// accumulate active contact Jacobians into moment
mj_mulJacTVec(m, d, moment+i*nv, efc_force);
mj_mulJacTVec(m, d, moment+adr, efc_force);
// add Jacobians from excluded contacts
mju_addTo(moment+i*nv, moment_exclude, nv);
mju_addTo(moment+adr, moment_exclude, nv);
// normalize by total contacts, flip sign
mju_scl(moment+i*nv, moment+i*nv, -1.0/counter, nv);
mju_scl(moment+adr, moment+adr, -1.0/counter, nv);
}
}
break;
default:
src/engine/engine_derivative.c
+10 -1
View File
@@ -827,6 +827,10 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) {
return;
}
// allocate dense actuator_moment row
mj_markStack(d);
mjtNum* moment = mj_stackAllocNum(d, nv);
// process actuators
for (int i=0; i < nu; i++) {
// skip if disabled
@@ -870,9 +874,14 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) {
// add
if (bias_vel != 0) {
addJTBJ(m, d, d->actuator_moment+i*nv, &bias_vel, 1);
mju_sparse2dense(moment, d->actuator_moment, 1, nv, d->moment_rownnz + i,
d->moment_rowadr + i, d->moment_colind);
addJTBJ(m, d, moment, &bias_vel, 1);
}
}
// free space
mj_freeStack(d);
}
src/engine/engine_forward.c
+8 -4
View File
@@ -208,8 +208,11 @@ void mj_fwdVelocity(const mjModel* m, mjData* d) {
mju_mulMatVec(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, m->nv);
}
// actuator velocity: always dense
mju_mulMatVec(d->actuator_velocity, d->actuator_moment, d->qvel, m->nu, m->nv);
// actuator velocity: always sparse
if (!mjDISABLED(mjDSBL_ACTUATION)) {
mju_mulMatVecSparse(d->actuator_velocity, d->actuator_moment, d->qvel, m->nu,
d->moment_rownnz, d->moment_rowadr, d->moment_colind, NULL);
}
// com-based velocities, passive forces, constraint references
mj_comVel(m, d);
@@ -270,7 +273,7 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
TM_START;
int nv = m->nv, nu = m->nu;
mjtNum gain, bias, tau;
mjtNum *prm, *moment = d->actuator_moment, *force = d->actuator_force;
mjtNum *prm, *force = d->actuator_force;
// clear actuator_force
mju_zero(force, nu);
@@ -475,7 +478,8 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
clampVec(force, m->actuator_forcerange, m->actuator_forcelimited, nu, NULL);
// qfrc_actuator = moment' * force
mju_mulMatTVec(d->qfrc_actuator, moment, force, nu, nv);
mju_mulMatTVecSparse(d->qfrc_actuator, d->actuator_moment, force, nu, nv,
d->moment_rownnz, d->moment_rowadr, d->moment_colind);
// actuator-level gravity compensation
if (m->ngravcomp && !mjDISABLED(mjDSBL_GRAVITY) && mju_norm3(m->opt.gravity)) {
src/engine/engine_io.c
-3
View File
@@ -1857,9 +1857,6 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
mju_zero(d->mocap_pos, 3*m->nmocap);
mju_zero(d->mocap_quat, 4*m->nmocap);
// zero out actuator_moment, mj_transmission touches it selectively
mju_zero(d->actuator_moment, m->nv*m->nu);
// copy qpos0 from model
if (m->qpos0) {
memcpy(d->qpos, m->qpos0, m->nq*sizeof(mjtNum));
src/engine/engine_print.c
+6 -3
View File
@@ -93,7 +93,7 @@ static void printSparse(const char* str, const mjtNum* mat, int nr,
const int* rownnz, const int* rowadr,
const int* colind, FILE* fp, const char* float_format) {
// if no data, or too many rows to be visually useful, return
if (!mat || nr > 300) {
if (!mat || !nr || nr > 300) {
return;
}
fprintf(fp, "%s\n", str);
@@ -147,7 +147,7 @@ static void printSparsity(const char* str, int nr, int nc,
// print vector
static void printVector(const char* str, const mjtNum* data, int n, FILE* fp,
const char* float_format) {
if (!data) {
if (!data || !n) {
return;
}
// print str
@@ -1005,7 +1005,10 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
}
printArray("ACTUATOR_LENGTH", m->nu, 1, d->actuator_length, fp, float_format);
printArray("ACTUATOR_MOMENT", m->nu, m->nv, d->actuator_moment, fp, float_format);
printSparsity("actuator_moments", m->nu, m->nv,
d->moment_rowadr, d->moment_rownnz, d->moment_colind, fp);
printSparse("ACTUATOR_MOMENT", d->actuator_moment, m->nu, d->moment_rownnz,
d->moment_rowadr, d->moment_colind, fp, float_format);
printArray("CRB", m->nbody, 10, d->crb, fp, float_format);
if (M) {
src/engine/engine_setconst.c
+21 -7
View File
@@ -29,6 +29,7 @@
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "engine/engine_util_sparse.h"
#include "engine/engine_util_spatial.h"
@@ -66,6 +67,7 @@ static void set0(mjModel* m, mjData* d) {
mj_markStack(d);
mjtNum* jac = mj_stackAllocNum(d, 6*nv);
mjtNum* tmp = mj_stackAllocNum(d, 6*nv);
mjtNum* moment = mj_stackAllocNum(d, nv);
int* cammode = 0;
int* lightmode = 0;
@@ -278,7 +280,9 @@ static void set0(mjModel* m, mjData* d) {
// compute actuator_acc0
for (int i=0; i < m->nu; i++) {
mj_solveM(m, d, tmp, d->actuator_moment+i*nv, 1);
mju_sparse2dense(moment, d->actuator_moment, 1, nv, d->moment_rownnz + i,
d->moment_rowadr + i, d->moment_colind);
mj_solveM(m, d, tmp, moment, 1);
m->actuator_acc0[i] = mju_norm(tmp, nv);
}
} else {
@@ -395,13 +399,16 @@ static void set0(mjModel* m, mjData* d) {
// === interpret biasprm[2] > 0 as dampratio for position-like actuators
// "reflected" inertia (inversely scaled by transmission squared)
mjtNum* transmission = d->actuator_moment + i*nv;
int rownnz = d->moment_rownnz[i];
int rowadr = d->moment_rowadr[i];
mjtNum* transmission = d->actuator_moment + rowadr;
mjtNum mass = 0;
for (int j=0; j < nv; j++) {
for (int j=0; j < rownnz; j++) {
mjtNum trn = mju_abs(transmission[j]);
mjtNum trn2 = trn*trn; // transmission squared
if (trn2 > mjMINVAL) {
mass += m->dof_M0[j] / trn2;
int dof = d->moment_colind[rowadr + j];
mass += m->dof_M0[dof] / trn2;
}
}
@@ -598,11 +605,16 @@ static mjtNum evalAct(const mjModel* m, mjData* d, int index, int side,
// step1: compute inertia and actuator moments
mj_step1(m, d);
// dense actuator_moment row
mj_markStack(d);
mjtNum* moment = mj_stackAllocNum(d, nv);
mju_sparse2dense(moment, d->actuator_moment, 1, nv, d->moment_rownnz + index,
d->moment_rowadr + index, d->moment_colind);
// set force to generate desired acceleration
mj_solveM(m, d, d->qfrc_applied, d->actuator_moment+index*nv, 1);
mj_solveM(m, d, d->qfrc_applied, moment, 1);
mjtNum nrm = mju_norm(d->qfrc_applied, nv);
mju_scl(d->qfrc_applied, d->actuator_moment+index*nv,
(2*side-1)*opt->accel/mjMAX(mjMINVAL, nrm), nv);
mju_scl(d->qfrc_applied, moment, (2*side-1)*opt->accel/mjMAX(mjMINVAL, nrm), nv);
// impose maxforce
nrm = mju_norm(d->qfrc_applied, nv);
@@ -613,6 +625,8 @@ static mjtNum evalAct(const mjModel* m, mjData* d, int index, int side,
// step2: apply force
mj_step2(m, d);
mj_freeStack(d);
// return actuator length
return d->actuator_length[index];
}
src/engine/engine_util_sparse.h
+2 -2
View File
@@ -39,8 +39,8 @@ MJAPI int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
int* rownnz, int* rowadr, int* colind, int nnz);
// convert matrix from sparse to dense
MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind);
MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc, const int* rownnz,
const int* rowadr, const int* colind);
// multiply sparse matrix and dense vector: res = mat * vec
MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,