diff --git a/doc/changelog.rst b/doc/changelog.rst index 77726e1f..3e40dc9a 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,9 +2,17 @@ Changelog ========= -Version 3.3.2 (April 28, 2025) +Upcoming version (not yet release) ---------------------------------- +General +^^^^^^^ +- Refactored island implementation so that island data is memory-contiguous. This speeds up island processing in the + solver and clears the way for the addition of the Newton and PGS solvers (currently only CG is supported). + +Version 3.3.2 (April 28, 2025) +------------------------------ + MJX ^^^ 1. Added inverse dynamics. diff --git a/doc/includes/references.h b/doc/includes/references.h index ce9bd670..6ebc759a 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -171,6 +171,7 @@ struct mjData_ { int nJ; // number of non-zeros in constraint Jacobian int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands + int nidof; // number of dofs in all islands // global properties mjtNum time; // simulation time @@ -381,16 +382,51 @@ struct mjData_ { mjtNum* efc_R; // inverse constraint mass (nefc x 1) int* tendon_efcadr; // first efc address involving tendon; -1: none (ntendon x 1) - // computed by mj_island + // computed by mj_island (island dof structure) int* dof_island; // island id of this dof; -1: none (nv x 1) - int* island_dofnum; // number of dofs in island (nisland x 1) - int* island_dofadr; // start address in island_dofind (nisland x 1) - int* island_dofind; // island dof indices; -1: none (nv x 1) - int* dof_islandind; // dof island indices; -1: none (nv x 1) + int* island_nv; // number of dofs in this island (nisland x 1) + int* island_idofadr; // island start address in idof vector (nisland x 1) + int* island_dofadr; // island start address in dof vector (nisland x 1) + int* map_dof2idof; // map from dof to idof (nv x 1) + int* map_idof2dof; // map from idof to dof; idof >= ni: unconstrained (nv x 1) + + // computed by mj_island (dofs sorted by island) + mjtNum* ifrc_smooth; // net unconstrained force (nidof x 1) + mjtNum* iacc_smooth; // unconstrained acceleration (nidof x 1) + int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1) + int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1) + int* iM_diagnum; // inertia: num of consecutive diagonal elements (nidof x 1) + int* iM_colind; // inertia: column indices of non-zeros (nM x 1) + mjtNum* iM; // total inertia (sparse) (nM x 1) + mjtNum* iLD; // L'*D*L factorization of M (sparse) (nM x 1) + mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1) + mjtNum* iacc; // acceleration (nidof x 1) + + // computed by mj_island (island constraint structure) int* efc_island; // island id of this constraint (nefc x 1) - int* island_efcnum; // number of constraints in island (nisland x 1) - int* island_efcadr; // start address in island_efcind (nisland x 1) - int* island_efcind; // island constraint indices (nefc x 1) + int* island_ne; // number of equality constraints in island (nisland x 1) + int* island_nf; // number of friction constraints in island (nisland x 1) + int* island_nefc; // number of constraints in island (nisland x 1) + int* island_iefcadr; // start address in iefc vector (nisland x 1) + int* map_efc2iefc; // map from efc to iefc (nefc x 1) + int* map_iefc2efc; // map from iefc to efc (nefc x 1) + + // computed by mj_island (constraints sorted by island) + int* iefc_type; // constraint type (mjtConstraint) (nefc x 1) + int* iefc_id; // id of object of specified type (nefc x 1) + int* iefc_J_rownnz; // number of non-zeros in constraint Jacobian row (nefc x 1) + int* iefc_J_rowadr; // row start address in colind array (nefc x 1) + int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) + int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1) + int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1) + int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1) + int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1) + int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1) + mjtNum* iefc_J; // constraint Jacobian (nJ x 1) + mjtNum* iefc_JT; // constraint Jacobian transposed (nJ x 1) + mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1) + mjtNum* iefc_D; // constraint mass (nefc x 1) + mjtNum* iefc_R; // inverse constraint mass (nefc x 1) // computed by mj_projectConstraint (PGS solver) int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1) @@ -408,8 +444,12 @@ struct mjData_ { // computed by mj_fwdConstraint/mj_inverse mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) - mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + mjtNum* iefc_aref; // reference pseudo-acceleration (nefc x 1) + int* iefc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* iefc_force; // constraint force in constraint space (nefc x 1) int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + mjtNum* ifrc_constraint; // constraint force (nidof x 1) // thread pool pointer uintptr_t threadpool; @@ -3174,7 +3214,6 @@ struct mjvSceneState_ { mjtNum* bvh_aabb_dyn; mjtByte* bvh_active; int* island_dofadr; - int* island_dofind; int* dof_island; int* efc_island; int* tendon_efcadr; diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index d6f98775..839b3aed 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -199,6 +199,7 @@ struct mjData_ { int nJ; // number of non-zeros in constraint Jacobian int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands + int nidof; // number of dofs in all islands // global properties mjtNum time; // simulation time @@ -409,16 +410,51 @@ struct mjData_ { mjtNum* efc_R; // inverse constraint mass (nefc x 1) int* tendon_efcadr; // first efc address involving tendon; -1: none (ntendon x 1) - // computed by mj_island + // computed by mj_island (island dof structure) int* dof_island; // island id of this dof; -1: none (nv x 1) - int* island_dofnum; // number of dofs in island (nisland x 1) - int* island_dofadr; // start address in island_dofind (nisland x 1) - int* island_dofind; // island dof indices; -1: none (nv x 1) - int* dof_islandind; // dof island indices; -1: none (nv x 1) + int* island_nv; // number of dofs in this island (nisland x 1) + int* island_idofadr; // island start address in idof vector (nisland x 1) + int* island_dofadr; // island start address in dof vector (nisland x 1) + int* map_dof2idof; // map from dof to idof (nv x 1) + int* map_idof2dof; // map from idof to dof; idof >= ni: unconstrained (nv x 1) + + // computed by mj_island (dofs sorted by island) + mjtNum* ifrc_smooth; // net unconstrained force (nidof x 1) + mjtNum* iacc_smooth; // unconstrained acceleration (nidof x 1) + int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1) + int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1) + int* iM_diagnum; // inertia: num of consecutive diagonal elements (nidof x 1) + int* iM_colind; // inertia: column indices of non-zeros (nM x 1) + mjtNum* iM; // total inertia (sparse) (nM x 1) + mjtNum* iLD; // L'*D*L factorization of M (sparse) (nM x 1) + mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1) + mjtNum* iacc; // acceleration (nidof x 1) + + // computed by mj_island (island constraint structure) int* efc_island; // island id of this constraint (nefc x 1) - int* island_efcnum; // number of constraints in island (nisland x 1) - int* island_efcadr; // start address in island_efcind (nisland x 1) - int* island_efcind; // island constraint indices (nefc x 1) + int* island_ne; // number of equality constraints in island (nisland x 1) + int* island_nf; // number of friction constraints in island (nisland x 1) + int* island_nefc; // number of constraints in island (nisland x 1) + int* island_iefcadr; // start address in iefc vector (nisland x 1) + int* map_efc2iefc; // map from efc to iefc (nefc x 1) + int* map_iefc2efc; // map from iefc to efc (nefc x 1) + + // computed by mj_island (constraints sorted by island) + int* iefc_type; // constraint type (mjtConstraint) (nefc x 1) + int* iefc_id; // id of object of specified type (nefc x 1) + int* iefc_J_rownnz; // number of non-zeros in constraint Jacobian row (nefc x 1) + int* iefc_J_rowadr; // row start address in colind array (nefc x 1) + int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) + int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1) + int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1) + int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1) + int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1) + int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1) + mjtNum* iefc_J; // constraint Jacobian (nJ x 1) + mjtNum* iefc_JT; // constraint Jacobian transposed (nJ x 1) + mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1) + mjtNum* iefc_D; // constraint mass (nefc x 1) + mjtNum* iefc_R; // inverse constraint mass (nefc x 1) // computed by mj_projectConstraint (PGS solver) int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1) @@ -436,8 +472,12 @@ struct mjData_ { // computed by mj_fwdConstraint/mj_inverse mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) - mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + mjtNum* iefc_aref; // reference pseudo-acceleration (nefc x 1) + int* iefc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* iefc_force; // constraint force in constraint space (nefc x 1) int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + mjtNum* ifrc_constraint; // constraint force (nidof x 1) // thread pool pointer uintptr_t threadpool; diff --git a/include/mujoco/mjvisualize.h b/include/mujoco/mjvisualize.h index 0a757cf1..fa0aec4f 100644 --- a/include/mujoco/mjvisualize.h +++ b/include/mujoco/mjvisualize.h @@ -677,7 +677,6 @@ struct mjvSceneState_ { mjtNum* bvh_aabb_dyn; mjtByte* bvh_active; int* island_dofadr; - int* island_dofind; int* dof_island; int* efc_island; int* tendon_efcadr; diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 7c73c598..88e6c0f7 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -739,23 +739,56 @@ X( int, efc_state, MJ_D(nefc), 1 ) // array fields of mjData that are used in the dual problem -#define MJDATA_ARENA_POINTERS_DUAL \ - X( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \ - X( int, efc_AR_rowadr, MJ_D(nefc), 1 ) \ - X( int, efc_AR_colind, MJ_D(nA), 1 ) \ - X( mjtNum, efc_AR, MJ_D(nA), 1 ) +#define MJDATA_ARENA_POINTERS_DUAL \ + X( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \ + X( int, efc_AR_rowadr, MJ_D(nefc), 1 ) \ + X( int, efc_AR_colind, MJ_D(nA), 1 ) \ + X( mjtNum, efc_AR, MJ_D(nA), 1 ) // array fields of mjData that are used for constraint islands -#define MJDATA_ARENA_POINTERS_ISLAND \ - X( int, dof_island, MJ_M(nv), 1 ) \ - X( int, island_dofnum, MJ_D(nisland), 1 ) \ - X( int, island_dofadr, MJ_D(nisland), 1 ) \ - X( int, island_dofind, MJ_M(nv), 1 ) \ - X( int, dof_islandind, MJ_M(nv), 1 ) \ - X( int, efc_island, MJ_D(nefc), 1 ) \ - X( int, island_efcnum, MJ_D(nisland), 1 ) \ - X( int, island_efcadr, MJ_D(nisland), 1 ) \ - X( int, island_efcind, MJ_D(nefc), 1 ) +#define MJDATA_ARENA_POINTERS_ISLAND \ + X( int, dof_island, MJ_M(nv), 1 ) \ + X( int, island_nv, MJ_D(nisland), 1 ) \ + X( int, island_idofadr, MJ_D(nisland), 1 ) \ + X( int, island_dofadr, MJ_D(nisland), 1 ) \ + X( int, map_dof2idof, MJ_M(nv), 1 ) \ + X( int, map_idof2dof, MJ_M(nv), 1 ) \ + X( mjtNum, ifrc_smooth, MJ_D(nidof), 1 ) \ + X( mjtNum, iacc_smooth, MJ_D(nidof), 1 ) \ + X( int, iM_rownnz, MJ_D(nidof), 1 ) \ + X( int, iM_rowadr, MJ_D(nidof), 1 ) \ + X( int, iM_diagnum, MJ_D(nidof), 1 ) \ + X( int, iM_colind, MJ_M(nM), 1 ) \ + X( mjtNum, iM, MJ_M(nM), 1 ) \ + X( mjtNum, iLD, MJ_M(nM), 1 ) \ + X( mjtNum, iLDiagInv, MJ_D(nidof), 1 ) \ + X( mjtNum, iacc, MJ_D(nidof), 1 ) \ + X( int, efc_island, MJ_D(nefc), 1 ) \ + X( int, island_ne, MJ_D(nisland), 1 ) \ + X( int, island_nf, MJ_D(nisland), 1 ) \ + X( int, island_nefc, MJ_D(nisland), 1 ) \ + X( int, island_iefcadr, MJ_D(nisland), 1 ) \ + X( int, map_efc2iefc, MJ_D(nefc), 1 ) \ + X( int, map_iefc2efc, MJ_D(nefc), 1 ) \ + X( int, iefc_type, MJ_D(nefc), 1 ) \ + X( int, iefc_id, MJ_D(nefc), 1 ) \ + X( int, iefc_J_rownnz, MJ_D(nefc), 1 ) \ + X( int, iefc_J_rowadr, MJ_D(nefc), 1 ) \ + X( int, iefc_J_rowsuper, MJ_D(nefc), 1 ) \ + X( int, iefc_J_colind, MJ_D(nJ), 1 ) \ + X( int, iefc_JT_rownnz, MJ_D(nidof), 1 ) \ + X( int, iefc_JT_rowadr, MJ_D(nidof), 1 ) \ + X( int, iefc_JT_rowsuper, MJ_D(nidof), 1 ) \ + X( int, iefc_JT_colind, MJ_D(nJ), 1 ) \ + X( mjtNum, iefc_J, MJ_D(nJ), 1 ) \ + X( mjtNum, iefc_JT, MJ_D(nJ), 1 ) \ + X( mjtNum, iefc_frictionloss, MJ_D(nefc), 1 ) \ + X( mjtNum, iefc_D, MJ_D(nefc), 1 ) \ + X( mjtNum, iefc_R, MJ_D(nefc), 1 ) \ + X( mjtNum, iefc_aref, MJ_D(nefc), 1 ) \ + X( int, iefc_state, MJ_D(nefc), 1 ) \ + X( mjtNum, iefc_force, MJ_D(nefc), 1 ) \ + X( mjtNum, ifrc_constraint, MJ_D(nidof), 1 ) // array fields of mjData that live in d->arena #define MJDATA_ARENA_POINTERS \ @@ -785,6 +818,7 @@ X( int, nJ ) \ X( int, nA ) \ X( int, nisland ) \ + X( int, nidof ) \ X( mjtNum, time ) \ X( uintptr_t, threadpool ) diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index db237c74..c7c3f4d1 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -4896,6 +4896,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='number of detected constraint islands', ), + StructFieldDecl( + name='nidof', + type=ValueType(name='int'), + doc='number of dofs in all islands', + ), StructFieldDecl( name='time', type=ValueType(name='mjtNum'), @@ -5940,11 +5945,19 @@ STRUCTS: Mapping[str, StructDecl] = dict([ array_extent=('nv',), ), StructFieldDecl( - name='island_dofnum', + name='island_nv', type=PointerType( inner_type=ValueType(name='int'), ), - doc='number of dofs in island', + doc='number of dofs in this island', + array_extent=('nisland',), + ), + StructFieldDecl( + name='island_idofadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='island start address in idof vector', array_extent=('nisland',), ), StructFieldDecl( @@ -5952,25 +5965,105 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='int'), ), - doc='start address in island_dofind', + doc='island start address in dof vector', array_extent=('nisland',), ), StructFieldDecl( - name='island_dofind', + name='map_dof2idof', type=PointerType( inner_type=ValueType(name='int'), ), - doc='island dof indices; -1: none', + doc='map from dof to idof', array_extent=('nv',), ), StructFieldDecl( - name='dof_islandind', + name='map_idof2dof', type=PointerType( inner_type=ValueType(name='int'), ), - doc='dof island indices; -1: none', + doc='map from idof to dof; idof >= ni: unconstrained', array_extent=('nv',), ), + StructFieldDecl( + name='ifrc_smooth', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='net unconstrained force', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iacc_smooth', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='unconstrained acceleration', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iM_rownnz', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='inertia: non-zeros in each row', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iM_rowadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='inertia: address of each row in iM_colind', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iM_diagnum', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='inertia: num of consecutive diagonal elements', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iM_colind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='inertia: column indices of non-zeros', + array_extent=('nM',), + ), + StructFieldDecl( + name='iM', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='total inertia (sparse)', + array_extent=('nM',), + ), + StructFieldDecl( + name='iLD', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc="L'*D*L factorization of M (sparse)", + array_extent=('nM',), + ), + StructFieldDecl( + name='iLDiagInv', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='1/diag(D)', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iacc', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='acceleration', + array_extent=('nidof',), + ), StructFieldDecl( name='efc_island', type=PointerType( @@ -5980,7 +6073,23 @@ STRUCTS: Mapping[str, StructDecl] = dict([ array_extent=('nefc',), ), StructFieldDecl( - name='island_efcnum', + name='island_ne', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of equality constraints in island', + array_extent=('nisland',), + ), + StructFieldDecl( + name='island_nf', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of friction constraints in island', + array_extent=('nisland',), + ), + StructFieldDecl( + name='island_nefc', type=PointerType( inner_type=ValueType(name='int'), ), @@ -5988,19 +6097,147 @@ STRUCTS: Mapping[str, StructDecl] = dict([ array_extent=('nisland',), ), StructFieldDecl( - name='island_efcadr', + name='island_iefcadr', type=PointerType( inner_type=ValueType(name='int'), ), - doc='start address in island_efcind', + doc='start address in iefc vector', array_extent=('nisland',), ), StructFieldDecl( - name='island_efcind', + name='map_efc2iefc', type=PointerType( inner_type=ValueType(name='int'), ), - doc='island constraint indices', + doc='map from efc to iefc', + array_extent=('nefc',), + ), + StructFieldDecl( + name='map_iefc2efc', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='map from iefc to efc', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_type', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='constraint type (mjtConstraint)', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_id', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='id of object of specified type', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_J_rownnz', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of non-zeros in constraint Jacobian row', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_J_rowadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='row start address in colind array', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_J_rowsuper', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of subsequent rows in supernode', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_J_colind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='column indices in constraint Jacobian', + array_extent=('nJ',), + ), + StructFieldDecl( + name='iefc_JT_rownnz', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of non-zeros in constraint Jacobian row T', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iefc_JT_rowadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='row start address in colind array T', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iefc_JT_rowsuper', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of subsequent rows in supernode T', + array_extent=('nidof',), + ), + StructFieldDecl( + name='iefc_JT_colind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='column indices in constraint Jacobian T', + array_extent=('nJ',), + ), + StructFieldDecl( + name='iefc_J', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='constraint Jacobian', + array_extent=('nJ',), + ), + StructFieldDecl( + name='iefc_JT', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='constraint Jacobian transposed', + array_extent=('nJ',), + ), + StructFieldDecl( + name='iefc_frictionloss', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='frictionloss (friction)', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_D', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='constraint mass', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_R', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='inverse constraint mass', array_extent=('nefc',), ), StructFieldDecl( @@ -6060,7 +6297,23 @@ STRUCTS: Mapping[str, StructDecl] = dict([ array_extent=('nefc',), ), StructFieldDecl( - name='efc_force', + name='iefc_aref', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='reference pseudo-acceleration', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_state', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='constraint state (mjtConstraintState)', + array_extent=('nefc',), + ), + StructFieldDecl( + name='iefc_force', type=PointerType( inner_type=ValueType(name='mjtNum'), ), @@ -6075,6 +6328,22 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='constraint state (mjtConstraintState)', array_extent=('nefc',), ), + StructFieldDecl( + name='efc_force', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='constraint force in constraint space', + array_extent=('nefc',), + ), + StructFieldDecl( + name='ifrc_constraint', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='constraint force', + array_extent=('nidof',), + ), StructFieldDecl( name='threadpool', type=ValueType(name='uintptr_t'), @@ -8642,13 +8911,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='', ), - StructFieldDecl( - name='island_dofind', - type=PointerType( - inner_type=ValueType(name='int'), - ), - doc='', - ), StructFieldDecl( name='dof_island', type=PointerType( diff --git a/simulate/simulate.cc b/simulate/simulate.cc index 78a5e312..bd255ce3 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -328,7 +328,7 @@ void UpdateProfiler(mj::Simulate* sim, const mjModel* m, const mjData* d) { sim->figconstraint.linedata[start + 4][2*i] = i; // y - int nefc = nisland == 1 ? d->nefc : d->island_efcnum[k]; + int nefc = nisland == 1 ? d->nefc : d->island_nefc[k]; sim->figconstraint.linedata[start + 0][2*i+1] = nefc; const mjSolverStat* stat = d->solver + k*mjNSOLVER + i; sim->figconstraint.linedata[start + 1][2*i+1] = stat->nactive; diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 47918d48..391ed4d2 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -378,50 +378,6 @@ void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* -// multiply Jacobian by vector, for one island -// flg_resunc and flg_vecunc denote whether res/vec are uncompressed -void mj_mulJacVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_resunc, int flg_vecunc) { - // no island, call regular function - if (island < 0) { - mj_mulJacVec(m, d, res, vec); - return; - } - - // sizes - int vecnnz = d->island_dofnum[island]; - int resnnz = d->island_efcnum[island]; - - // indices - int* vecind = d->island_dofind + d->island_dofadr[island]; - int* resind = d->island_efcind + d->island_efcadr[island]; - - // sparse Jacobian - if (mj_isSparse(m)) { - for (int i=0; i < resnnz; i++) { - int row = resind[i]; - int Jnnz = d->efc_J_rownnz[row]; - int Jrowadr = d->efc_J_rowadr[row]; - int* Jind = d->efc_J_colind + Jrowadr; - mjtNum* J = d->efc_J + Jrowadr; - int j = flg_resunc ? row : i; - res[j] = mju_dotSparse2(J, vec, Jnnz, Jind, vecnnz, vecind, flg_vecunc); - } - } - - // dense Jacobian - else { - int nv = m->nv; - for (int i=0; i < resnnz; i++) { - int row = resind[i]; - int j = flg_resunc ? row : i; - res[j] = mju_dotSparse(vec, d->efc_J + nv*row, vecnnz, vecind, flg_vecunc); - } - } -} - - - // multiply JacobianT by vector void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { // exit if no constraints @@ -443,50 +399,6 @@ void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* -// multiply Jacobian transpose by vector, for one island -// flg_resunc and flg_vecunc denote whether res/vec are uncompressed -void mj_mulJacTVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_resunc, int flg_vecunc) { - // no island, call regular function - if (island < 0) { - mj_mulJacTVec(m, d, res, vec); - return; - } - - // sizes - int vecnnz = d->island_efcnum[island]; - int resnnz = d->island_dofnum[island]; - - // indices - int* vecind = d->island_efcind + d->island_efcadr[island]; - int* resind = d->island_dofind + d->island_dofadr[island]; - - // sparse Jacobian - if (mj_isSparse(m)) { - for (int i=0; i < resnnz; i++) { - int row = resind[i]; - int JTnnz = d->efc_JT_rownnz[row]; - int JTrowadr = d->efc_JT_rowadr[row]; - int* JTind = d->efc_JT_colind + JTrowadr; - mjtNum* JT = d->efc_JT + JTrowadr; - int j = flg_resunc ? row : i; - res[j] = mju_dotSparse2(JT, vec, JTnnz, JTind, vecnnz, vecind, flg_vecunc); - } - } - - // dense Jacobian - else { - int nefc = d->nefc; - for (int i=0; i < resnnz; i++) { - int row = resind[i]; - int j = flg_resunc ? row : i; - res[j] = mju_dotSparse(vec, d->efc_JT + nefc*row, vecnnz, vecind, flg_vecunc); - } - } -} - - - //--------------------- instantiate constraints by type -------------------------------------------- // equality constraints @@ -2102,10 +2014,6 @@ void mj_makeConstraint(const mjModel* m, mjData* d) { // supernodes of JT mju_superSparse(m->nv, d->efc_JT_rowsuper, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind); - } else { - if (mjENABLED(mjENBL_ISLAND)) { - mju_transpose(d->efc_JT, d->efc_J, d->nefc, m->nv); - } } // compute diagApprox @@ -2377,25 +2285,17 @@ void mj_referenceConstraint(const mjModel* m, mjData* d) { //---------------------------- update constraint state --------------------------------------------- -// compute efc_state, efc_force, qfrc_constraint, optionally restricted to one island -// island < 0: update all d->nefc constraints -// island >= 0: update only d->island_efcnum[island] constraints -// jar = Jac*qacc-aref is restricted to the island, in the above sense +// compute efc_state, efc_force // optional: cost(qacc) = shat(jar); cone Hessians -void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, - mjtNum cost[1], int flg_coneHessian, int island) { - int ne = d->ne, nf = d->nf; - const mjtNum *D = d->efc_D, *R = d->efc_R, *floss = d->efc_frictionloss; - mjtNum* force = d->efc_force; +void mj_constraintUpdate_impl(int ne, int nf, int nefc, + const mjtNum* D, const mjtNum* R, const mjtNum* floss, + const mjtNum* jar, const int* type, const int* id, + mjContact* contact, int* state, mjtNum* force, mjtNum cost[1], + int flg_coneHessian) { mjtNum s = 0; - int nefc = island < 0 ? d->nefc : d->island_efcnum[island]; - int* efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island]; - - // no constraints: clear qfrc_constraint and cost, return + // no constraints: clear cost, return if (!nefc) { - // can only occur for island == -1 - mju_zero(d->qfrc_constraint, m->nv); if (cost) { *cost = 0; } @@ -2403,55 +2303,49 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, } // compute unconstrained efc_force - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - force[i] = -D[i]*jar[c]; + for (int i=0; i < nefc; i++) { + force[i] = -D[i]*jar[i]; } // update constraints - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - + for (int i=0; i < nefc; i++) { // ==== equality if (i < ne) { if (cost) { - s += 0.5*D[i]*jar[c]*jar[c]; + s += 0.5*D[i]*jar[i]*jar[i]; } - d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; + state[i] = mjCNSTRSTATE_QUADRATIC; continue; } // ==== friction if (i < ne + nf) { // linear negative - if (jar[c] <= -R[i]*floss[i]) { + if (jar[i] <= -R[i]*floss[i]) { if (cost) { - s += -0.5*R[i]*floss[i]*floss[i] - floss[i]*jar[c]; + s += -0.5*R[i]*floss[i]*floss[i] - floss[i]*jar[i]; } force[i] = floss[i]; - - d->efc_state[i] = mjCNSTRSTATE_LINEARNEG; + state[i] = mjCNSTRSTATE_LINEARNEG; } // linear positive - else if (jar[c] >= R[i]*floss[i]) { + else if (jar[i] >= R[i]*floss[i]) { if (cost) { - s += -0.5*R[i]*floss[i]*floss[i] + floss[i]*jar[c]; + s += -0.5*R[i]*floss[i]*floss[i] + floss[i]*jar[i]; } force[i] = -floss[i]; - - d->efc_state[i] = mjCNSTRSTATE_LINEARPOS; + state[i] = mjCNSTRSTATE_LINEARPOS; } // quadratic else { if (cost) { - s += 0.5*D[i]*jar[c]*jar[c]; + s += 0.5*D[i]*jar[i]*jar[i]; } - - d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; + state[i] = mjCNSTRSTATE_QUADRATIC; } continue; } @@ -2459,36 +2353,35 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // ==== contact // non-negative constraint - if (d->efc_type[i] != mjCNSTR_CONTACT_ELLIPTIC) { + if (type[i] != mjCNSTR_CONTACT_ELLIPTIC) { // constraint is satisfied: no cost - if (jar[c] >= 0) { + if (jar[i] >= 0) { force[i] = 0; - d->efc_state[i] = mjCNSTRSTATE_SATISFIED; + state[i] = mjCNSTRSTATE_SATISFIED; } // quadratic else { if (cost) { - s += 0.5*D[i]*jar[c]*jar[c]; + s += 0.5*D[i]*jar[i]*jar[i]; } - - d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; + state[i] = mjCNSTRSTATE_QUADRATIC; } } // contact with elliptic cone else { // get contact - mjContact* con = d->contact + d->efc_id[i]; + mjContact* con = contact + id[i]; mjtNum mu = con->mu, *friction = con->friction; int dim = con->dim; // map to regular dual cone space mjtNum U[6]; - U[0] = jar[c]*mu; + U[0] = jar[i]*mu; for (int j=1; j < dim; j++) { - U[j] = jar[c+j]*friction[j-1]; + U[j] = jar[i+j]*friction[j-1]; } // decompose into normal and tangent @@ -2498,19 +2391,17 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // top zone if (N >= mu*T || (T <= 0 && N >= 0)) { mju_zero(force+i, dim); - - d->efc_state[i] = mjCNSTRSTATE_SATISFIED; + state[i] = mjCNSTRSTATE_SATISFIED; } // bottom zone else if (mu*N+T <= 0 || (T <= 0 && N < 0)) { if (cost) { for (int j=0; j < dim; j++) { - s += 0.5*D[i+j]*jar[c+j]*jar[c+j]; + s += 0.5*D[i+j]*jar[i+j]*jar[i+j]; } } - - d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; + state[i] = mjCNSTRSTATE_QUADRATIC; } // middle zone @@ -2530,12 +2421,12 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, } // set state - d->efc_state[i] = mjCNSTRSTATE_CONE; + state[i] = mjCNSTRSTATE_CONE; // cone Hessian if (flg_coneHessian) { // get Hessian pointer - mjtNum* H = d->contact[d->efc_id[i]].H; + mjtNum* H = contact[id[i]].H; // set first row: (1, -mu/T * U) mjtNum scl = -mu/T; @@ -2546,10 +2437,11 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // set upper block: mu*N/T^3 * U*U' scl = mu*N/(T*T*T); - for (int k=1; k < dim; k++) + for (int k=1; k < dim; k++) { for (int j=k; j < dim; j++) { H[k*dim+j] = scl*U[j]*U[k]; } + } // add to diagonal: (mu^2 - mu*N/T) * I scl = mu*mu - mu*N/T; @@ -2576,19 +2468,14 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // replicate state in all cone dimensions for (int j=1; j < dim; j++) { - d->efc_state[i+j] = d->efc_state[i]; + state[i+j] = state[i]; } // advance to end of contact - c += (dim-1); + i += (dim-1); } } - // compute qfrc_constraint - int flg_vecunc = 1; - int flg_resunc = 1; - mj_mulJacTVec_island(m, d, d->qfrc_constraint, d->efc_force, island, flg_vecunc, flg_resunc); - // assign cost if (cost) { *cost = s; @@ -2601,5 +2488,8 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, // optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, mjtNum cost[1], int flg_coneHessian) { - mj_constraintUpdate_island(m, d, jar, cost, flg_coneHessian, -1); + mj_constraintUpdate_impl(d->ne, d->nf, d->nefc, d->efc_D, d->efc_R, d->efc_frictionloss, + jar, d->efc_type, d->efc_id, d->contact, d->efc_state, d->efc_force, + cost, flg_coneHessian); + mj_mulJacTVec(m, d, d->qfrc_constraint, d->efc_force); } diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h index 05c5fd57..f752b143 100644 --- a/src/engine/engine_core_constraint.h +++ b/src/engine/engine_core_constraint.h @@ -24,6 +24,7 @@ extern "C" { #endif + //-------------------------- Jacobian-related ------------------------------------------------------ // determine type of friction cone @@ -38,16 +39,9 @@ MJAPI int mj_isDual(const mjModel* m); // multiply Jacobian by vector MJAPI void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); -// multiply Jacobian by vector, for one island -MJAPI void mj_mulJacVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_resunc, int flg_vecunc); - // multiply JacobianT by vector MJAPI void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); -// multiply JacobianT by vector, for one island -MJAPI void mj_mulJacTVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_resunc, int flg_vecunc); //-------------------------- utility functions ----------------------------------------------------- @@ -90,6 +84,7 @@ void mj_diagApprox(const mjModel* m, mjData* d); // compute efc_R, efc_D, efc_KDIP, adjust diagApprox void mj_makeImpedance(const mjModel* m, mjData* d); + //---------------------------- top-level API for constraint construction --------------------------- // main driver: call all functions above @@ -101,14 +96,19 @@ MJAPI void mj_projectConstraint(const mjModel* m, mjData* d); // compute efc_vel, efc_aref MJAPI void mj_referenceConstraint(const mjModel* m, mjData* d); +// compute efc_state, efc_force +// optional: cost(qacc) = shat(jar); cone Hessians +MJAPI void mj_constraintUpdate_impl(int ne, int nf, int nefc, + const mjtNum* D, const mjtNum* R, const mjtNum* floss, + const mjtNum* jar, const int* type, const int* id, + mjContact* contact, int* state, mjtNum* force, mjtNum cost[1], + int flg_coneHessian); + // compute efc_state, efc_force, qfrc_constraint // 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); -// compute efc_state, efc_force, qfrc_constraint for one island -MJAPI void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, - mjtNum cost[1], int flg_coneHessian, int island); #ifdef __cplusplus } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 09ea01c2..efd74614 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1803,7 +1803,7 @@ void mj_solveLD_legacy(const mjModel* m, mjtNum* restrict x, int n, // in-place sparse backsubstitution: x = inv(L'*D*L)*x -void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n, +void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n, const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) { // x <- L^-T x for (int i=nv-1; i > 0; i--) { @@ -1819,7 +1819,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int start = rowadr[i]; int end = start + rownnz[i] - 1; for (int adr=start; adr < end; adr++) { - x[colind[adr]] -= qLDs[adr] * x_i; + x[colind[adr]] -= qLD[adr] * x_i; } } } @@ -1832,7 +1832,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, mjtNum x_i; if ((x_i = x[i+offset])) { for (int adr=start; adr < end; adr++) { - x[offset + colind[adr]] -= qLDs[adr] * x_i; + x[offset + colind[adr]] -= qLD[adr] * x_i; } } } @@ -1870,13 +1870,13 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, // one vector if (n == 1) { - x[i] -= mju_dotSparse(qLDs+adr, x, d, colind+adr, /*flg_unc1=*/0); + x[i] -= mju_dotSparse(qLD+adr, x, d, colind+adr, /*flg_unc1=*/0); } // multiple vectors else { for (int offset=0; offset < n*nv; offset+=nv) { - x[i+offset] -= mju_dotSparse(qLDs+adr, x+offset, d, colind+adr, /*flg_unc1=*/0); + x[i+offset] -= mju_dotSparse(qLD+adr, x+offset, d, colind+adr, /*flg_unc1=*/0); } } } @@ -1896,65 +1896,6 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { } -// in-place sparse backsubstitution for one island: x = inv(L'*D*L)*x -// L is in lower triangle of qLD; D is on diagonal of qLD -void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) { - // if no islands, call mj_solveLD - const mjtNum* qLD = d->qLD; - const mjtNum* qLDiagInv = d->qLDiagInv; - if (island < 0) { - mj_solveLD(x, qLD, qLDiagInv, m->nv, 1, - d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); - return; - } - - // local copies of key variables - const int* rownnz = d->M_rownnz; - const int* rowadr = d->M_rowadr; - const int* colind = d->M_colind; - const int* diagnum = m->dof_simplenum; - - // local constants: island specific - int ndof = d->island_dofnum[island]; - const int* dofind = d->island_dofind + d->island_dofadr[island]; - const int* islandind = d->dof_islandind; - - // x <- inv(L') * x; skip simple, exploit sparsity of input vector - for (int k=ndof-1; k >= 0; k--) { - int i = dofind[k]; - mjtNum x_k; - if (!diagnum[i] && (x_k = x[k])) { - int start = rowadr[i]; - int end = start + rownnz[i] - 1; - for (int adr=end-1; adr >= start; adr--) { - x[islandind[colind[adr]]] -= qLD[adr] * x_k; - } - } - } - - // x <- inv(D) * x - for (int k=ndof-1; k >= 0; k--) { - x[k] *= qLDiagInv[dofind[k]]; // x(i) /= L(i,i) - } - - // x <- inv(L) * x; skip simple - for (int k=0; k < ndof; k++) { - int i = dofind[k]; - - // skip diagonal rows - if (diagnum[i]) { - continue; - } - - int start = rowadr[i]; - int end = start + rownnz[i] - 1; - for (int adr=end-1; adr >= start; adr--) { - x[k] -= x[islandind[colind[adr]]] * qLD[adr]; - } - } -} - - // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index 4bc124b0..38decdca 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -71,15 +71,12 @@ MJAPI void mj_solveLD_legacy(const mjModel* m, mjtNum* x, int n, // in-place sparse backsubstitution: x = inv(L'*D*L)*x // handle n vectors at once -MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n, +MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n, const int* rownnz, const int* rowadr, const int* diagnum, const int* colind); // sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); -// sparse backsubstitution for one island: x = inv(L'*D*L)*x, use factorization in d -MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island); - // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, const mjtNum* sqrtInvD, int n); diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index a720f156..b409e1ea 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -631,10 +631,10 @@ static void warmstart(const mjModel* m, mjData* d) { // have island structure: unconstrained qacc = qacc_smooth if (d->nisland > 0) { - for (int i=0; i < nv; i++) { - if (d->dof_island[i] < 0) { - d->qacc[i] = d->qacc_smooth[i]; - } + // loop over unconstrained dofs in map_idof2dof[nidof, nv) + for (int i=d->nidof; i < nv; i++) { + int dof = d->map_idof2dof[i]; + d->qacc[dof] = d->qacc_smooth[dof]; } } @@ -723,22 +723,37 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) { // check if islands are supported int islands_supported = mjENABLED(mjENBL_ISLAND) && - d->nisland > 0 && + nisland > 0 && m->opt.solver == mjSOL_CG && m->opt.noslip_iterations == 0; // run solver over constraint islands if (islands_supported) { - // no threadpool, loop over islands + int nidof = d->nidof; + + // copy CG inputs to islands (vel+acc deps, pos-dependent already copied in mj_island) + mju_gather(d->ifrc_smooth, d->qfrc_smooth, d->map_idof2dof, nidof); + mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof); + mju_gather(d->iacc_smooth, d->qacc_smooth, d->map_idof2dof, nidof); + mju_gather(d->iacc, d->qacc, d->map_idof2dof, nidof); + mju_gather(d->iefc_force, d->efc_force, d->map_iefc2efc, nefc); + mju_gather(d->iefc_aref, d->efc_aref, d->map_iefc2efc, nefc); + + // solve per island if (!d->threadpool) { + // no threadpool, loop over islands for (int island=0; island < nisland; island++) { mj_solCG_island(m, d, island, m->opt.iterations); } - } - else { - // solve using threads + } else { + // have threadpool, solve using threads mj_solCG_island_multithreaded(m, d); } + + // copy back solver outputs (scatter dofs since ni <= nv) + mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof); + mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof); + mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc); } // run solver over all constraints diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 78aed97d..31084c46 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1917,6 +1917,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { d->nJ = 0; d->nA = 0; d->nisland = 0; + d->nidof = 0; // clear global properties d->time = 0; diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 9a8be763..108844d8 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -16,6 +16,7 @@ #include #include +#include #include #include @@ -26,12 +27,65 @@ #include "engine/engine_support.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" +#include "engine/engine_util_sparse.h" #ifdef MEMORY_SANITIZER #include #endif +//-------------------------- local utilities ------------------------------------------------------- + +// clear island-related arena pointers in mjData +static void clearIsland(mjData* d, size_t parena) { +#define X(type, name, nr, nc) d->name = NULL; + MJDATA_ARENA_POINTERS_ISLAND +#undef X + d->nefc = 0; + d->nisland = 0; + d->nidof = 0; + d->parena = parena; + + // poison remaining memory +#ifdef ADDRESS_SANITIZER + ASAN_POISON_MEMORY_REGION( + (char*)d->arena + d->parena, d->narena - d->pstack - d->parena); +#endif +} + + + +// allocate island arrays on arena, return 1 on success, 0 on failure +static int arenaAllocIsland(const mjModel* m, mjData* d) { +#undef MJ_M +#define MJ_M(n) m->n +#undef MJ_D +#define MJ_D(n) d->n + + size_t parena_old = d->parena; + +#define X(type, name, nr, nc) \ + d->name = mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \ + if (!d->name) { \ + mj_warning(d, mjWARN_CNSTRFULL, d->narena); \ + clearIsland(d, parena_old); \ + return 0; \ + } + + MJDATA_ARENA_POINTERS_ISLAND + +#undef X + +#undef MJ_M +#define MJ_M(n) n +#undef MJ_D +#define MJ_D(n) n + return 1; +} + + + +//-------------------------- flood-fill and graph construction ------------------------------------ // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix // arguments: @@ -87,54 +141,6 @@ int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, cons -// clear island-related arena pointers in mjData -static void clearIsland(mjData* d, size_t parena) { -#define X(type, name, nr, nc) d->name = NULL; - MJDATA_ARENA_POINTERS_ISLAND -#undef X - d->nefc = 0; - d->nisland = 0; - d->parena = parena; - - // poison remaining memory -#ifdef ADDRESS_SANITIZER - ASAN_POISON_MEMORY_REGION( - (char*)d->arena + d->parena, d->narena - d->pstack - d->parena); -#endif -} - - - -// allocate island arrays on arena, return 1 on success, 0 on failure -static int arenaAllocIsland(const mjModel* m, mjData* d) { -#undef MJ_M -#define MJ_M(n) m->n -#undef MJ_D -#define MJ_D(n) d->n - - size_t parena_old = d->parena; - -#define X(type, name, nr, nc) \ - d->name = mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \ - if (!d->name) { \ - mj_warning(d, mjWARN_CNSTRFULL, d->narena); \ - clearIsland(d, parena_old); \ - return 0; \ - } - - MJDATA_ARENA_POINTERS_ISLAND - -#undef X - -#undef MJ_M -#define MJ_M(n) n -#undef MJ_D -#define MJ_D(n) n - return 1; -} - - - // return upper bound on number of tree-tree edges static int countMaxEdge(const mjModel* m, const mjData* d) { int nedge_max = 0; @@ -411,14 +417,17 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg +//-------------------------- main entry-point ----------------------------------------------------- + // discover islands: -// nisland, island_dofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext +// nisland, island_idofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext void mj_island(const mjModel* m, mjData* d) { int nv = m->nv, nefc = d->nefc, ntree=m->ntree; // no constraints: quick return if (!nefc || m->nflex) { // TODO: add flex support to island discovery d->nisland = 0; + d->nidof = 0; return; } @@ -454,86 +463,201 @@ void mj_island(const mjModel* m, mjData* d) { int* stack = mjSTACKALLOC(d, nedge, int); d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack); + // no islands found: quick return + if (!d->nisland) { + d->nidof = 0; + mj_freeStack(d); + return; + } + + // count ni: total number of dofs in islands + int nidof = 0; + for (int i=0; i < nv; i++) { + nidof += (tree_island[m->dof_treeid[i]] >= 0); + } + d->nidof = nidof; + // allocate island arrays on arena if (!arenaAllocIsland(m, d)) { mj_freeStack(d); return; } - int nisland = d->nisland; // local copy + // local copy + int nisland = d->nisland; - // compute dof_island, island_dofnum - int num_dof_unc = 0; // number of unconstrained dofs - mju_zeroInt(d->island_dofnum, nisland); + + // ------------------------------------- degrees of freedom -------------------------------------- + + // compute dof_island, island_nv + mju_zeroInt(d->island_nv, nisland); for (int i=0; i < nv; i++) { - // dof_island - int island = tree_island[m->dof_treeid[i]]; + // assign dofs to islands + int island = tree_island[m->dof_treeid[i]]; // -1 if unconstrained d->dof_island[i] = island; - // island_dofnum + // increment island_nv if (island >= 0) { - d->island_dofnum[island]++; - } else { - num_dof_unc++; + d->island_nv[island]++; } } - // compute island_dofadr - if (nisland) d->island_dofadr[0] = 0; + // compute island_idofadr (cumsum of island_nv) + d->island_idofadr[0] = 0; for (int i=1; i < nisland; i++) { - d->island_dofadr[i] = d->island_dofadr[i-1] + d->island_dofnum[i-1]; + d->island_idofadr[i] = d->island_idofadr[i-1] + d->island_nv[i-1]; } - // reset island_dofnum - mju_zeroInt(d->island_dofnum, nisland); - - // compute dof_islandind, island_dofind - int num_dof_island = 0; - for (int i=0; i < nv; i++) { - int island = d->dof_island[i]; + // compute dof <-> idof maps + int* island_nv2 = mjSTACKALLOC(d, nisland + 1, int); // last element counts unconstrained dofs + mju_zeroInt(island_nv2, nisland + 1); + for (int dof=0; dof < nv; dof++) { + int island = d->dof_island[dof]; + int idof; if (island >= 0) { - d->island_dofind[d->island_dofadr[island] + d->island_dofnum[island]] = i; - d->dof_islandind[i] = d->island_dofnum[island]++; - num_dof_island++; + // constrained dof + idof = d->island_idofadr[island] + island_nv2[island]++; } else { - d->dof_islandind[i] = -1; + // unconstrained dof + idof = nidof + island_nv2[nisland]++; } + + d->map_dof2idof[dof] = idof; + d->map_idof2dof[idof] = dof; // only the first ni elements of map_idof2dof are in some island } - // sanity check, SHOULD NOT OCCUR - if (num_dof_island + num_dof_unc != nv) { - mjERROR("not all islands assigned to dofs"); + // SHOULD NOT OCCUR + if (!mju_compare(island_nv2, d->island_nv, nisland)) mjERROR("island_nv miscount"); + if (nidof + island_nv2[nisland] != nv) mjERROR("miscount of unconstrained dofs"); + + // compute island_dofadr (used for visualization) + for (int i=0; i < nisland; i++) { + d->island_dofadr[i] = d->map_idof2dof[d->island_idofadr[i]]; } - // finalize dof_islandind: set remaining indices to -1 - for (int i=num_dof_island; i < nv; i++) { - d->island_dofind[i] = -1; + // local CSR copy of qM + mjtNum* qM = mjSTACKALLOC(d, m->nM, mjtNum); + mju_gather(qM, d->qM, d->mapM2M, m->nM); + + // inertia: block-diagonalize both iLD <- qLD and iM <- qM + mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind, + d->qLD, d->M_rownnz, d->M_rowadr, d->M_colind, + nidof, nisland, + d->map_idof2dof, d->map_dof2idof, + d->island_idofadr, d->island_idofadr, + d->iM, qM); + mju_gather(d->iLDiagInv, d->qLDiagInv, d->map_idof2dof, nidof); + + // compute iM_diagnum (dof_simplenum per island) + int count = 0; + int dof_next = d->map_idof2dof[nidof-1]; + for (int i=nidof-1; i >= 0; i--) { + // check if island boundary was crossed + int dof = d->map_idof2dof[i]; + int island_boundary = (d->dof_island[dof] != d->dof_island[dof_next]); + dof_next = dof; + + // accumulate and set simple dof (diagonal row) counter + if (m->dof_simplenum[dof] && !island_boundary) { + count++; // increment counter + } else { + count = 0; // reset + } + d->iM_diagnum[i] = count; } - // compute efc_island, island_efcnum - mju_zeroInt(d->island_efcnum, nisland); + + + // ------------------------------------- constraints --------------------------------------------- + + // compute efc_island, island_{ne,nf,nefc} + mju_zeroInt(d->island_ne, nisland); + mju_zeroInt(d->island_nf, nisland); + mju_zeroInt(d->island_nefc, nisland); for (int i=0; i < nefc; i++) { int tree[2]; treeFirst(m, d, tree, i); int island = tree_island[tree[0]]; d->efc_island[i] = island; - d->island_efcnum[island]++; + d->island_nefc[island]++; + switch (d->efc_type[i]) { + case mjCNSTR_EQUALITY: + d->island_ne[island]++; + break; + case mjCNSTR_FRICTION_DOF: + case mjCNSTR_FRICTION_TENDON: + d->island_nf[island]++; + break; + default: + break; + } } - // compute island_efcadr - if (nisland) d->island_efcadr[0] = 0; + // compute island_iefcadr (cumsum of island_nefc) + d->island_iefcadr[0] = 0; for (int i=1; i < nisland; i++) { - d->island_efcadr[i] = d->island_efcadr[i-1] + d->island_efcnum[i-1]; + d->island_iefcadr[i] = d->island_iefcadr[i-1] + d->island_nefc[i-1]; } - // reset island_efcnum - mju_zeroInt(d->island_efcnum, nisland); - - // compute efc_islandind - for (int i=0; i < nefc; i++) { - int island = d->efc_island[i]; - d->island_efcind[d->island_efcadr[island] + (d->island_efcnum[island]++)] = i; + // compute efc <-> iefc maps + int* island_nefc2 = island_nv2; // reuse island_nv2 + mju_zeroInt(island_nefc2, nisland); + for (int c=0; c < nefc; c++) { + int island = d->efc_island[c]; + int ic = d->island_iefcadr[island] + island_nefc2[island]++; + d->map_efc2iefc[c] = ic; + d->map_iefc2efc[ic] = c; } + // SHOULD NOT OCCUR + if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount"); + + // dense: block-diagonalize Jacobian + if (!mj_isSparse(m)) { + mju_blockDiag(d->iefc_J, d->efc_J, + nv, nidof, nisland, + d->map_iefc2efc, d->map_idof2dof, + d->island_nefc, d->island_nv, + d->island_iefcadr, d->island_idofadr); + } + + // sparse + else { + // block-diagonalize Jacobian + mju_blockDiagSparse(d->iefc_J, d->iefc_J_rownnz, d->iefc_J_rowadr, d->iefc_J_colind, + d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, + nefc, nisland, + d->map_iefc2efc, d->map_dof2idof, + d->island_iefcadr, d->island_idofadr, NULL, NULL); + + // recompute rowsuper per island + for (int island=0; island < nisland; island++) { + int adr = d->island_iefcadr[island]; + mju_superSparse(d->island_nefc[island], d->iefc_J_rowsuper + adr, + d->iefc_J_rownnz + adr, d->iefc_J_rowadr + adr, d->iefc_J_colind); + } + + // block-diagonalize Jacobian-transpose + mju_blockDiagSparse(d->iefc_JT, d->iefc_JT_rownnz, d->iefc_JT_rowadr, d->iefc_JT_colind, + d->efc_JT, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, + nidof, nisland, + d->map_idof2dof, d->map_efc2iefc, + d->island_idofadr, d->island_iefcadr, NULL, NULL); + + // recompute rowsuper per island + for (int island=0; island < nisland; island++) { + int adr = d->island_idofadr[island]; + mju_superSparse(d->island_nv[island], d->iefc_JT_rowsuper + adr, + d->iefc_JT_rownnz + adr, d->iefc_JT_rowadr + adr, d->iefc_JT_colind); + } + } + + // copy position-dependent efc vectors required by solver + mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, nefc); + mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, nefc); + mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, nefc); + mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, nefc); + mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, nefc); + mj_freeStack(d); } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 49d0afa3..61a9ef02 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1392,27 +1392,30 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_DOFNUM"); + fprintf(fp, NAME_FORMAT, "ISLAND_NV"); for (int i = 0; i < d->nisland; i++) { - fprintf(fp, " %d", d->island_dofnum[i]); + fprintf(fp, " %d", d->island_nv[i]); } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_DOFADR"); + fprintf(fp, NAME_FORMAT, "ISLAND_IDOFADR"); for (int i = 0; i < d->nisland; i++) { - fprintf(fp, " %d", d->island_dofadr[i]); + fprintf(fp, " %d", d->island_idofadr[i]); } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_DOFIND"); + fprintf(fp, NAME_FORMAT, "MAP_IDOF2DOF"); for (int i = 0; i < m->nv; i++) { - fprintf(fp, " %d", d->island_dofind[i]); - } - fprintf(fp, "\n\n"); + int dof = d->map_idof2dof[i]; + if (i > 0) { + int dofprev = d->map_idof2dof[i-1]; - fprintf(fp, NAME_FORMAT, "DOF_ISLANDIND"); - for (int i = 0; i < m->nv; i++) { - fprintf(fp, " %d", d->dof_islandind[i]); + // print '|' at island boundaries + if (d->dof_island[dof] != d->dof_island[dofprev]) { + fprintf(fp, " |"); + } + } + fprintf(fp, " %d", dof); } fprintf(fp, "\n\n"); @@ -1422,21 +1425,30 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_EFCNUM"); + fprintf(fp, NAME_FORMAT, "ISLAND_NEFC"); for (int i = 0; i < d->nisland; i++) { - fprintf(fp, " %d", d->island_efcnum[i]); + fprintf(fp, " %d", d->island_nefc[i]); } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_EFCADR"); + fprintf(fp, NAME_FORMAT, "ISLAND_IEFCADR"); for (int i = 0; i < d->nisland; i++) { - fprintf(fp, " %d", d->island_efcadr[i]); + fprintf(fp, " %d", d->island_iefcadr[i]); } fprintf(fp, "\n\n"); - fprintf(fp, NAME_FORMAT, "ISLAND_EFCIND"); + fprintf(fp, NAME_FORMAT, "MAP_IEFC2EFC"); for (int i = 0; i < d->nefc; i++) { - fprintf(fp, " %d", d->island_efcind[i]); + int efc = d->map_iefc2efc[i]; + if (i > 0) { + int efcprev = d->map_iefc2efc[i-1]; + + // print '|' at island boundaries + if (d->efc_island[efc] != d->efc_island[efcprev]) { + fprintf(fp, " |"); + } + } + fprintf(fp, " %d", efc); } fprintf(fp, "\n\n"); } diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 8e83ace5..fb9dedd1 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -766,13 +766,55 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) { // CG context struct _mjCGContext { int flg_Newton; // 1: Newton, 0: CG - - // island-related int island; // current island index, -1 if monolithic + + // sizes int nv; // number of dofs - int nefc; // number of constraints - int* dofind; // dof indices of this island, NULL if monolithic - int* efcind; // constraint indices of this island, NULL if monolithic + int ne; // number of equalities + int nf; // number of friction constraints + int nefc; // number of all constraints + + // contact array + mjContact* contact; + + // dof arrays + const mjtNum* qfrc_smooth; + const mjtNum* qacc_smooth; + mjtNum* qfrc_constraint; + mjtNum* qacc; + + // inertia + const int* M_rownnz; + const int* M_rowadr; + const int* M_diagnum; + const int* M_colind; + const int* dof_Madr; + const int* dof_parentid; + const mjtNum* qM; + const mjtNum* qLD; + const mjtNum* qLDiagInv; + + // efc arrays + const mjtNum* efc_D; + const mjtNum* efc_R; + const mjtNum* efc_frictionloss; + const mjtNum* efc_aref; + const int* efc_id; + const int* efc_type; + mjtNum* efc_force; + int* efc_state; + + // Jacobians + const int* J_rownnz; + const int* J_rowadr; + const int* J_rowsuper; + const int* J_colind; + const int* JT_rownnz; + const int* JT_rowadr; + const int* JT_rowsuper; + const int* JT_colind; + const mjtNum* J; + const mjtNum* JT; // common arrays (CGallocate) mjtNum* Jaref; // Jac*qacc - aref (nefc x 1) @@ -793,7 +835,7 @@ struct _mjCGContext { int* L_rownnz; // Hessian factor row nonzeros (nv x 1) int* L_rowadr; // Hessian factor row addresses (nv x 1) - // Newton arrays, computed-size (HessianMake) + // Newton arrays, computed-size (MakeHessian) int nH; // number of nonzeros in Hessian H int* H_colind; // Hessian column indices (nH x 1) mjtNum* H; // Hessian (nH x 1) @@ -818,23 +860,130 @@ struct _mjCGContext { typedef struct _mjCGContext mjCGContext; + +// set sizes and pointers to mjData arrays in mjCGContext +static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int island) { + int is_sparse = mj_isSparse(m); + ctx->contact = d->contact; + ctx->island = island; + + // set sizes and pointers (monolithic) + if (island < 0) { + // sizes + ctx->nv = m->nv; + ctx->ne = d->ne; + ctx->nf = d->nf; + ctx->nefc = d->nefc; + + // dof arrays + ctx->qfrc_smooth = d->qfrc_smooth; + ctx->qfrc_constraint = d->qfrc_constraint; + ctx->qacc_smooth = d->qacc_smooth; + ctx->qacc = d->qacc; + + // inertia + ctx->M_rownnz = d->M_rownnz; + ctx->M_rowadr = d->M_rowadr; + ctx->M_diagnum = m->dof_simplenum; + ctx->M_colind = d->M_colind; + ctx->dof_Madr = m->dof_Madr; + ctx->dof_parentid = m->dof_parentid; + ctx->qM = d->qM; + ctx->qLD = d->qLD; + ctx->qLDiagInv = d->qLDiagInv; + + // efc arrays + ctx->efc_D = d->efc_D; + ctx->efc_R = d->efc_R; + ctx->efc_frictionloss = d->efc_frictionloss; + ctx->efc_aref = d->efc_aref; + ctx->efc_id = d->efc_id; + ctx->efc_type = d->efc_type; + ctx->efc_force = d->efc_force; + ctx->efc_state = d->efc_state; + + // Jacobians + ctx->J = d->efc_J; + if (is_sparse) { + ctx->J_rownnz = d->efc_J_rownnz; + ctx->J_rowadr = d->efc_J_rowadr; + ctx->J_rowsuper = d->efc_J_rowsuper; + ctx->J_colind = d->efc_J_colind; + ctx->JT_rownnz = d->efc_JT_rownnz; + ctx->JT_rowadr = d->efc_JT_rowadr; + ctx->JT_rowsuper = d->efc_JT_rowsuper; + ctx->JT_colind = d->efc_JT_colind; + ctx->JT = d->efc_JT; + } + } + + // set sizes and pointers (per-island) + else { + // sizes + ctx->nv = d->island_nv[island]; + ctx->ne = d->island_ne[island]; + ctx->nf = d->island_nf[island]; + ctx->nefc = d->island_nefc[island]; + + // dof arrays + int idofadr = d->island_idofadr[island]; + ctx->qfrc_smooth = d->ifrc_smooth + idofadr; + ctx->qfrc_constraint = d->ifrc_constraint + idofadr; + ctx->qacc_smooth = d->iacc_smooth + idofadr; + ctx->qacc = d->iacc + idofadr; + + // inertia + ctx->M_rownnz = d->iM_rownnz + idofadr; + ctx->M_rowadr = d->iM_rowadr + idofadr; + ctx->M_diagnum = d->iM_diagnum + idofadr; + ctx->M_colind = d->iM_colind; + ctx->qM = d->iM; + ctx->qLD = d->iLD; + ctx->qLDiagInv = d->iLDiagInv + idofadr; + + // efc arrays + int iefcadr = d->island_iefcadr[island]; + ctx->efc_D = d->iefc_D + iefcadr; + ctx->efc_R = d->iefc_R + iefcadr; + ctx->efc_frictionloss = d->iefc_frictionloss + iefcadr; + ctx->efc_aref = d->iefc_aref + iefcadr; + ctx->efc_id = d->iefc_id + iefcadr; + ctx->efc_type = d->iefc_type + iefcadr; + ctx->efc_force = d->iefc_force + iefcadr; + ctx->efc_state = d->iefc_state + iefcadr; + + // Jacobians + if (!is_sparse) { + ctx->J = d->iefc_J + d->nidof * iefcadr; + } else { + ctx->J_rownnz = d->iefc_J_rownnz + iefcadr; + ctx->J_rowadr = d->iefc_J_rowadr + iefcadr; + ctx->J_rowsuper = d->iefc_J_rowsuper + iefcadr; + ctx->J_colind = d->iefc_J_colind; + ctx->JT_rownnz = d->iefc_JT_rownnz + idofadr; + ctx->JT_rowadr = d->iefc_JT_rowadr + idofadr; + ctx->JT_rowsuper = d->iefc_JT_rowsuper + idofadr; + ctx->JT_colind = d->iefc_JT_colind; + ctx->J = d->iefc_J; + ctx->JT = d->iefc_JT; + } + } +} + + + // allocate fixed-size arrays in mjCGContext // mj_{mark/free}Stack in calling function! -static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, - int island, int flg_Newton) { +static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, int island, int flg_Newton) { // clear everything memset(ctx, 0, sizeof(mjCGContext)); - // get sizes - int nv = island < 0 ? m->nv : d->island_dofnum[island]; - int nefc = island < 0 ? d->nefc : d->island_efcnum[island]; + // set sizes and pointers + CGpointers(m, d, ctx, island); - // island-related - ctx->island = island; - ctx->nv = nv; - ctx->nefc = nefc; - ctx->dofind = island < 0 ? NULL : d->island_dofind + d->island_dofadr[island]; - ctx->efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island]; + // local sizes + int nv = ctx->nv; + int nefc = ctx->nefc; // common arrays ctx->Jaref = mjSTACKALLOC(d, nefc, mjtNum); @@ -849,7 +998,7 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, // Newton only, known-size arrays ctx->flg_Newton = flg_Newton; if (flg_Newton) { - ctx->D = mjSTACKALLOC(d, nefc, mjtNum); + ctx->D = mjSTACKALLOC(d, nefc, mjtNum); // sparse Newton only if (mj_isSparse(m)) { @@ -866,28 +1015,35 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, // update efc_force, qfrc_constraint, cost-related -static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) { +static void CGupdateConstraint(mjCGContext* ctx) { int nefc = ctx->nefc, nv = ctx->nv; - const int* dofind = ctx->dofind; - const int* efcind = ctx->efcind; // update constraints - mj_constraintUpdate_island(m, d, ctx->Jaref, &(ctx->cost), ctx->flg_Newton, ctx->island); + mj_constraintUpdate_impl(ctx->ne, ctx->nf, ctx->nefc, ctx->efc_D, ctx->efc_R, + ctx->efc_frictionloss, ctx->Jaref, ctx->efc_type, ctx->efc_id, + ctx->contact, ctx->efc_state, ctx->efc_force, + &(ctx->cost), ctx->flg_Newton); + + // compute qfrc_constraint (dense or sparse) + if (!ctx->JT) { + mju_mulMatTVec(ctx->qfrc_constraint, ctx->J, ctx->efc_force, nefc, nv); + } else { + mju_mulMatVecSparse(ctx->qfrc_constraint, ctx->JT, ctx->efc_force, nv, + ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper); + } // count active and cone ctx->nactive = 0; ctx->ncone = 0; - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - ctx->nactive += (d->efc_state[i] != mjCNSTRSTATE_SATISFIED); - ctx->ncone += (d->efc_state[i] == mjCNSTRSTATE_CONE); + for (int i=0; i < nefc; i++) { + ctx->nactive += (ctx->efc_state[i] != mjCNSTRSTATE_SATISFIED); + ctx->ncone += (ctx->efc_state[i] == mjCNSTRSTATE_CONE); } // add Gauss cost, set in quadratic[0] mjtNum Gauss = 0; - for (int c=0; c < nv; c++) { - int i = dofind ? dofind[c] : c; - Gauss += 0.5 * (ctx->Ma[c] - d->qfrc_smooth[i]) * (d->qacc[i] - d->qacc_smooth[i]); + for (int i=0; i < nv; i++) { + Gauss += 0.5 * (ctx->Ma[i] - ctx->qfrc_smooth[i]) * (ctx->qacc[i] - ctx->qacc_smooth[i]); } ctx->quadGauss[0] = Gauss; @@ -895,22 +1051,20 @@ static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) { } -// TODO(tassa): Restore mjData const-ness. + // update grad, Mgrad -static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) { +static void CGupdateGradient(mjCGContext* ctx) { int nv = ctx->nv; - const int* dofind = ctx->dofind; // grad = M*qacc - qfrc_smooth - qfrc_constraint - for (int c=0; c < nv; c++) { - int i = dofind ? dofind[c] : c; - ctx->grad[c] = ctx->Ma[c] - d->qfrc_smooth[i] - d->qfrc_constraint[i]; + for (int i=0; i < nv; i++) { + ctx->grad[i] = ctx->Ma[i] - ctx->qfrc_smooth[i] - ctx->qfrc_constraint[i]; } // Newton: Mgrad = H \ grad // TODO: b/295296178 - add island support to Newton solver if (ctx->flg_Newton) { - if (mj_isSparse(m)) { + if (ctx->L_rowadr) { mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L), ctx->grad, nv, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind); } else { @@ -921,44 +1075,32 @@ static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) { // CG: Mgrad = M \ grad else { mju_copy(ctx->Mgrad, ctx->grad, nv); - mj_solveM_island(m, d, ctx->Mgrad, ctx->island); + mj_solveLD(ctx->Mgrad, ctx->qLD, ctx->qLDiagInv, nv, 1, + ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind); } } // prepare quadratic polynomials and contact cone quantities -static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) { - int nv = ctx->nv, nefc = ctx->nefc, island = ctx->island; - const int* dofind = ctx->dofind; - const int* efcind = ctx->efcind; +static void CGprepare(mjCGContext* ctx) { + int nv = ctx->nv, nefc = ctx->nefc; const mjtNum* v = ctx->search; // Gauss: alpha^2*0.5*v'*M*v + alpha*v'*(Ma-qfrc_smooth) + 0.5*(a-qacc_smooth)'*(Ma-qfrc_smooth) // quadGauss[0] already computed in CGupdateConstraint - mjtNum v_dot_smooth; - if (island < 0) { - v_dot_smooth = mju_dot(d->qfrc_smooth, v, nv); - } else { - v_dot_smooth = 0; - for (int c=0; c < nv; c++) { - v_dot_smooth += d->qfrc_smooth[dofind[c]] * v[c]; - } - } - ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - v_dot_smooth; + ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - mju_dot(ctx->qfrc_smooth, v, nv); ctx->quadGauss[2] = 0.5*mju_dot(v, ctx->Mv, nv); // process constraints - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - + for (int i=0; i < nefc; i++) { // pointers to numeric data - const mjtNum* Jv = ctx->Jv + c; - const mjtNum* Jaref = ctx->Jaref + c; - const mjtNum* D = d->efc_D + i; + const mjtNum* Jv = ctx->Jv + i; + const mjtNum* Jaref = ctx->Jaref + i; + const mjtNum* D = ctx->efc_D + i; // pointer to this quadratic - mjtNum* quad = ctx->quad + 3*c; + mjtNum* quad = ctx->quad + 3*i; // init with scalar quadratic mjtNum DJ0 = D[0]*Jaref[0]; @@ -967,12 +1109,12 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) { quad[2] = Jv[0]*D[0]*Jv[0]; // elliptic cone: extra processing - if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { + if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // extract contact info - mjContact* con = d->contact + d->efc_id[i]; + const mjContact* con = ctx->contact + ctx->efc_id[i]; int dim = con->dim; mjtNum U[6], V[6], UU = 0, UV = 0, VV = 0, mu = con->mu; - mjtNum* friction = con->friction; + const mjtNum* friction = con->friction; // complete vector quadratic (for bottom zone) for (int j=1; j < dim; j++) { @@ -1006,7 +1148,7 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) { quad[8] = D[0] / ((mu*mu) * (1 + (mu*mu))); // advance to next constraint - c += (dim-1); + i += (dim-1); } // apply scaling @@ -1028,9 +1170,8 @@ typedef struct _mjCGPnt mjCGPnt; // evaluate linesearch cost, return first and second derivatives -static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* p) { - int ne = d->ne, nf = d->nf, nefc = ctx->nefc; - const int* efcind = ctx->efcind; +static void CGeval(mjCGContext* ctx, mjCGPnt* p) { + int ne = ctx->ne, nf = ctx->nf, nefc = ctx->nefc; // clear result mjtNum cost = 0, alpha = p->alpha; @@ -1041,26 +1182,24 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* mju_copy3(quadTotal, ctx->quadGauss); // process constraints - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - + for (int i=0; i < nefc; i++) { // equality if (i < ne) { - mju_addTo3(quadTotal, ctx->quad+3*c); + mju_addTo3(quadTotal, ctx->quad+3*i); continue; } // friction if (i < ne + nf) { // search point, friction loss, bound (Rf) - mjtNum start = ctx->Jaref[c], dir = ctx->Jv[c]; + mjtNum start = ctx->Jaref[i], dir = ctx->Jv[i]; mjtNum x = start + alpha*dir; - mjtNum f = d->efc_frictionloss[i]; - mjtNum Rf = d->efc_R[i]*f; + mjtNum f = ctx->efc_frictionloss[i]; + mjtNum Rf = ctx->efc_R[i]*f; // -bound < x < bound : quadratic if (-Rf < x && x < Rf) { - mju_addTo3(quadTotal, ctx->quad+3*c); + mju_addTo3(quadTotal, ctx->quad+3*i); } // x < -bound : linear negative @@ -1078,10 +1217,10 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* } // limit and contact - if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone + if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone // extract contact info - mjContact* con = d->contact + d->efc_id[i]; - mjtNum* quad = ctx->quad + 3*c; + const mjContact* con = ctx->contact + ctx->efc_id[i]; + mjtNum* quad = ctx->quad + 3*i; int dim = con->dim; mjtNum mu = con->mu; @@ -1137,14 +1276,14 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* } // advance to next constraint - c += (dim-1); + i += (dim-1); } else { // inequality // search point - mjtNum x = ctx->Jaref[c] + alpha*ctx->Jv[c]; + mjtNum x = ctx->Jaref[i] + alpha*ctx->Jv[i]; // active if (x < 0) { - mju_addTo3(quadTotal, ctx->quad+3*c); + mju_addTo3(quadTotal, ctx->quad+3*i); } } } @@ -1170,7 +1309,7 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* // update bracket point given 3 candidate points -static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx, +static int updateBracket(mjCGContext* ctx, mjCGPnt* p, const mjCGPnt candidates[3], mjCGPnt* pnext) { int flag = 0; for (int i=0; i < 3; i++) { @@ -1192,7 +1331,7 @@ static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx, // compute next point if updated if (flag) { pnext->alpha = p->alpha - p->deriv[0]/p->deriv[1]; - CGeval(m, d, ctx, pnext); + CGeval(ctx, pnext); } return flag; @@ -1201,8 +1340,8 @@ static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx, // line search -static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { - int nv = ctx->nv; +static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations) { + int nv = ctx->nv, nefc = ctx->nefc; mjCGPnt p0, p1, p2, pmid, p1next, p2next; // clear results @@ -1218,23 +1357,36 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { } // compute scaled gradtol and slope scaling - mjtNum gtol = m->opt.tolerance * m->opt.ls_tolerance * snorm / ctx->scale; + mjtNum gtol = tolerance * snorm / ctx->scale; mjtNum slopescl = ctx->scale / snorm; - // compute Mv, Jv - mj_mulM_island(m, d, ctx->Mv, ctx->search, ctx->island, /*flg_vecunc=*/0); - mj_mulJacVec_island(m, d, ctx->Jv, ctx->search, ctx->island, /*flg_resunc=*/0, /*flg_vecunc=*/0); + // compute Mv = M * v (island or monolithic) + if (ctx->island >= 0) { + mju_mulSymVecSparse(ctx->Mv, ctx->qM, ctx->search, nv, + ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind); + } else { + mj_mulM_impl(ctx->Mv, ctx->search, nv, ctx->qM, + ctx->dof_Madr, ctx->dof_parentid, ctx->M_diagnum); + } + + // compute Jv = J * search (dense or sparse) + if (!ctx->J_rowadr) { + mju_mulMatVec(ctx->Jv, ctx->J, ctx->search, nefc, nv); + } else { + mju_mulMatVecSparse(ctx->Jv, ctx->J, ctx->search, nefc, + ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->J_rowsuper); + } // prepare quadratics and cones - CGprepare(m, d, ctx); + CGprepare(ctx); // init at alpha = 0, save p0.alpha = 0; - CGeval(m, d, ctx, &p0); + CGeval(ctx, &p0); // always attempt one Newton step p1.alpha = p0.alpha - p0.deriv[0]/p0.deriv[1]; - CGeval(m, d, ctx, &p1); + CGeval(ctx, &p1); if (p0.cost < p1.cost) { p1 = p0; } @@ -1289,14 +1441,14 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { // one-sided search int p2update = 0; - while (p1.deriv[0]*dir <= -gtol && ctx->LSiter < m->opt.ls_iterations) { + while (p1.deriv[0]*dir <= -gtol && ctx->LSiter < ls_iterations) { // save current p2 = p1; p2update = 1; // move to Newton point w.r.t current p1.alpha -= p1.deriv[0]/p1.deriv[1]; - CGeval(m, d, ctx, &p1); + CGeval(ctx, &p1); // check for convergence if (mju_abs(p1.deriv[0]) < gtol) { @@ -1306,7 +1458,7 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { } // check for failure to bracket - if (ctx->LSiter >= m->opt.ls_iterations) { + if (ctx->LSiter >= ls_iterations) { ctx->LSresult = 3; // could not bracket ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; return p1.alpha; @@ -1322,13 +1474,13 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { // compute next-points for bracket p2next = p1; p1next.alpha = p1.alpha - p1.deriv[0]/p1.deriv[1]; - CGeval(m, d, ctx, &p1next); + CGeval(ctx, &p1next); // bracketed search - while (ctx->LSiter < m->opt.ls_iterations) { + while (ctx->LSiter < ls_iterations) { // evaluate at midpoint pmid.alpha = 0.5*(p1.alpha + p2.alpha); - CGeval(m, d, ctx, &pmid); + CGeval(ctx, &pmid); // make list of candidates mjCGPnt candidates[3] = {p1next, p2next, pmid}; @@ -1349,8 +1501,8 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { } // update brackets - int b1 = updateBracket(m, d, ctx, &p1, candidates, &p1next); - int b2 = updateBracket(m, d, ctx, &p2, candidates, &p2next); + int b1 = updateBracket(ctx, &p1, candidates, &p1next); + int b2 = updateBracket(ctx, &p2, candidates, &p2next); // no update possible: numerical accuracy reached, use midpoint if (!b1 && !b2) { @@ -1730,8 +1882,6 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, // local copies int nv = ctx.nv; int nefc = ctx.nefc; - const int* dofind = ctx.dofind; - const int* efcind = ctx.efcind; // allocate local storage if (!flg_Newton) { @@ -1741,27 +1891,32 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, } int* oldstate = mjSTACKALLOC(d, nefc, int); - // initialize matrix-vector products - int flg_vecunc = 1; // d->qacc is uncompressed - mj_mulM_island(m, d, ctx.Ma, d->qacc, island, flg_vecunc); - int flg_resunc = 0; // ctx.Jaref is compressed - mj_mulJacVec_island(m, d, ctx.Jaref, d->qacc, island, flg_resunc, flg_vecunc); - if (island < 0) { - mju_subFrom(ctx.Jaref, d->efc_aref, nefc); + // compute Ma = M * qacc (island or monolithic) + if (island >= 0) { + mju_mulSymVecSparse(ctx.Ma, ctx.qM, ctx.qacc, nv, + ctx.M_rownnz, ctx.M_rowadr, ctx.M_diagnum, ctx.M_colind); } else { - for (int c=0; c < nefc; c++) { - ctx.Jaref[c] -= d->efc_aref[efcind[c]]; - } + mj_mulM_impl(ctx.Ma, ctx.qacc, nv, ctx.qM, + ctx.dof_Madr, ctx.dof_parentid, ctx.M_diagnum); } + // compute Jaref = J * qacc - aref (dense or sparse) + if (!ctx.J_rownnz) { + mju_mulMatVec(ctx.Jaref, ctx.J, ctx.qacc, nefc, nv); + } else { + mju_mulMatVecSparse(ctx.Jaref, ctx.J, ctx.qacc, nefc, + ctx.J_rownnz, ctx.J_rowadr, ctx.J_colind, ctx.J_rowsuper); + } + mju_subFrom(ctx.Jaref, ctx.efc_aref, nefc); + // first update - CGupdateConstraint(m, d, &ctx); + CGupdateConstraint(&ctx); if (flg_Newton) { // compute and factorize Hessian MakeHessian(m, d, &ctx); FactorizeHessian(m, d, &ctx, /*flg_recompute=*/0); } - CGupdateGradient(m, d, &ctx); + CGupdateGradient(&ctx); // start both with preconditioned gradient mju_scl(ctx.search, ctx.Mgrad, -1, nv); @@ -1772,8 +1927,9 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv)); } else { mjtNum island_inertia = 0; - for (int c=0; c < nv; c++) { - island_inertia += d->qM[m->dof_Madr[dofind[c]]]; + for (int i=0; i < nv; i++) { + int* map2dof = d->map_idof2dof + d->island_idofadr[island]; + island_inertia += d->qM[m->dof_Madr[map2dof[i]]]; } scale = 1 / island_inertia; } @@ -1782,7 +1938,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, // main loop while (iter < maxiter) { // perform linesearch - alpha = CGsearch(m, d, &ctx); + alpha = CGsearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations); // no improvement: done if (alpha == 0) { @@ -1790,13 +1946,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, } // move to new solution - if (island < 0) { - mju_addToScl(d->qacc, ctx.search, alpha, nv); - } else { - for (int c=0; c < nv; c++) { - d->qacc[dofind[c]] += alpha * ctx.search[c]; - } - } + mju_addToScl(ctx.qacc, ctx.search, alpha, nv); mju_addToScl(ctx.Ma, ctx.Mv, alpha, nv); mju_addToScl(ctx.Jaref, ctx.Jv, alpha, nefc); @@ -1805,27 +1955,20 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, mju_copy(gradold, ctx.grad, nv); mju_copy(Mgradold, ctx.Mgrad, nv); } - if (island < 0) { - mju_copyInt(oldstate, d->efc_state, nefc); - } else { - for (int c=0; c < nefc; c++) { - oldstate[c] = d->efc_state[efcind[c]]; - } - } + mju_copyInt(oldstate, ctx.efc_state, nefc); mjtNum oldcost = ctx.cost; // update - CGupdateConstraint(m, d, &ctx); + CGupdateConstraint(&ctx); if (flg_Newton) { HessianIncremental(m, d, &ctx, oldstate); } - CGupdateGradient(m, d, &ctx); + CGupdateGradient(&ctx); // count state changes int nchange = 0; - for (int c=0; c < nefc; c++) { - int i = efcind ? efcind[c] : c; - nchange += (d->efc_state[i] != oldstate[c]); + for (int i=0; i < nefc; i++) { + nchange += (ctx.efc_state[i] != oldstate[i]); } // scale improvement, gradient, save stats @@ -1857,8 +2000,8 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, } // update - for (int c=0; c < nv; c++) { - ctx.search[c] = -ctx.Mgrad[c] + beta*ctx.search[c]; + for (int i=0; i < nv; i++) { + ctx.search[i] = -ctx.Mgrad[i] + beta*ctx.search[i]; } } } diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 2b3960e8..541cf31f 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -972,14 +972,9 @@ void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M) { -// multiply vector by inertia matrix -void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { - int nv = m->nv; - const mjtNum* M = d->qM; - const int* Madr = m->dof_Madr; - const int* parentid = m->dof_parentid; - const int* simplenum = m->dof_simplenum; - +// multiply vector by inertia matrix (implementation) +void mj_mulM_impl(mjtNum* res, const mjtNum* vec, int nv, const mjtNum* M, + const int* Madr, const int* parentid, const int* simplenum) { mju_zero(res, nv); for (int i=0; i < nv; i++) { @@ -1031,64 +1026,9 @@ void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) -// multiply vector by inertia matrix for one dof island -void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_vecunc) { - // if no island, call regular function - if (island < 0) { - mj_mulM(m, d, res, vec); - return; - } - - // local constants: general - const mjtNum* M = d->qM; - const int* Madr = m->dof_Madr; - const int* parentid = m->dof_parentid; - const int* simplenum = m->dof_simplenum; - - // local constants: island specific - int ndof = d->island_dofnum[island]; - const int* dofind = d->island_dofind + d->island_dofadr[island]; - const int* islandind = d->dof_islandind; - - mju_zero(res, ndof); - - for (int k=0; k < ndof; k++) { - // address in full dof vector - int i = dofind[k]; - - // address in M - int adr = Madr[i]; - - // diagonal - if (flg_vecunc) { - res[k] = M[adr]*vec[i]; - } else { - res[k] = M[adr]*vec[k]; - } - - // simple dof: continue - if (simplenum[i]) { - continue; - } - - // off-diagonal - int j = parentid[i]; - while (j >= 0) { - adr++; - int l = islandind[j]; - if (flg_vecunc) { - res[k] += M[adr]*vec[j]; - res[l] += M[adr]*vec[i]; - } else { - res[k] += M[adr]*vec[l]; - res[l] += M[adr]*vec[k]; - } - - // advance to parent - j = parentid[j]; - } - } +// multiply vector by inertia matrix +void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { + mj_mulM_impl(res, vec, m->nv, d->qM, m->dof_Madr, m->dof_parentid, m->dof_simplenum); } diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index c188ea38..220979da 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -120,13 +120,13 @@ MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body); // convert sparse inertia matrix M into full matrix MJAPI void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M); +// multiply vector by inertia matrix (implementation) +MJAPI void mj_mulM_impl(mjtNum* res, const mjtNum* vec, int nv, const mjtNum* M, + const int* Madr, const int* parentid, const int* simplenum); + // multiply vector by inertia matrix MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); -// multiply vector by inertia matrix for one dof island -MJAPI void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, - int island, int flg_vecunc); - // multiply vector by (inertia matrix)^(1/2) MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c index bec28aa3..b7f06601 100644 --- a/src/engine/engine_util_misc.c +++ b/src/engine/engine_util_misc.c @@ -1408,6 +1408,24 @@ void mju_scatter(mjtNum* restrict res, const mjtNum* restrict vec, const int* re +// gather integers +void mju_gatherInt(int* restrict res, const int* restrict vec, const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[i] = vec[ind[i]]; + } +} + + + +// scatter integers +void mju_scatterInt(int* restrict res, const int* restrict vec, const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[ind[i]] = vec[i]; + } +} + + + // insertion sort, increasing order void mju_insertionSort(mjtNum* list, int n) { for (int i=1; i < n; i++) { diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h index de10f907..f7f23782 100644 --- a/src/engine/engine_util_misc.h +++ b/src/engine/engine_util_misc.h @@ -156,12 +156,18 @@ MJAPI void mju_d2n(mjtNum* res, const double* vec, int n); // convert from mjtNum to double MJAPI void mju_n2d(double* res, const mjtNum* vec, int n); -// gather +// gather mjtNums MJAPI void mju_gather(mjtNum* res, const mjtNum* vec, const int* ind, int n); -// scatter +// scatter mjtNums MJAPI void mju_scatter(mjtNum* res, const mjtNum* vec, const int* ind, int n); +// gather integers +MJAPI void mju_gatherInt(int* res, const int* vec, const int* ind, int n); + +// scatter integers +MJAPI void mju_scatterInt(int* res, const int* vec, const int* ind, int n); + // insertion sort, increasing order MJAPI void mju_insertionSort(mjtNum* list, int n); diff --git a/src/engine/engine_vis_state.c b/src/engine/engine_vis_state.c index 9a0ea210..ff305f05 100644 --- a/src/engine/engine_vis_state.c +++ b/src/engine/engine_vis_state.c @@ -98,7 +98,6 @@ void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnsta // buffer space required for islands scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->island_dofadr) * m->ntree); - scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->island_dofind) * m->nv); scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->dof_island) * m->nv); scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->efc_island) * maxgeom * condimmax); scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->tendon_efcadr) * m->ntendon); @@ -136,9 +135,6 @@ void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnsta scnstate->data.island_dofadr = (int*)ptr; ptr += roundUpToCacheLine(sizeof(*scnstate->data.island_dofadr) * scnstate->model.ntree); - scnstate->data.island_dofind = (int*)ptr; - ptr += roundUpToCacheLine(sizeof(*scnstate->data.island_dofind) * scnstate->model.nv); - scnstate->data.dof_island = (int*)ptr; ptr += roundUpToCacheLine(sizeof(*scnstate->data.dof_island) * scnstate->model.nv); @@ -224,7 +220,6 @@ void mjv_assignFromSceneState(const mjvSceneState* scnstate, mjModel* m, mjData* d->contact = scnstate->data.contact; d->efc_force = scnstate->data.efc_force; d->island_dofadr = scnstate->data.island_dofadr; - d->island_dofind = scnstate->data.island_dofind; d->dof_island = scnstate->data.dof_island; d->efc_island = scnstate->data.efc_island; d->tendon_efcadr = scnstate->data.tendon_efcadr; @@ -385,7 +380,6 @@ void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt, scnstate->data.nisland = d->nisland; if (d->nisland) { memcpy(scnstate->data.island_dofadr, d->island_dofadr, sizeof(*d->island_dofadr) * d->nisland); - memcpy(scnstate->data.island_dofind, d->island_dofind, sizeof(*d->island_dofind) * m->nv); memcpy(scnstate->data.dof_island, d->dof_island, sizeof(*d->dof_island) * m->nv); memcpy(scnstate->data.tendon_efcadr, d->tendon_efcadr, sizeof(*d->tendon_efcadr) * m->ntendon); } diff --git a/src/engine/engine_vis_visualize.c b/src/engine/engine_vis_visualize.c index 1e665a2c..012cb966 100644 --- a/src/engine/engine_vis_visualize.c +++ b/src/engine/engine_vis_visualize.c @@ -91,9 +91,9 @@ static void makeLabel(const mjModel* m, mjtObj type, int id, char* label) { // assign pseudo-random rgba to constraint island using Halton sequence static void islandColor(float rgba[4], int islanddofadr) { - rgba[0] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 2); - rgba[1] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 3); - rgba[2] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 5); + rgba[0] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 2); + rgba[1] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 3); + rgba[2] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 5); rgba[3] = 1; } @@ -152,7 +152,7 @@ static void addContactGeom(const mjModel* m, mjData* d, const mjtByte* flags, // override standard colors if visualizing islands if (vopt->flags[mjVIS_ISLAND] && d->nisland && efc_adr >= 0) { // set color using island's first dof - islandColor(thisgeom->rgba, d->island_dofind[d->island_dofadr[d->efc_island[efc_adr]]]); + islandColor(thisgeom->rgba, d->island_dofadr[d->efc_island[efc_adr]]); } // otherwise regular colors (different for included and excluded contacts) @@ -1344,7 +1344,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, int island = d->dof_island[m->body_dofadr[weld_id]]; if (island > -1) { // color using island's first dof - islandColor(rgba_island, d->island_dofind[d->island_dofadr[island]]); + islandColor(rgba_island, d->island_dofadr[island]); } } } @@ -1835,7 +1835,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, if (d->tendon_efcadr[i] != -1) { // set color using island's first dof int island = d->efc_island[d->tendon_efcadr[i]]; - islandColor(rgba_island, d->island_dofind[d->island_dofadr[island]]); + islandColor(rgba_island, d->island_dofadr[island]); } } setMaterial(m, thisgeom, tendon_matid, rgba, vopt->flags); diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index a0a8f149..3704bdfc 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -25,6 +25,7 @@ #include #include "src/engine/engine_core_constraint.h" #include "src/engine/engine_support.h" +#include "src/engine/engine_util_misc.h" #include "test/fixture.h" namespace mujoco { @@ -284,205 +285,15 @@ TEST_F(CoreConstraintTest, EqualityBodySite) { mj_deleteModel(model); } - static const char* const kIlslandEfcPath = "engine/testdata/island/island_efc.xml"; -TEST_F(CoreConstraintTest, MulJacVecIsland) { +// validate mj_constraintUpdate_impl +TEST_F(CoreConstraintTest, ConstraintUpdateImpl) { const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data = mj_makeData(model); - - // allocate vec_nv, fill with arbitrary values - mjtNum* vec_nv = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); - for (int i=0; i < model->nv; i++) { - vec_nv[i] = 0.2 + 0.3*i; - } - - // iterate through dense and sparse - for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { - model->opt.jacobian = sparsity; - - // simulate for 0.2 seconds - mj_resetData(model, data); - while (data->time < 0.2) { - mj_step(model, data); - } - mj_forward(model, data); - - // multiply by Jacobian: vec_nefc = J * vec_nv - mjtNum* vec_nefc = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); - mj_mulJacVec(model, data, vec_nefc, vec_nv); - mjtNum* vec_nefc_tmp = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); - - // iterate over islands - for (int i=0; i < data->nisland; i++) { - // allocate dof and efc vectors for island - int dofnum = data->island_dofnum[i]; - mjtNum* vec_nvi = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - int efcnum = data->island_efcnum[i]; - mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); - - // get indices - int* dofind = data->island_dofind + data->island_dofadr[i]; - int* efcind = data->island_efcind + data->island_efcadr[i]; - - // copy values into vec_nvi - for (int j=0; j < dofnum; j++) { - vec_nvi[j] = vec_nv[dofind[j]]; - } - - // ===== both compressed - int flg_resunc = 0; - int flg_vecunc = 0; - mju_zero(vec_nefci, efcnum); // clear output - mj_mulJacVec_island(model, data, vec_nefci, vec_nvi, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < efcnum; j++) { - EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12)); - } - - // ===== input uncompressed: read from vec_nv - flg_resunc = 0; - flg_vecunc = 1; - mju_zero(vec_nefci, efcnum); // clear output - mj_mulJacVec_island(model, data, vec_nefci, vec_nv, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < efcnum; j++) { - EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12)); - } - - // ===== output uncompressed: write to vec_nefc_tmp - flg_resunc = 1; - flg_vecunc = 0; - mju_zero(vec_nefc_tmp, data->nefc); // clear output - mj_mulJacVec_island(model, data, vec_nefc_tmp, vec_nvi, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < efcnum; j++) { - EXPECT_THAT(vec_nefc_tmp[efcind[j]], - DoubleNear(vec_nefc[efcind[j]], 1e-12)); - } - - mju_free(vec_nvi); - mju_free(vec_nefci); - } - - mju_free(vec_nefc_tmp); - mju_free(vec_nefc); - } - - mju_free(vec_nv); - mj_deleteData(data); - mj_deleteModel(model); -} - -TEST_F(CoreConstraintTest, MulJacTVecIsland) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data = mj_makeData(model); - - // allocate vec_nv - mjtNum* vec_nv = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); - mjtNum* vec_nv_tmp = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); - - // iterate through dense and sparse - for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { - model->opt.jacobian = sparsity; - - // simulate for 0.3 seconds - mj_resetData(model, data); - while (data->time < 0.3) { - mj_step(model, data); - } - mj_forward(model, data); - - // allocate vec_nefc, fill with arbitrary values - mjtNum* vec_nefc = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); - for (int i=0; i < data->nefc; i++) { - vec_nefc[i] = 0.2 + 0.3*i; - } - - // multiply by Jacobian: vec_nv = J^T * vec_nefc - mj_mulJacTVec(model, data, vec_nv, vec_nefc); - - // iterate over islands - for (int i=0; i < data->nisland; i++) { - // allocate dof and efc vectors for island - int dofnum = data->island_dofnum[i]; - mjtNum* vec_nvi = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - int efcnum = data->island_efcnum[i]; - mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); - - // get indices - int* efcind = data->island_efcind + data->island_efcadr[i]; - int* dofind = data->island_dofind + data->island_dofadr[i]; - - // copy values into vec_nefci - for (int j=0; j < efcnum; j++) { - vec_nefci[j] = vec_nefc[efcind[j]]; - } - - // ==== both compressed - int flg_resunc = 0; - int flg_vecunc = 0; - mju_zero(vec_nvi, dofnum); // clear output - mj_mulJacTVec_island(model, data, vec_nvi, vec_nefci, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(vec_nvi[j], DoubleNear(vec_nv[dofind[j]], 1e-12)); - } - - // ===== input uncompressed: read from vec_nefc - flg_resunc = 0; - flg_vecunc = 1; - mju_zero(vec_nvi, dofnum); // clear output - mj_mulJacTVec_island(model, data, vec_nvi, vec_nefc, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(vec_nvi[j], DoubleNear(vec_nv[dofind[j]], 1e-12)); - } - - // ===== output uncompressed: write to vec_nv_tmp - flg_resunc = 1; - flg_vecunc = 0; - mju_zero(vec_nv_tmp, model->nv); // clear output - mj_mulJacTVec_island(model, data, vec_nv_tmp, vec_nefci, - i, flg_resunc, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(vec_nv_tmp[dofind[j]], - DoubleNear(vec_nv[dofind[j]], 1e-12)); - } - - mju_free(vec_nvi); - mju_free(vec_nefci); - } - mju_free(vec_nefc); - } - - mju_free(vec_nv_tmp); - mju_free(vec_nv); - mj_deleteData(data); - mj_deleteModel(model); -} - -// compare mj_constraintUpdate and mj_constraintUpdate_island -TEST_F(CoreConstraintTest, ConstraintUpdateIsland) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data1 = mj_makeData(model); - mjData* data2 = mj_makeData(model); + mjData* d1 = mj_makeData(model); + mjData* d2 = mj_makeData(model); // iterate over sparsity and cone for (mjtJacobian sparsity : {mjJAC_SPARSE, mjJAC_DENSE}) { @@ -491,81 +302,84 @@ TEST_F(CoreConstraintTest, ConstraintUpdateIsland) { model->opt.cone = cone; // simulate for 0.2 seconds - mj_resetData(model, data1); - mj_resetData(model, data2); - while (data1->time < 0.2) { - mj_step(model, data1); - mj_step(model, data2); + mj_resetData(model, d1); + mj_resetData(model, d2); + while (d1->time < 0.2) { + mj_step(model, d1); + mj_step(model, d2); } - mj_forward(model, data1); - mj_forward(model, data2); + mj_forward(model, d1); + mj_forward(model, d2); // get sizes - int nefc = data1->nefc; + int nefc = d1->nefc; int nv = model->nv; - int nisland = data1->nisland; + int nisland = d1->nisland; EXPECT_GT(nisland, 0); // get jar = J*a - aref mjtNum* jar = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc); - mj_mulJacVec(model, data1, jar, data1->qacc); - mju_subFrom(jar, data1->efc_aref, nefc); + mj_mulJacVec(model, d1, jar, d1->qacc); + mju_subFrom(jar, d1->efc_aref, nefc); // constraint update for data1 given jar mjtNum cost1; - mj_constraintUpdate(model, data1, jar, &cost1, /*flg_coneHessian=*/1); + mj_constraintUpdate(model, d1, jar, &cost1, /*flg_coneHessian=*/1); // iterate over islands, check match mjtNum cost2 = 0; for (int island=0; island < nisland; island++) { // clear outputs from data2 - for (int i=0; i < nefc; i++) data2->efc_state[i] = -1; - mju_zero(data2->efc_force, nefc); - mju_zero(data2->qfrc_constraint, nv); - for (int i=0; i < data2->ncon; i++) mju_zero(data2->contact[i].H, 36); + for (int i=0; i < nefc; i++) d2->efc_state[i] = -1; + mju_zero(d2->efc_force, nefc); + mju_zero(d2->qfrc_constraint, nv); + for (int i=0; i < d2->ncon; i++) mju_zero(d2->contact[i].H, 36); // sizes and indices, in this island - int dofnum = data2->island_dofnum[island]; - int efcnum = data2->island_efcnum[island]; - int* dofind = data2->island_dofind + data2->island_dofadr[island]; - int* efcind = data2->island_efcind + data2->island_efcadr[island]; + int efcnum = d2->island_nefc[island]; - // get jar restricted to island + // gather values into jari mjtNum* jari = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); - for (int c=0; c < efcnum; c++) { - jari[c] = jar[efcind[c]]; - } + int* map2efc = d2->map_iefc2efc + d2->island_iefcadr[island]; + mju_gather(jari, jar, map2efc, efcnum); // update constraints for this island mjtNum cost2i; - mj_constraintUpdate_island(model, data2, jari, &cost2i, - /*flg_coneHessian=*/1, island); + int ne = d2->island_ne[island]; + int nf = d2->island_nf[island]; + int adr = d2->island_iefcadr[island]; + int* state = d2->iefc_state + adr; + mjtNum *force = d2->iefc_force + adr; + mj_constraintUpdate_impl(ne, nf, efcnum, + d2->iefc_D + adr, + d2->iefc_R + adr, + d2->iefc_frictionloss + adr, + jari, + d2->iefc_type + adr, + d2->iefc_id + adr, + d2->contact, + state, + force, + &cost2i, + /*flg_coneHessian=*/1); // compare nefc vectors for (int c=0; c < efcnum; c++) { - int i = efcind[c]; - EXPECT_EQ(data2->efc_island[i], island); - EXPECT_EQ(data2->efc_state[i], data1->efc_state[i]); - EXPECT_THAT(data2->efc_force[i], - DoubleNear(data1->efc_force[i], 1e-12)); - } - - // compare qfrc_constraint - for (int c=0; c < dofnum; c++) { - int i = dofind[c]; - EXPECT_THAT(data2->qfrc_constraint[i], - DoubleNear(data1->qfrc_constraint[i], 1e-12)); + int i = map2efc[c]; + EXPECT_EQ(d2->efc_island[i], island); + EXPECT_EQ(state[c], d1->efc_state[i]); + EXPECT_THAT(force[c], DoubleNear(d1->efc_force[i], 1e-12)); } // compare cone Hessians if (cone == mjCONE_ELLIPTIC) { - for (int c=0; c < data2->ncon; c++) { - int efcadr = data2->contact[c].efc_address; - if (data2->efc_island[efcadr] == island && - data2->efc_state[efcadr] == mjCNSTRSTATE_CONE) { + for (int c=0; c < d2->ncon; c++) { + int efcadr = d2->contact[c].efc_address; + if (d2->efc_island[efcadr] == island && + d2->efc_state[efcadr] == mjCNSTRSTATE_CONE) { for (int j=0; j < 36; j++) { - EXPECT_THAT(data2->contact[c].H[j], - DoubleNear(data1->contact[c].H[j], 1e-12)); + EXPECT_THAT(d2->contact[c].H[j], + DoubleNear(d1->contact[c].H[j], 1e-12)); } } } @@ -584,8 +398,8 @@ TEST_F(CoreConstraintTest, ConstraintUpdateIsland) { } } - mj_deleteData(data2); - mj_deleteData(data1); + mj_deleteData(d2); + mj_deleteData(d1); mj_deleteModel(model); } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 1c71475a..7e4c113a 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -634,66 +634,6 @@ TEST_F(CoreSmoothTest, RefsiteConservesMomentum) { mj_deleteModel(model); } -static const char* const kIlslandEfcPath = - "engine/testdata/island/island_efc.xml"; -static const char* const kModelPath = - "testdata/model.xml"; - -TEST_F(CoreSmoothTest, SolveMIsland) { - for (auto model_path : {kModelPath, kIlslandEfcPath}) { - const std::string xml_path = GetTestDataFilePath(model_path); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data = mj_makeData(model); - int nv = model->nv; - - // allocate vec, fill with arbitrary values, copy to sol - mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv); - mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv); - for (int i=0; i < nv; i++) { - vec[i] = 0.2 + 0.3*i; - } - mju_copy(res, vec, nv); - - if (model->nkey > 0) mj_resetDataKeyframe(model, data, 0); - - for (int i=0; i < 6; i++) { - mj_step(model, data); - } - - mj_forward(model, data); - - // divide by mass matrix: sol = M^-1 * vec - mj_solveM(model, data, res, res, 1); - - // iterate over islands - for (int i=0; i < data->nisland; i++) { - // allocate dof vectors for island - int dofnum = data->island_dofnum[i]; - mjtNum* res_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - - // copy values into sol_i - int* dofind = data->island_dofind + data->island_dofadr[i]; - for (int j=0; j < dofnum; j++) { - res_i[j] = vec[dofind[j]]; - } - - // divide by mass matrix, for this island - mj_solveM_island(model, data, res_i, i); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-12)); - } - mju_free(res_i); - } - - mju_free(res); - mju_free(vec); - mj_deleteData(data); - mj_deleteModel(model); - } -} - static const char* const kInertiaPath = "engine/testdata/inertia.xml"; TEST_F(CoreSmoothTest, FactorI) { diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 3688c15b..e59ac3b4 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -208,17 +208,19 @@ TEST_F(IslandTest, Abacus) { int nv = model->nv; int nefc = data->nefc; int nisland = data->nisland; + int nidof = data->nidof; // 4 dofs, 12 constraints, 2 islands EXPECT_EQ(nv, 4); + EXPECT_EQ(nidof, 3); EXPECT_EQ(nefc, 12); // 3 pyramidal contacts EXPECT_EQ(nisland, 2); // the islands begin at dofs 0 and 1 - EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->island_idofadr, nisland), ElementsAre(0, 1)); // number of dofs in the 2 islands - EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(1, 2)); + EXPECT_THAT(AsVector(data->island_nv, nisland), ElementsAre(1, 2)); // dof 0 in island 0 // dof 1 in no island @@ -228,19 +230,19 @@ TEST_F(IslandTest, Abacus) { // dof 0 constitutes first island // dofs 2, 3 are the second island // last index is unassigned since dof 1 is unconstrained - EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 2, 3, -1)); + EXPECT_THAT(AsVector(data->map_idof2dof, nv), ElementsAre(0, 2, 3, 1)); // dof 0 constitutes first island // dofs 1 is unassigned // dofs 2, 3 are second island - EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, -1, 0, 1)); + EXPECT_THAT(AsVector(data->map_dof2idof, nv), ElementsAre(0, 3, 1, 2)); // island 0 starts at constraint 0 // island 1 starts at constraint 4 - EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4)); + EXPECT_THAT(AsVector(data->island_iefcadr, nisland), ElementsAre(0, 4)); // number of constraints in the 2 islands - EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(4, 8)); + EXPECT_THAT(AsVector(data->island_nefc, nisland), ElementsAre(4, 8)); // first contact (4 constraints) is in island 0 // second contact (8 constraints) is in island 1 @@ -248,7 +250,7 @@ TEST_F(IslandTest, Abacus) { ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1)); // index lists for islands 0 and 1 - EXPECT_THAT(AsVector(data->island_efcind, nefc), + EXPECT_THAT(AsVector(data->map_iefc2efc, nefc), ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)); // reset, push 0 to the left, 3 to the right, 1,2 to the middle @@ -266,18 +268,20 @@ TEST_F(IslandTest, Abacus) { // local variables nefc = data->nefc; nisland = data->nisland; + nidof = data->nidof; EXPECT_EQ(nisland, 3); - EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 1, 3)); - EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(1, 2, 1)); + EXPECT_EQ(nidof, 4); + EXPECT_THAT(AsVector(data->island_idofadr, nisland), ElementsAre(0, 1, 3)); + EXPECT_THAT(AsVector(data->island_nv, nisland), ElementsAre(1, 2, 1)); EXPECT_THAT(AsVector(data->dof_island, nv), ElementsAre(0, 1, 1, 2)); - EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 1, 2, 3)); - EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, 0, 1, 0)); - EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4, 8)); - EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(4, 4, 4)); + EXPECT_THAT(AsVector(data->map_idof2dof, nv), ElementsAre(0, 1, 2, 3)); + EXPECT_THAT(AsVector(data->map_dof2idof, nv), ElementsAre(0, 1, 2, 3)); + EXPECT_THAT(AsVector(data->island_iefcadr, nisland), ElementsAre(0, 4, 8)); + EXPECT_THAT(AsVector(data->island_nefc, nisland), ElementsAre(4, 4, 4)); EXPECT_THAT(AsVector(data->efc_island, nefc), ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2)); - EXPECT_THAT(AsVector(data->island_efcind, nefc), + EXPECT_THAT(AsVector(data->map_iefc2efc, nefc), ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)); mj_deleteData(data); @@ -311,27 +315,30 @@ TEST_F(IslandTest, DenseSparse) { int nisland = data1->nisland; // expect sparse and dense to be identical + EXPECT_EQ(data1->nidof, data2->nidof); EXPECT_EQ(data1->nefc, data2->nefc); EXPECT_EQ(data1->nisland, data2->nisland); EXPECT_EQ(data1->nefc, data2->nefc); - EXPECT_EQ(AsVector(data1->island_dofadr, nisland), - AsVector(data2->island_dofadr, nisland)); - EXPECT_EQ(AsVector(data1->island_dofnum, nisland), - AsVector(data2->island_dofnum, nisland)); + EXPECT_EQ(AsVector(data1->island_idofadr, nisland), + AsVector(data2->island_idofadr, nisland)); + EXPECT_EQ(AsVector(data1->island_nv, nisland), + AsVector(data2->island_nv, nisland)); EXPECT_EQ(AsVector(data1->dof_island, nv), AsVector(data2->dof_island, nv)); - EXPECT_EQ(AsVector(data1->island_dofind, nv), - AsVector(data2->island_dofind, nv)); - EXPECT_EQ(AsVector(data1->dof_islandind, nv), - AsVector(data2->dof_islandind, nv)); - EXPECT_EQ(AsVector(data1->island_efcadr, nisland), - AsVector(data2->island_efcadr, nisland)); - EXPECT_EQ(AsVector(data1->island_efcnum, nisland), - AsVector(data2->island_efcnum, nisland)); + EXPECT_EQ(AsVector(data1->map_idof2dof, nv), + AsVector(data2->map_idof2dof, nv)); + EXPECT_EQ(AsVector(data1->map_dof2idof, nv), + AsVector(data2->map_dof2idof, nv)); + EXPECT_EQ(AsVector(data1->island_iefcadr, nisland), + AsVector(data2->island_iefcadr, nisland)); + EXPECT_EQ(AsVector(data1->island_nefc, nisland), + AsVector(data2->island_nefc, nisland)); EXPECT_EQ(AsVector(data1->efc_island, nefc), AsVector(data2->efc_island, nefc)); - EXPECT_EQ(AsVector(data1->island_efcind, nefc), - AsVector(data2->island_efcind, nefc)); + EXPECT_EQ(AsVector(data1->map_iefc2efc, nefc), + AsVector(data2->map_iefc2efc, nefc)); + EXPECT_EQ(AsVector(data1->map_efc2iefc, nefc), + AsVector(data2->map_efc2iefc, nefc)); mj_deleteData(data2); mj_deleteData(data1); @@ -361,6 +368,156 @@ TEST_F(IslandTest, IslandEfc) { mj_deleteModel(model); } +static const char* const k2H100Path = "engine/testdata/island/2humanoid100.xml"; + +TEST_F(IslandTest, IslandJacobian) { + for (const char* local_path : {kIlslandEfcPath, k2H100Path}) { + const std::string xml_path = GetTestDataFilePath(local_path); + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + int jac0 = m->opt.jacobian; + mjData* d = mj_makeData(m); + + for (mjtNum t_stop : {0.0, 0.2, 2.0}) { + while (d->time < t_stop) { + mj_step(m, d); + } + + for (mjtJacobian jac : {mjJAC_DENSE, mjJAC_SPARSE}) { + m->opt.jacobian = jac; + mj_forward(m, d); + + int nv = m->nv; + int nefc = d->nefc; + int nisland = d->nisland; + int nidof = d->nidof; + + mjtNum* J = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc * nv); + mjtNum* iJ = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc * nidof); + + // get local dense Jacobian + if (jac == mjJAC_DENSE) { + mju_copy(J, d->efc_J, nefc * nv); + mju_copy(iJ, d->iefc_J, nefc * nidof); + } else { + mju_sparse2dense(J, d->efc_J, nefc, nv, d->efc_J_rownnz, + d->efc_J_rowadr, d->efc_J_colind); + } + + // compare random access in efc_J to contiguous memory in iefc_J + for (int island=0; island < nisland; island++) { + int idof = d->island_idofadr[island]; + int iefc = d->island_iefcadr[island]; + int nefc_island = d->island_nefc[island]; + int nv_island = d->island_nv[island]; + + // === test J + + // get pointer to J_island, dense (nefc_island x nv_island) submatrix + mjtNum* J_island; + if (jac == mjJAC_DENSE) { + // point to starting address of island in efc_J + J_island = iJ + iefc * nidof; + } else { + // dense copy of island in iJ (here used as scratch) + mju_sparse2dense(iJ, d->iefc_J, nefc_island, nv_island, + d->iefc_J_rownnz + iefc, + d->iefc_J_rowadr + iefc, + d->iefc_J_colind); + J_island = iJ; + } + + // sequential memory in J_island equals random access memory in J + for (int i=0; i < nefc_island; i++) { + for (int j=0; j < nv_island; j++) { + int efc = d->map_iefc2efc[iefc + i]; + int dof = d->map_idof2dof[idof + j]; + EXPECT_EQ(J_island[i * nv_island + j], J[efc * nv + dof]); + } + } + + // === test JT (if sparse) + + // get pointer to J_island, dense (nefc_island x nv_island) submatrix + if (jac == mjJAC_SPARSE) { + // dense copy of island in iJ (here used as scratch) + mju_sparse2dense(iJ, d->iefc_JT, nv_island, nefc_island, + d->iefc_JT_rownnz + idof, + d->iefc_JT_rowadr + idof, + d->iefc_JT_colind); + J_island = iJ; + + // sequential memory in J_island equals random access memory in J + for (int i=0; i < nv_island; i++) { + for (int j=0; j < nefc_island; j++) { + int dof = d->map_idof2dof[idof + i]; + int efc = d->map_iefc2efc[iefc + j]; + EXPECT_EQ(J_island[i * nefc_island + j], J[efc * nv + dof]); + } + } + } + } + + mju_free(iJ); + mju_free(J); + } + + // reset opt.jacobian to initial value + m->opt.jacobian = jac0; + } + + mj_deleteData(d); + mj_deleteModel(m); + } +} + +TEST_F(IslandTest, IslandInertia) { + for (const char* local_path : {kIlslandEfcPath, k2H100Path}) { + const std::string xml_path = GetTestDataFilePath(local_path); + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + int nv = m->nv; + mjData* d = mj_makeData(m); + mjtNum* M = (mjtNum*)mju_malloc(sizeof(mjtNum) * nv * nv); + + for (mjtNum t_stop : {0.0, 0.2, 2.0}) { + while (d->time < t_stop) { + mj_step(m, d); + } + mj_forward(m, d); + + int nisland = d->nisland; + + // get dense inertia (lower only) + mj_fullM(m, M, d->qM); + + // compare iM sub-matrix to full M + for (int island=0; island < nisland; island++) { + int nvi = d->island_nv[island]; + mjtNum* Mi = (mjtNum*)mju_malloc(sizeof(mjtNum) * nvi * nvi); + + int adr = d->island_idofadr[island]; + mju_sparse2dense(Mi, d->iM, nvi, nvi, + d->iM_rownnz + adr, + d->iM_rowadr + adr, + d->iM_colind); + + // compare Mi to M (lower triangle only) + for (int i=0; i < nvi; i++) { + for (int j=0; j <= i; j++) { + int dofi = d->map_idof2dof[adr + j]; + int dofj = d->map_idof2dof[adr + i]; + EXPECT_EQ(Mi[i * nvi + j], M[dofi * nv + dofj]); + } + } + mju_free(Mi); + } + } + + mju_free(M); + mj_deleteData(d); + mj_deleteModel(m); + } +} + TEST_F(IslandTest, IslandEfcElliptic) { const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); diff --git a/test/engine/engine_solver_test.cc b/test/engine/engine_solver_test.cc index b5b2db6d..75b42cf7 100644 --- a/test/engine/engine_solver_test.cc +++ b/test/engine/engine_solver_test.cc @@ -17,7 +17,6 @@ #include #include #include -#include #include #include @@ -29,19 +28,9 @@ namespace { using ::testing::DoubleNear; using ::testing::NotNull; -using ::std::vector; using ::std::abs; using ::std::max; -// compare two vectors, relative error (increase tolerance for large elements) -inline void ExpectEqRel(vector v1, vector v2, mjtNum rtol) { - ASSERT_TRUE(v1.size() == v2.size()); - for (int i = 0; i < v1.size(); i++) { - mjtNum scale = 0.5 * max(2.0, abs(v1[i]) + abs(v2[i])); - EXPECT_THAT(v1[i], DoubleNear(v2[i], scale*rtol)); - } -} - using SolverTest = MujocoTest; static const char* const kModelPath = @@ -169,79 +158,5 @@ TEST_F(SolverTest, IslandsEquivalentForward) { mj_deleteModel(model); } -static const char* const kIlslandEfcPath = - "engine/testdata/island/island_efc.xml"; - -// compare qacc from 1 iteration of monolithic CG solver and one big island -TEST_F(SolverTest, OneBigIsland) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - ASSERT_THAT(model, NotNull()); - model->opt.solver = mjSOL_CG; // use CG solver - model->opt.disableflags |= mjDSBL_WARMSTART; // disable warmstart - model->opt.tolerance = 0; // set tolerance to 0 - model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands - - int state_size = mj_stateSize(model, mjSTATE_INTEGRATION); - mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size); - - mjData* data_island = mj_makeData(model); - mjData* data_noisland = mj_makeData(model); - - int nv = model->nv; - mjtNum rtol = 1e-7; - - // save current (default) iterations - int iterations_default = model->opt.iterations; - - while (data_noisland->time < .2) { - // step and copy the state to data_island - mj_step(model, data_noisland); - mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION); - mj_setState(model, data_island, state, mjSTATE_INTEGRATION); - - // set small number of iterations - model->opt.iterations = 1; - - // call forward on data_noisland - mj_forward(model, data_noisland); - - // enable islands - model->opt.enableflags |= mjENBL_ISLAND; - - // call forward (just for smooth dynamics and to allocate islands) - mj_forward(model, data_island); - - // overwrite island structure with one big island - data_island->nisland = 1; - data_island->island_dofnum[0] = nv; - data_island->island_dofadr[0] = 0; - for (int i = 0; i < nv; i++) { - data_island->island_dofind[i] = data_island->dof_islandind[i] = i; - } - int nefc = data_island->nefc; - data_island->island_efcnum[0] = nefc; - data_island->island_efcadr[0] = 0; - for (int i = 0; i < nefc; i++) data_island->island_efcind[i] = i; - - // solve using using one big island - mj_fwdConstraint(model, data_island); - - // re-disable islands and reset iterations - model->opt.enableflags &= ~mjENBL_ISLAND; - model->opt.iterations = iterations_default; - - // compare accelerations (relative error) - ExpectEqRel(AsVector(data_noisland->qacc, nv), - AsVector(data_island->qacc, nv), rtol); - } - - mj_deleteData(data_noisland); - mj_deleteData(data_island); - mju_free(state); - mj_deleteModel(model); -} - - } // namespace } // namespace mujoco diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index c7d4f018..71b8e820 100644 --- a/test/engine/engine_support_test.cc +++ b/test/engine/engine_support_test.cc @@ -830,77 +830,6 @@ TEST_F(InertiaTest, mulM2) { mj_deleteModel(model); } -static const char* const kIlslandEfcPath = - "engine/testdata/island/island_efc.xml"; - -TEST_F(SupportTest, MulMIsland) { - const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); - mjData* data = mj_makeData(model); - - // allocate vec, fill with arbitrary values - mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); - for (int i=0; i < model->nv; i++) { - vec[i] = 0.2 + 0.3*i; - } - - // simulate for 0.2 seconds - mj_resetData(model, data); - while (data->time < 0.2) { - mj_step(model, data); - } - mj_forward(model, data); - - // multiply by Mass matrix: Mvec = M * vec - mjtNum* Mvec = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); - mj_mulM(model, data, Mvec, vec); - - // iterate over islands - for (int i=0; i < data->nisland; i++) { - // allocate dof vectors for island - int dofnum = data->island_dofnum[i]; - mjtNum* vec_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - mjtNum* Mvec_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum); - - // copy values into vec_i - int* dofind = data->island_dofind + data->island_dofadr[i]; - for (int j=0; j < dofnum; j++) { - vec_i[j] = vec[dofind[j]]; - } - - // === compressed: use vec_i - - // multiply by Jacobian, for this island - int flg_vecunc = 0; - mj_mulM_island(model, data, Mvec_i, vec_i, i, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(Mvec_i[j], DoubleNear(Mvec[dofind[j]], 1e-12)); - } - - // === uncompressed: use vec - mju_zero(Mvec_i, dofnum); // clear output - - // multiply by Jacobian, for this island - flg_vecunc = 1; - mj_mulM_island(model, data, Mvec_i, vec, i, flg_vecunc); - - // expect corresponding values to match - for (int j=0; j < dofnum; j++) { - EXPECT_THAT(Mvec_i[j], DoubleNear(Mvec[dofind[j]], 1e-12)); - } - - mju_free(vec_i); - mju_free(Mvec_i); - } - - mju_free(Mvec); - mju_free(vec); - mj_deleteData(data); - mj_deleteModel(model); -} - static constexpr char GeomDistanceTestingModel[] = R"(