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

Add mjtState enum and related mj_stateSize, mj_getState and mj_setState functions.

Browse Source 
PiperOrigin-RevId: 539667943
Change-Id: I14c34be8ce4c287e529380257b5c5bffc3ab45ce
This commit is contained in:
Yuval Tassa
2023-06-12 08:55:58 -07:00
committed by Copybara-Service
parent 84c33e53df
commit f67e359532
16 changed files with 520 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/APItypes.rst
+12
View File
@@ -346,6 +346,18 @@ Data
The enums below are defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`_.
.. _mjtState:
mjtState
~~~~~~~~
State component elements as integer bitflags and several convenient combinations of these flags. Used by
:ref:`mj_getState`, :ref:`mj_setState` and :ref:`mj_stateSize`.
.. mujoco-include:: mjtState
.. _mjtWarning:
mjtWarning
doc/APIreference/functions.rst
+30
View File
@@ -151,6 +151,36 @@ These are support functions that need access to :ref:`mjModel` and :ref:`mjData`
not need such access. Support functions are called within the simulator but some of them can also be useful for custom
computations, and are documented in more detail below.
.. _mj_stateSize:
mj_stateSize
~~~~~~~~~~~~
.. mujoco-include:: mj_stateSize
Returns the number of :ref:`mjtNum` |-| s required for a given state specification. The bits of the integer ``spec``
correspond to element fields of :ref:`mjtState`.
.. _mj_getState:
mj_getState
~~~~~~~~~~~
.. mujoco-include:: mj_getState
Copy concatenated state components specified by ``spec`` from ``d`` into ``state``. The bits of the integer
``spec`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``spec`` is invalid.
.. _mj_setState:
mj_setState
~~~~~~~~~~~
.. mujoco-include:: mj_setState
Copy concatenated state components specified by ``spec`` from ``state`` into ``d``. The bits of the integer
``spec`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``spec`` is invalid.
.. _mj_addContact:
mj_addContact
doc/APIreference/functions_override.rst
+15
View File
@@ -96,6 +96,21 @@ These are support functions that need access to :ref:`mjModel` and :ref:`mjData`
not need such access. Support functions are called within the simulator but some of them can also be useful for custom
computations, and are documented in more detail below.
.. _mj_stateSize:
Returns the number of :ref:`mjtNum` |-| s required for a given state specification. The bits of the integer ``spec``
correspond to element fields of :ref:`mjtState`.
.. _mj_getState:
Copy concatenated state components specified by ``spec`` from ``d`` into ``state``. The bits of the integer
``spec`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``spec`` is invalid.
.. _mj_setState:
Copy concatenated state components specified by ``spec`` from ``state`` into ``d``. The bits of the integer
``spec`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``spec`` is invalid.
.. _mj_mulJacVec:
This function multiplies the constraint Jacobian mjData.efc_J by a vector. Note that the Jacobian can be either dense or
doc/changelog.rst
+2
View File
@@ -63,6 +63,8 @@ Simulate
General
^^^^^^^
- Added :ref:`mj_getState` and :ref:`mj_setState` for getting and setting the simulation state as a concatenated vector
of floating point numbers. See the :ref:`State<geState>` section for details.
- Added :ref:`mjContact.solreffriction<mjContact>`, allowing different :ref:`solref<CSolver>` parameters for the normal
and frictional axes of contacts when using :ref:`elliptic friction cones<option-cone>`. This attribute is required
for elastic frictional collisions, see associated
doc/computation.rst
+49 -37
View File
@@ -569,41 +569,60 @@ Fast implicit-in-velocity (``implicitfast``)
.. _geState:
The **state**
~~~~~~~~~~~~~
The State
~~~~~~~~~
To complete our description of the general framework we will now discuss the notion of *state*. MuJoCo has a compact,
well-defined internal state which, together with the deterministic computational pipeline, means that operations like
resetting the state and computing dynamics derivatives are also well-defined. The state is entirely encapsulated in the
``mjData`` struct and consists of several components:
well-defined internal state which, together with the :ref:`deterministic computational pipeline<piReproducibility>`,
means that operations like resetting the state and computing dynamics derivatives are also well-defined.
The state is entirely encapsulated in the :ref:`mjData` struct and consists of several components. The components are
enumerated in :ref:`mjtState` as bit flags, along with several common combinations, corresponding to the groupings
below. Concatenated state vectors can be conveniently read from and written into :ref:`mjData` using :ref:`mj_getState`
and :ref:`mj_setState`, respectively.
.. _gePhysicsState:
Physics state
^^^^^^^^^^^^^
The *physics state* contains all quantities which are time-integrated during stepping.
These are ``mjData.{qpos, qvel, act, time}``:
The *physics state* (:ref:`mjSTATE_PHYSICS<mjtState>`) contains the main quantities which are time-integrated during
stepping. These are ``mjData.{qpos, qvel, act}``:
Mechanical state: ``qpos`` and ``qvel``
The *mechanical state* of a simulation is given by the generalized position (``mjData.qpos``) and velocity
(``mjData.qvel``) vectors, denoted above as :math:`q` and :math:`v`, respectively.
Position: ``qpos``
The configuration in generalized coodinates, denoted above as :math:`q`.
Actuator activations: ``act``
Velocity: ``qvel``
The generalized velocities, denoted above as :math:`v`.
Actuator activation: ``act``
``mjData.act`` contains the internal states of stateful actuators, denoted above as :math:`w`.
.. _geFullPhysics:
Full physics state
^^^^^^^^^^^^^^^^^^
The *full physics state* (:ref:`mjSTATE_FULLPHYSICS<mjtState>`) contains the physics state and two additional
components:
Time: ``time``
The time of the simulation is given by the scalar ``mjData.time``. Since physics is time-invariant, it is
often excluded from the *physics state*; an exception could be a time-dependent user callback (e.g., an open-loop
The simulation time is given by the scalar ``mjData.time``. Since physics is time-invariant, it is
excluded from the *physics state*; exceptions include time-dependent user callbacks and plugins (e.g., an open-loop
controller), in which case time should be included.
Plugin state: ``plugin_state``
``mjData.plugin_state`` are states declared by :ref:`engine plugins<exPlugin>`. Please see the :ref:`exPluginState`
section for more details.
.. _geInput:
User inputs
^^^^^^^^^^^
These input fields are set by the user and affect the physics simulation, but are untouched by the simulator. All input
fields except for MoCap poses default to 0.
Controls: ``ctrl``
These input fields (:ref:`mjSTATE_USER<mjtState>`) are set by the user and affect the physics simulation, but are
untouched by the simulator. All input fields except for MoCap poses default to 0.
Control: ``ctrl``
Controls are defined by the :ref:`actuator<actuator>` section of the XML. ``mjData.ctrl`` values either produce
generalized forces directly (stateless actuators), or affect the actuator activations in ``mjData.act``, which then
produce forces.
@@ -627,8 +646,8 @@ User data: ``userdata``
.. _geWarmstart:
Warmstart accelerations
^^^^^^^^^^^^^^^^^^^^^^^
Warmstart acceleration
^^^^^^^^^^^^^^^^^^^^^^
``qacc_warmstart``
``mjData.qacc_warmstart`` are accelerations used to warmstart the constraint solver, saved from the previous step.
@@ -642,34 +661,27 @@ Warmstart accelerations
<https://en.wikipedia.org/wiki/Lyapunov_exponent>`__ when time-stepping, quickly leading to divergent trajectories
for different warmstarts.
.. _gePlugin:
Plugin state
^^^^^^^^^^^^
``plugin_state``
``mjData.plugin_state`` are states declared by :ref:`engine plugins<exPlugin>`. Please see the :ref:`exPluginState`
section for more details.
.. _geIntegrationState:
Integration state
^^^^^^^^^^^^^^^^^
The *integration state* is the union of all the above ``mjData`` fields and constitutes the entire set of inputs to
the *forward dynamics*. In the case of *inverse dynamics*, ``mjData.qacc`` is also treated as an input variable. All
other ``mjData`` fields are functions of the integration state.
|br| When saving the integration state in order to reload it elsewhere, it is sensible to avoid saving unused fields
that always remain in their default values. Specifically, ``xfrc_applied`` can be quite large (``6 x nbody``) yet is
often unused.
The *integration state* (:ref:`mjSTATE_INTEGRATION<mjtState>`) is the union of all the above :ref:`mjData` fields and
constitutes the entire set of inputs to the *forward dynamics*. In the case of *inverse dynamics*, ``mjData.qacc`` is
also treated as an input variable. All other :ref:`mjData` fields are functions of the integration state.
Note that the full integration state as given by :ref:`mjSTATE_INTEGRATION<mjtState>` is maximalist and includes fields
which are often unused. If a small state size is desired, it might be sensible to avoid saving unused fields.
In particular `xfrc_applied`` can be quite large (``6 x nbody``) yet is often unused.
.. _geSimulationState:
Simulation state: ``mjData``
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
The *simulation state* is the entirety of the ``mjData`` struct and associated memory buffer. This state includes
all derived quantities computed during dynamics computation. Because the ``mjData`` buffers are preallocated for the
The *simulation state* is the entirety of the :ref:`mjData` struct and associated memory buffer. This state includes
all derived quantities computed during dynamics computation. Because the :ref:`mjData` buffers are preallocated for the
worst case, it is often significantly faster to recompute derived quantities from the *integration state* rather than
using ``mj_copyData``.
using :ref:`mj_copyData`.
.. _Constraint:
@@ -1591,7 +1603,7 @@ MuJoCo's simulation pipeline is entirely deterministic and reproducible -- if a
saved and reloaded and :ref:`mj_step` called again, the resulting next state will be identical. However, there are some
important caveats:
- Save all the required :ref:`integration state<IntegrationState>` components. In particular :ref:`warmstart
- Save all the required :ref:`integration state<geIntegrationState>` components. In particular :ref:`warmstart
accelerations<geWarmstart>` have only a very small effect on the next state, but should be saved if bit-wise equality
is required.
- Any numerical difference between states, no matter how small, will become significant upon integration, especially for
doc/includes/references.h
+26
View File
@@ -16,6 +16,29 @@
// Error: C reference not found
// NOLINTBEGIN
typedef enum mjtState_ { // state elements
mjSTATE_TIME = 1<<0, // time
mjSTATE_QPOS = 1<<1, // position
mjSTATE_QVEL = 1<<2, // velocity
mjSTATE_ACT = 1<<3, // actuator activation
mjSTATE_WARMSTART = 1<<4, // acceleration used for warmstart
mjSTATE_CTRL = 1<<5, // control
mjSTATE_QFRC_APPLIED = 1<<6, // applied generalized force
mjSTATE_XFRC_APPLIED = 1<<7, // applied Cartesian force/torque
mjSTATE_MOCAP_POS = 1<<8, // positions of mocap bodies
mjSTATE_MOCAP_QUAT = 1<<9, // orientations of mocap bodies
mjSTATE_USERDATA = 1<<10, // user data
mjSTATE_PLUGIN = 1<<11, // plugin state
mjNSTATE = 12, // number of state elements
// convenience values for commonly used state specifications
mjSTATE_PHYSICS = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT,
mjSTATE_FULLPHYSICS = mjSTATE_PHYSICS | mjSTATE_TIME | mjSTATE_PLUGIN,
mjSTATE_USER = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED |
mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT | mjSTATE_USERDATA,
mjSTATE_INTEGRATION = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART
} mjtState;
typedef enum mjtWarning_ { // warning types
mjWARN_INERTIA = 0, // (near) singular inertia matrix
mjWARN_CONTACTFULL, // too many contacts in contact list
@@ -2176,6 +2199,9 @@ void mj_projectConstraint(const mjModel* m, mjData* d);
void mj_referenceConstraint(const mjModel* m, mjData* d);
void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
mjtNum cost[1], int flg_coneHessian);
int mj_stateSize(const mjModel* m, unsigned int spec);
void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int spec);
void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int spec);
int mj_addContact(const mjModel* m, mjData* d, const mjContact* con);
int mj_isPyramidal(const mjModel* m);
int mj_isSparse(const mjModel* m);
include/mujoco/mjdata.h
+25
View File
@@ -23,6 +23,31 @@
//---------------------------------- primitive types (mjt) -----------------------------------------
typedef enum mjtState_ { // state elements
mjSTATE_TIME = 1<<0, // time
mjSTATE_QPOS = 1<<1, // position
mjSTATE_QVEL = 1<<2, // velocity
mjSTATE_ACT = 1<<3, // actuator activation
mjSTATE_WARMSTART = 1<<4, // acceleration used for warmstart
mjSTATE_CTRL = 1<<5, // control
mjSTATE_QFRC_APPLIED = 1<<6, // applied generalized force
mjSTATE_XFRC_APPLIED = 1<<7, // applied Cartesian force/torque
mjSTATE_MOCAP_POS = 1<<8, // positions of mocap bodies
mjSTATE_MOCAP_QUAT = 1<<9, // orientations of mocap bodies
mjSTATE_USERDATA = 1<<10, // user data
mjSTATE_PLUGIN = 1<<11, // plugin state
mjNSTATE = 12, // number of state elements
// convenience values for commonly used state specifications
mjSTATE_PHYSICS = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT,
mjSTATE_FULLPHYSICS = mjSTATE_PHYSICS | mjSTATE_TIME | mjSTATE_PLUGIN,
mjSTATE_USER = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED |
mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT | mjSTATE_USERDATA,
mjSTATE_INTEGRATION = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART
} mjtState;
typedef enum mjtWarning_ { // warning types
mjWARN_INERTIA = 0, // (near) singular inertia matrix
mjWARN_CONTACTFULL, // too many contacts in contact list
include/mujoco/mujoco.h
+9
View File
@@ -357,6 +357,15 @@ MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
//---------------------------------- Support -------------------------------------------------------
// Return size of state specification.
MJAPI int mj_stateSize(const mjModel* m, unsigned int spec);
// Get state.
MJAPI void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int spec);
// Set state.
MJAPI void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int spec);
// Add contact to d->contact list; return 0 if success; 1 if buffer full.
MJAPI int mj_addContact(const mjModel* m, mjData* d, const mjContact* con);
introspect/enums.py
+24
View File
@@ -373,6 +373,30 @@ ENUMS: Mapping[str, EnumDecl] = dict([
('mjLRMODE_ALL', 3),
]),
)),
('mjtState',
EnumDecl(
name='mjtState',
declname='enum mjtState_',
values=dict([
('mjSTATE_TIME', 1),
('mjSTATE_QPOS', 2),
('mjSTATE_QVEL', 4),
('mjSTATE_ACT', 8),
('mjSTATE_WARMSTART', 16),
('mjSTATE_CTRL', 32),
('mjSTATE_QFRC_APPLIED', 64),
('mjSTATE_XFRC_APPLIED', 128),
('mjSTATE_MOCAP_POS', 256),
('mjSTATE_MOCAP_QUAT', 512),
('mjSTATE_USERDATA', 1024),
('mjSTATE_PLUGIN', 2048),
('mjNSTATE', 12),
('mjSTATE_PHYSICS', 14),
('mjSTATE_FULLPHYSICS', 2063),
('mjSTATE_USER', 2016),
('mjSTATE_INTEGRATION', 4095),
]),
)),
('mjtWarning',
EnumDecl(
name='mjtWarning',
introspect/functions.py
+78
View File
@@ -1783,6 +1783,84 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians. If cost is not NULL, set *cost = s(jar) where jar = Jac*qacc-aref.', # pylint: disable=line-too-long
)),
('mj_stateSize',
FunctionDecl(
name='mj_stateSize',
return_type=ValueType(name='int'),
parameters=(
FunctionParameterDecl(
name='m',
type=PointerType(
inner_type=ValueType(name='mjModel', is_const=True),
),
),
FunctionParameterDecl(
name='spec',
type=ValueType(name='unsigned int'),
),
),
doc='Return size of state specification.',
)),
('mj_getState',
FunctionDecl(
name='mj_getState',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='m',
type=PointerType(
inner_type=ValueType(name='mjModel', is_const=True),
),
),
FunctionParameterDecl(
name='d',
type=PointerType(
inner_type=ValueType(name='mjData', is_const=True),
),
),
FunctionParameterDecl(
name='state',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='spec',
type=ValueType(name='unsigned int'),
),
),
doc='Get state.',
)),
('mj_setState',
FunctionDecl(
name='mj_setState',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='m',
type=PointerType(
inner_type=ValueType(name='mjModel', is_const=True),
),
),
FunctionParameterDecl(
name='d',
type=PointerType(
inner_type=ValueType(name='mjData'),
),
),
FunctionParameterDecl(
name='state',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='spec',
type=ValueType(name='unsigned int'),
),
),
doc='Set state.',
)),
('mj_addContact',
FunctionDecl(
name='mj_addContact',
python/mujoco/bindings_test.py
+37
View File
@@ -667,6 +667,43 @@ class MuJoCoBindingsTest(parameterized.TestCase):
with self.assertRaises(TypeError):
mujoco.mju_rotVecQuat(vec, quat, res=np.zeros(3, int))
def test_getsetstate(self): # pylint: disable=invalid-name
mujoco.mj_step(self.model, self.data)
# Test for invalid state spec
invalid_spec = 2**mujoco.mjtState.mjNSTATE.value
expected_message = (
f'mj_stateSize: invalid state spec {invalid_spec} >= 2^mjNSTATE'
)
with self.assertRaisesWithLiteralMatch(mujoco.FatalError, expected_message):
mujoco.mj_stateSize(self.model, invalid_spec)
spec = mujoco.mjtState.mjSTATE_INTEGRATION
size = mujoco.mj_stateSize(self.model, spec)
state_bad_size = np.empty(size + 1, np.float64)
expected_message = ('state size should equal mj_stateSize(m, spec)')
with self.assertRaisesWithLiteralMatch(TypeError, expected_message):
mujoco.mj_getState(self.model, self.data, state_bad_size, spec)
# Get initial state.
state0 = np.empty(size, np.float64)
mujoco.mj_getState(self.model, self.data, state0, spec)
# Step, get next state.
mujoco.mj_step(self.model, self.data)
state1a = np.empty(size, np.float64)
mujoco.mj_getState(self.model, self.data, state1a, spec)
# Reset to initial state, step again, get state again.
mujoco.mj_setState(self.model, self.data, state0, spec)
mujoco.mj_step(self.model, self.data)
state1b = np.empty(size, np.float64)
mujoco.mj_getState(self.model, self.data, state1b, spec)
# Expect next states to be equal.
np.testing.assert_array_equal(state1a, state1b)
def test_mj_jacSite(self): # pylint: disable=invalid-name
mujoco.mj_forward(self.model, self.data)
site_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, 'mysite')
python/mujoco/functions.cc
+19
View File
@@ -269,6 +269,25 @@ PYBIND11_MODULE(_functions, pymodule) {
});
// Support
Def<traits::mj_stateSize>(pymodule);
Def<traits::mj_getState>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d,
Eigen::Ref<EigenVectorX> state, unsigned int spec) {
if (state.size() != mj_stateSize(m, spec)) {
throw py::type_error("state size should equal mj_stateSize(m, spec)");
}
return InterceptMjErrors(::mj_getState)(m, d, state.data(), spec);
});
Def<traits::mj_setState>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
const Eigen::Ref<EigenVectorX> state, unsigned int spec) {
if (state.size() != mj_stateSize(m, spec)) {
throw py::type_error("state size should equal mj_stateSize(m, spec)");
}
return InterceptMjErrors(::mj_setState)(m, d, state.data(), spec);
});
Def<traits::mj_addContact>(pymodule);
Def<traits::mj_isPyramidal>(pymodule);
Def<traits::mj_isSparse>(pymodule);
src/engine/engine_support.c
+115
View File
@@ -36,6 +36,8 @@
#endif
#endif
//-------------------------- Constants -------------------------------------------------------------
#define mjVERSION 236
#define mjVERSIONSTRING "2.3.6"
@@ -87,6 +89,119 @@ const char* mjTIMERSTRING[mjNTIMER]= {
//-------------------------- get/set state ---------------------------------------------------------
// return size of a single state element
static inline int mj_stateElemSize(const mjModel* m, mjtState spec) {
switch (spec) {
case mjSTATE_TIME: return 1;
case mjSTATE_QPOS: return m->nq;
case mjSTATE_QVEL: return m->nv;
case mjSTATE_ACT: return m->na;
case mjSTATE_WARMSTART: return m->nv;
case mjSTATE_CTRL: return m->nu;
case mjSTATE_QFRC_APPLIED: return m->nv;
case mjSTATE_XFRC_APPLIED: return 6*m->nbody;
case mjSTATE_MOCAP_POS: return 3*m->nmocap;
case mjSTATE_MOCAP_QUAT: return 4*m->nmocap;
case mjSTATE_USERDATA: return m->nuserdata;
case mjSTATE_PLUGIN: return m->npluginstate;
default:
mju_error("mju_stateElementSize: invalid state element %u", spec);
return 0;
}
}
// return pointer to a single state element
static inline mjtNum* mj_stateElemPtr(const mjModel* m, mjData* d, mjtState spec) {
switch (spec) {
case mjSTATE_TIME: return &d->time;
case mjSTATE_QPOS: return d->qpos;
case mjSTATE_QVEL: return d->qvel;
case mjSTATE_ACT: return d->act;
case mjSTATE_WARMSTART: return d->qacc_warmstart;
case mjSTATE_CTRL: return d->ctrl;
case mjSTATE_QFRC_APPLIED: return d->qfrc_applied;
case mjSTATE_XFRC_APPLIED: return d->xfrc_applied;
case mjSTATE_MOCAP_POS: return d->mocap_pos;
case mjSTATE_MOCAP_QUAT: return d->mocap_quat;
case mjSTATE_USERDATA: return d->userdata;
case mjSTATE_PLUGIN: return d->plugin_state;
default:
mju_error("mju_stateElemPtr: invalid state element %u", spec);
return NULL;
}
}
static inline const mjtNum* mj_stateElemConstPtr(const mjModel* m, const mjData* d, mjtState spec) {
return mj_stateElemPtr(m, (mjData*) d, spec); // discard const qualifier from d
}
// get size of state specification
int mj_stateSize(const mjModel* m, unsigned int spec) {
if (spec >= (1<<mjNSTATE)) {
mju_error("mj_stateSize: invalid state spec %u >= 2^mjNSTATE", spec);
}
int size = 0;
for (int i=0; i < mjNSTATE; i++) {
mjtState element = 1<<i;
if (element & spec) {
size += mj_stateElemSize(m, element);
}
}
return size;
}
// get state
void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int spec) {
if (spec >= (1<<mjNSTATE)) {
mju_error("mj_getState: invalid state spec %u >= 2^mjNSTATE", spec);
}
int adr = 0;
for (int i=0; i < mjNSTATE; i++) {
mjtState element = 1<<i;
if (element & spec) {
int size = mj_stateElemSize(m, element);
const mjtNum* ptr = mj_stateElemConstPtr(m, d, element);
mju_copy(state + adr, ptr, size);
adr += size;
}
}
}
// set state
void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int spec) {
if (spec >= (1<<mjNSTATE)) {
mju_error("mj_setState: invalid state spec %u >= 2^mjNSTATE", spec);
}
int adr = 0;
for (int i=0; i < mjNSTATE; i++) {
mjtState element = 1<<i;
if (element & spec) {
int size = mj_stateElemSize(m, element);
mjtNum* ptr = mj_stateElemPtr(m, d, element);
mju_copy(ptr, state + adr, size);
adr += size;
}
}
}
//-------------------------- Jacobians -------------------------------------------------------------
// compute 3/6-by-nv Jacobian of global point attached to given body
src/engine/engine_support.h
+12
View File
@@ -29,6 +29,18 @@ MJAPI extern const char* mjENABLESTRING[mjNENABLE];
MJAPI extern const char* mjTIMERSTRING[mjNTIMER];
//-------------------------- get/set state ---------------------------------------------------------
// return size of state specification
MJAPI int mj_stateSize(const mjModel* m, unsigned int spec);
// get state
MJAPI void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int spec);
// set state
MJAPI void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int spec);
//-------------------------- Jacobians -------------------------------------------------------------
// compute 3/6-by-nv Jacobian of global point attached to given body
test/engine/engine_support_test.cc
+58 -4
View File
@@ -78,7 +78,7 @@ TEST_F(JacobianTest, SubtreeJac) {
mjtNum* qpos = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nq);
mjtNum* nudge = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv);
// all we need for Jacobians are kinematics and CoM-related quantitites
// all we need for Jacobians are kinematics and CoM-related quantities
mj_kinematics(model, data);
mj_comPos(model, data);
@@ -93,7 +93,7 @@ TEST_F(JacobianTest, SubtreeJac) {
// compare analytic Jacobian to finite-difference approximation
static const mjtNum eps = 1e-6;
for (int i=0; i<nv; i++) {
for (int i=0; i < nv; i++) {
// reset qpos, nudge i-th dof, update data->qpos, reset nudge
mju_copy(data->qpos, qpos, model->nq);
nudge[i] = 1;
@@ -105,7 +105,7 @@ TEST_F(JacobianTest, SubtreeJac) {
mj_comPos(model, data);
// compare finite-differenced and analytic Jacobian
for (int j=0; j<3; j++) {
for (int j=0; j < 3; j++) {
mjtNum findiff = (data->subtree_com[3*bodyid+j] - subtree_com[j]) / eps;
EXPECT_THAT(jac_subtree[nv*j+i], DoubleNear(findiff, eps));
}
@@ -127,7 +127,7 @@ TEST_F(JacobianTest, SubtreeJacNoInternalAcc) {
mjData* data = mj_makeData(model);
mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
// all we need for Jacobians are kinematics and CoM-related quantitites
// all we need for Jacobians are kinematics and CoM-related quantities
mj_kinematics(model, data);
mj_comPos(model, data);
@@ -351,5 +351,59 @@ TEST_F(SupportTest, DifferentiatePosSubQuat) {
mj_deleteModel(model);
}
static const char* const kDefaultModel = "testdata/model.xml";
TEST_F(SupportTest, GetSetStateStepEqual) {
const std::string xml_path = GetTestDataFilePath(kDefaultModel);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
mjData* data = mj_makeData(model);
// make distribution using seed
std::mt19937_64 rng;
rng.seed(3);
std::normal_distribution<double> dist(0, .01);
// set controls and applied joint forces to random values
for (int i=0; i < model->nu; i++) data->ctrl[i] = dist(rng);
for (int i=0; i < model->nv; i++) data->qfrc_applied[i] = dist(rng);
// take one step
mj_step(model, data);
int spec = mjSTATE_INTEGRATION;
int size = mj_stateSize(model, spec);
// save the initial state and step
std::vector<mjtNum> state0a(size);
mj_getState(model, data, state0a.data(), spec);
// get the initial state, expect equality
std::vector<mjtNum> state0b(size);
mj_getState(model, data, state0b.data(), spec);
EXPECT_EQ(state0a, state0b);
// take one step
mj_step(model, data);
// save the resulting state
std::vector<mjtNum> state1a(size);
mj_getState(model, data, state1a.data(), spec);
// expect the state to be different after stepping
EXPECT_THAT(state0a, testing::Ne(state1a));
// reset to the saved state, step again, get the resulting state
mj_setState(model, data, state0a.data(), spec);
mj_step(model, data);
std::vector<mjtNum> state1b(size);
mj_getState(model, data, state1b.data(), spec);
// expect the state to be the same after re-stepping
EXPECT_EQ(state1a, state1b);
mj_deleteData(data);
mj_deleteModel(model);
}
} // namespace
} // namespace mujoco
unity/Runtime/Bindings/MjBindings.cs
+9
View File
@@ -3191,6 +3191,15 @@ public static unsafe extern void mj_referenceConstraint(mjModel_* m, mjData_* d)
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_constraintUpdate(mjModel_* m, mjData_* d, double* jar, double* cost, int flg_coneHessian);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern int mj_stateSize(mjModel_* m, uint spec);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_getState(mjModel_* m, mjData_* d, double* state, uint spec);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_setState(mjModel_* m, mjData_* d, double* state, uint spec);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern int mj_addContact(mjModel_* m, mjData_* d, mjContact_* con);