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

Split MjSpec out into header to allow use by other pybind modules.

Browse Source 
This should enable other pybind modules to accept arguments of type MjSpec.

PiperOrigin-RevId: 745580987
Change-Id: I435a8d264149369e3fbee14c5d2c628ef74b2541
This commit is contained in:
Sam Haves
2025-04-09 07:29:44 -07:00
committed by Copybara-Service
parent 85b323c7fb
commit 5c12ea871f
4 changed files with 172 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/CMakeLists.txt
+2
View File
@@ -417,6 +417,8 @@ endif()
mujoco_pybind11_module(
_specs
specs.cc
specs_wrapper.h
specs_wrapper.cc
specs.cc.inc
)
target_link_libraries(
python/mujoco/specs.cc
+1 -94
View File
@@ -30,6 +30,7 @@
#include <mujoco/mujoco.h>
#include "errors.h"
#include "indexers.h" // IWYU pragma: keep
#include "specs_wrapper.h" // IWYU pragma: keep
#include "raw.h"
#include "structs.h" // IWYU pragma: keep
#include <pybind11/cast.h>
@@ -70,100 +71,6 @@ using MjDoubleRef10 = Eigen::Ref<const Eigen::Matrix<double, 10, 1>>;
using MjDoubleRef11 = Eigen::Ref<const Eigen::Matrix<double, 11, 1>>;
using MjDoubleRefVec = Eigen::Ref<const Eigen::VectorXd>;
struct MjSpec {
MjSpec() : ptr(mj_makeSpec()) {}
MjSpec(raw::MjSpec* ptr, const py::dict& assets_ = {}) : ptr(ptr) {
for (const auto [key, value] : assets_) {
assets[key] = value;
}
}
// copy constructor and assignment
MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) {
override_assets = other.override_assets;
for (const auto [key, value] : other.assets) {
assets[key] = value;
}
parent = other.parent;
}
MjSpec& operator=(const MjSpec& other) {
override_assets = other.override_assets;
ptr = mj_copySpec(other.ptr);
for (const auto [key, value] : other.assets) {
assets[key] = value;
}
parent = other.parent;
return *this;
}
// move constructor and move assignment
MjSpec(MjSpec&& other) : ptr(other.ptr) {
override_assets = other.override_assets;
other.ptr = nullptr;
for (const auto [key, value] : other.assets) {
assets[key] = value;
}
other.assets.clear();
parent = other.parent;
other.parent = nullptr;
}
MjSpec& operator=(MjSpec&& other) {
override_assets = other.override_assets;
ptr = other.ptr;
other.ptr = nullptr;
for (const auto [key, value] : other.assets) {
assets[key] = value;
}
other.assets.clear();
parent = other.parent;
other.parent = nullptr;
return *this;
}
~MjSpec() {
mj_deleteSpec(ptr);
}
raw::MjModel* Compile() {
if (assets.empty()) {
auto m = mj_compile(ptr, 0);
if (!m || mjs_isWarning(ptr)) {
throw py::value_error(mjs_getError(ptr));
}
return m;
}
mjVFS vfs;
mj_defaultVFS(&vfs);
for (const auto& asset : assets) {
std::string buffer_name =
_impl::StripPath(py::cast<std::string>(asset.first).c_str());
std::string buffer = py::cast<std::string>(asset.second);
const int vfs_error = InterceptMjErrors(mj_addBufferVFS)(
&vfs, buffer_name.c_str(), buffer.c_str(), buffer.size());
if (vfs_error) {
mj_deleteVFS(&vfs);
if (vfs_error == 2) {
throw py::value_error("Repeated file name in assets dict: " +
buffer_name);
} else {
throw py::value_error("Asset failed to load: " + buffer_name);
}
}
}
auto m = mj_compile(ptr, &vfs);
if (!m || mjs_isWarning(ptr)) {
throw py::value_error(mjs_getError(ptr));
}
mj_deleteVFS(&vfs);
return m;
}
raw::MjSpec* ptr;
py::dict assets;
bool override_assets = true;
MjSpec* parent = nullptr;
};
template <typename LoadFunc>
static raw::MjSpec* LoadSpecFileImpl(
const std::string& filename,
python/mujoco/specs_wrapper.cc
+124
View File
@@ -0,0 +1,124 @@
// Copyright 2024 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "specs_wrapper.h"
#include <cstddef> // IWYU pragma: keep
#include <string>
#include <string_view> // IWYU pragma: keep
#include <vector> // IWYU pragma: keep
#include <mujoco/mjspec.h> // IWYU pragma: keep
#include <mujoco/mujoco.h>
#include "errors.h"
#include "indexers.h" // IWYU pragma: keep
#include "raw.h"
#include "structs.h" // IWYU pragma: keep
#include <pybind11/cast.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
namespace py = ::pybind11;
namespace mujoco::python {
MjSpec::MjSpec() : ptr(mj_makeSpec()) {}
MjSpec::MjSpec(raw::MjSpec* ptr, const py::dict& assets_) : ptr(ptr) {
for (const auto [key, value] : assets_) {
assets[key] = value;
}
}
// copy constructor and assignment
MjSpec::MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) {
override_assets = other.override_assets;
for (const auto [key, value] : other.assets) {
assets[key] = value;
}
parent = other.parent;
}
MjSpec& MjSpec::operator=(const MjSpec& other) {
override_assets = other.override_assets;
ptr = mj_copySpec(other.ptr);
for (const auto [key, value] : other.assets) {
assets[key] = value;
}
parent = other.parent;
return *this;
}
// move constructor and move assignment
MjSpec::MjSpec(MjSpec&& other) : ptr(other.ptr) {
override_assets = other.override_assets;
other.ptr = nullptr;
for (const auto [key, value] : other.assets) {
assets[key] = value;
}
other.assets.clear();
parent = other.parent;
other.parent = nullptr;
}
MjSpec& MjSpec::operator=(MjSpec&& other) {
override_assets = other.override_assets;
ptr = other.ptr;
other.ptr = nullptr;
for (const auto [key, value] : other.assets) {
assets[key] = value;
}
other.assets.clear();
parent = other.parent;
other.parent = nullptr;
return *this;
}
MjSpec::~MjSpec() { mj_deleteSpec(ptr); }
raw::MjModel* MjSpec::Compile() {
if (assets.empty()) {
auto m = mj_compile(ptr, 0);
if (!m || mjs_isWarning(ptr)) {
throw py::value_error(mjs_getError(ptr));
}
return m;
}
mjVFS vfs;
mj_defaultVFS(&vfs);
for (const auto& asset : assets) {
std::string buffer_name =
_impl::StripPath(py::cast<std::string>(asset.first).c_str());
std::string buffer = py::cast<std::string>(asset.second);
const int vfs_error = InterceptMjErrors(mj_addBufferVFS)(
&vfs, buffer_name.c_str(), buffer.c_str(), buffer.size());
if (vfs_error) {
mj_deleteVFS(&vfs);
if (vfs_error == 2) {
throw py::value_error("Repeated file name in assets dict: " +
buffer_name);
} else {
throw py::value_error("Asset failed to load: " + buffer_name);
}
}
}
auto m = mj_compile(ptr, &vfs);
if (!m || mjs_isWarning(ptr)) {
throw py::value_error(mjs_getError(ptr));
}
mj_deleteVFS(&vfs);
return m;
}
} // namespace mujoco::python
python/mujoco/specs_wrapper.h
+45
View File
@@ -0,0 +1,45 @@
// Copyright 2024 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "raw.h"
#include <pybind11/cast.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
namespace py = ::pybind11;
namespace mujoco::python {
struct MjSpec {
MjSpec();
MjSpec(raw::MjSpec* ptr, const py::dict& assets_ = {});
// copy constructor and assignment
MjSpec(const MjSpec& other);
MjSpec& operator=(const MjSpec& other);
// move constructor and move assignment
MjSpec(MjSpec&& other);
MjSpec& operator=(MjSpec&& other);
~MjSpec();
raw::MjModel* Compile();
raw::MjSpec* ptr;
py::dict assets;
bool override_assets = true;
MjSpec* parent = nullptr;
};
} // namespace mujoco::python