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

Add mj_extractState as a public API function.

Browse Source 
This function allows the caller to extract a subset of components of a state previously obtained via `mj_getState` without having to first write it back into `mjData`.

PiperOrigin-RevId: 822165077
Change-Id: I6433261f4f5ec2e8d024bb7dee8f7dbf1e62666c
This commit is contained in:
Saran Tunyasuvunakool
2025-10-21 10:00:24 -07:00
committed by Copybara-Service
parent 25f5be82ff
commit 2f65e23779
11 changed files with 171 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/functions.cc
+15 -1
View File
@@ -329,10 +329,24 @@ PYBIND11_MODULE(_functions, pymodule) {
}
return InterceptMjErrors(::mj_getState)(m, d, state.data(), sig);
});
Def<traits::mj_extractState>(
pymodule,
[](const raw::MjModel* m,
Eigen::Ref<const EigenVectorX> src, unsigned int srcsig,
Eigen::Ref<EigenVectorX> dst, unsigned int dstsig) {
if (src.size() != mj_stateSize(m, srcsig)) {
throw py::type_error("src size should equal mj_stateSize(m, srcsig)");
}
if (dst.size() != mj_stateSize(m, dstsig)) {
throw py::type_error("dst size should equal mj_stateSize(m, dstsig)");
}
return InterceptMjErrors(::mj_extractState)(m, src.data(), srcsig,
dst.data(), dstsig);
});
Def<traits::mj_setState>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
const Eigen::Ref<EigenVectorX> state, unsigned int sig) {
Eigen::Ref<const EigenVectorX> state, unsigned int sig) {
if (state.size() != mj_stateSize(m, sig)) {
throw py::type_error("state size should equal mj_stateSize(m, sig)");
}