b362cb4972
This change also requires user scripts to explicitly synchronize changes to physics state to the viewer. The Simulate class was reconfigured so that certain UI events are handled during this sync operation, outside of the render loop on the main thread. These correspond to operations that require access to the full mjModel/mjData. To support other, more interactive operations (e.g. camera movements), a new mjvSceneState struct is introduced which captures only the portion of the physics state required for scene re-rendering. The mjvSceneState is updated from mjModel/mjData during the viewer sync operation, and is significantly cheaper than a full mj_copyModel and mj_copyData. Fixes https://github.com/deepmind/mujoco/issues/796 PiperOrigin-RevId: 525723636 Change-Id: Id08d0210a2c067d5afe85e2bf104f276aeddd75e
2264 lines
70 KiB
C++
2264 lines
70 KiB
C++
// Copyright 2022 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 "structs.h"
|
|
|
|
#include <Python.h>
|
|
|
|
#include <array>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <ios>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <sstream>
|
|
#include <string>
|
|
#include <string_view>
|
|
#include <type_traits>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include <absl/container/flat_hash_map.h>
|
|
#include <mujoco/mjxmacro.h>
|
|
#include <mujoco/mujoco.h>
|
|
#include "errors.h"
|
|
#include "function_traits.h"
|
|
#include "indexers.h"
|
|
#include "mjdata_meta.h"
|
|
#include "raw.h"
|
|
#include "serialization.h"
|
|
#include <pybind11/numpy.h>
|
|
#include <pybind11/operators.h>
|
|
#include <pybind11/pybind11.h>
|
|
#include <pybind11/pytypes.h>
|
|
#include <pybind11/stl.h>
|
|
|
|
namespace mujoco::python::_impl {
|
|
|
|
namespace py = ::pybind11;
|
|
|
|
namespace {
|
|
#define PTRDIFF(x, y) \
|
|
reinterpret_cast<const char*>(x) - reinterpret_cast<const char*>(y)
|
|
|
|
// Returns the shape of a NumPy array given the dimensions from an X Macro.
|
|
// If dim1 is a _literal_ constant 1, the resulting array is 1-dimensional of
|
|
// length dim0, otherwise the resulting array is 2-dimensional of shape
|
|
// (dim0, dim1).
|
|
#define X_ARRAY_SHAPE(dim0, dim1) XArrayShapeImpl(#dim1)((dim0), (dim1))
|
|
|
|
std::vector<int> XArrayShapeImpl1D(int dim0, int dim1) {
|
|
return {dim0};
|
|
}
|
|
|
|
std::vector<int> XArrayShapeImpl2D(int dim0, int dim1) {
|
|
return {dim0, dim1};
|
|
}
|
|
|
|
constexpr auto XArrayShapeImpl(const std::string_view dim1_str) {
|
|
if (dim1_str == "1") {
|
|
return XArrayShapeImpl1D;
|
|
} else {
|
|
return XArrayShapeImpl2D;
|
|
}
|
|
}
|
|
|
|
inline std::size_t NConMax(const mjData* d) {
|
|
return d->nstack * sizeof(mjtNum) / sizeof(mjContact);
|
|
}
|
|
} // namespace
|
|
|
|
// ==================== MJOPTION ===============================================
|
|
#define X(var, dim) , var(InitPyArray(std::array{dim}, ptr_->var, owner_))
|
|
MjOptionWrapper::MjWrapper()
|
|
: WrapperBase([]() {
|
|
raw::MjOption* const opt = new raw::MjOption;
|
|
mj_defaultOption(opt);
|
|
return opt;
|
|
}())
|
|
MJOPTION_VECTORS {}
|
|
|
|
MjOptionWrapper::MjWrapper(raw::MjOption* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner)
|
|
MJOPTION_VECTORS {}
|
|
#undef X
|
|
|
|
MjOptionWrapper::MjWrapper(const MjOptionWrapper& other)
|
|
: MjOptionWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJVISUAL ===============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjVisualHeadlightWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjVisualHeadlight{}),
|
|
X(ambient),
|
|
X(diffuse),
|
|
X(specular) {}
|
|
|
|
MjVisualHeadlightWrapper::MjWrapper(
|
|
raw::MjVisualHeadlight* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner),
|
|
X(ambient),
|
|
X(diffuse),
|
|
X(specular) {}
|
|
#undef X
|
|
|
|
MjVisualHeadlightWrapper::MjWrapper(const MjVisualHeadlightWrapper& other)
|
|
: MjVisualHeadlightWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjVisualRgbaWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjVisualRgba{}),
|
|
X(fog),
|
|
X(haze),
|
|
X(force),
|
|
X(inertia),
|
|
X(joint),
|
|
X(actuator),
|
|
X(actuatornegative),
|
|
X(actuatorpositive),
|
|
X(com),
|
|
X(camera),
|
|
X(light),
|
|
X(selectpoint),
|
|
X(connect),
|
|
X(contactpoint),
|
|
X(contactforce),
|
|
X(contactfriction),
|
|
X(contacttorque),
|
|
X(contactgap),
|
|
X(rangefinder),
|
|
X(constraint),
|
|
X(slidercrank),
|
|
X(crankbroken) {}
|
|
|
|
MjVisualRgbaWrapper::MjWrapper(raw::MjVisualRgba* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner),
|
|
X(fog),
|
|
X(haze),
|
|
X(force),
|
|
X(inertia),
|
|
X(joint),
|
|
X(actuator),
|
|
X(actuatornegative),
|
|
X(actuatorpositive),
|
|
X(com),
|
|
X(camera),
|
|
X(light),
|
|
X(selectpoint),
|
|
X(connect),
|
|
X(contactpoint),
|
|
X(contactforce),
|
|
X(contactfriction),
|
|
X(contacttorque),
|
|
X(contactgap),
|
|
X(rangefinder),
|
|
X(constraint),
|
|
X(slidercrank),
|
|
X(crankbroken) {}
|
|
#undef X
|
|
|
|
MjVisualRgbaWrapper::MjWrapper(const MjVisualRgbaWrapper& other)
|
|
: MjVisualRgbaWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
MjVisualWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjVisual{}),
|
|
headlight(&ptr_->headlight, owner_),
|
|
rgba(&ptr_->rgba, owner_) {}
|
|
|
|
MjVisualWrapper::MjWrapper(raw::MjVisual* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner),
|
|
headlight(&ptr_->headlight, owner_),
|
|
rgba(&ptr_->rgba, owner_) {}
|
|
|
|
|
|
MjVisualWrapper::MjWrapper(const MjVisualWrapper& other)
|
|
: MjVisualWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJMODEL ================================================
|
|
static void MjModelCapsuleDestructor(PyObject* pyobj) {
|
|
mj_deleteModel(
|
|
static_cast<raw::MjModel*>(PyCapsule_GetPointer(pyobj, nullptr)));
|
|
}
|
|
|
|
static absl::flat_hash_map<raw::MjModel*, MjModelWrapper*>&
|
|
MjModelRawPointerMap() {
|
|
static auto* hash_map =
|
|
new absl::flat_hash_map<raw::MjModel*, MjModelWrapper*>();
|
|
return *hash_map;
|
|
}
|
|
|
|
MjModelWrapper* MjModelWrapper::FromRawPointer(raw::MjModel* m) noexcept {
|
|
try {
|
|
auto& map = MjModelRawPointerMap();
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
auto found = map.find(m);
|
|
return found != map.end() ? found->second : nullptr;
|
|
}
|
|
} catch (...) {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
#undef MJ_M
|
|
#define MJ_M(x) ptr_->x
|
|
#define X(dtype, var, dim0, dim1) \
|
|
, var (InitPyArray(X_ARRAY_SHAPE(ptr_->dim0, dim1), ptr_->var, owner_))
|
|
MjModelWrapper::MjWrapper(raw::MjModel* ptr)
|
|
: WrapperBase(ptr, &MjModelCapsuleDestructor),
|
|
opt(&ptr->opt, owner_),
|
|
vis(&ptr->vis, owner_),
|
|
stat(&ptr->stat, owner_)
|
|
MJMODEL_POINTERS,
|
|
text_data_bytes(ptr->text_data, ptr->ntextdata),
|
|
names_bytes(ptr->names, ptr->nnames),
|
|
indexer_(ptr, owner_) {
|
|
bool is_newly_inserted = false;
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
is_newly_inserted = MjModelRawPointerMap().insert({ptr_, this}).second;
|
|
}
|
|
if (!is_newly_inserted) {
|
|
throw UnexpectedError(
|
|
"MjModelWrapper(mjModel*): MjModelRawPointerMap already contains this "
|
|
"raw mjModel*");
|
|
}
|
|
}
|
|
|
|
MjModelWrapper::MjWrapper(MjModelWrapper&& other)
|
|
: WrapperBase(other.ptr_, other.owner_),
|
|
opt(&ptr_->opt, owner_),
|
|
vis(&ptr_->vis, owner_),
|
|
stat(&ptr_->stat, owner_)
|
|
MJMODEL_POINTERS,
|
|
text_data_bytes(ptr_->text_data, ptr_->ntextdata),
|
|
names_bytes(ptr_->names, ptr_->nnames),
|
|
indexer_(ptr_, owner_) {
|
|
bool is_newly_inserted = false;
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
is_newly_inserted =
|
|
MjModelRawPointerMap().insert_or_assign(ptr_, this).second;
|
|
}
|
|
if (is_newly_inserted) {
|
|
throw UnexpectedError(
|
|
"MjModelRawPointerMap does not contains the moved-from mjModel*");
|
|
}
|
|
other.ptr_ = nullptr;
|
|
}
|
|
#undef X
|
|
#undef MJ_M
|
|
#define MJ_M(x) x
|
|
|
|
// Delegating to the MjModelWrapper::MjWrapper(raw::MjModel*) constructor,
|
|
// no need to modify MjModelRawPointerMap here.
|
|
MjModelWrapper::MjWrapper(const MjModelWrapper& other)
|
|
: MjModelWrapper(InterceptMjErrors(mj_copyModel)(NULL, other.get())) {}
|
|
|
|
MjModelWrapper::~MjWrapper() {
|
|
if (ptr_) {
|
|
bool erased = false;
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
erased = MjModelRawPointerMap().erase(ptr_);
|
|
}
|
|
if (!erased) {
|
|
std::cerr << "MjModelRawPointerMap does not contain this raw mjModel*"
|
|
<< std::endl;
|
|
std::terminate();
|
|
}
|
|
}
|
|
}
|
|
|
|
namespace {
|
|
struct VfsAsset {
|
|
VfsAsset(const char* name, const void* content, std::size_t content_size)
|
|
: name(name), content(content), content_size(content_size) {}
|
|
const char* name;
|
|
const void* content;
|
|
std::size_t content_size;
|
|
};
|
|
}
|
|
|
|
// Helper function for both LoadXMLFile and LoadBinaryFile.
|
|
// Creates a temporary MJB from the assets dictionary if one is supplied.
|
|
template <typename LoadFunc>
|
|
static raw::MjModel* LoadModelFileImpl(
|
|
const std::string& filename,
|
|
const std::vector<VfsAsset>& assets,
|
|
LoadFunc&& loadfunc) {
|
|
std::unique_ptr<mjVFS, void(*)(mjVFS*)> vfs(nullptr, [](mjVFS*){});
|
|
if (!assets.empty()) {
|
|
// mjVFS should be allocated on the heap, because it's ~2MB
|
|
vfs = decltype(vfs)(new mjVFS, [](mjVFS* vfs) {
|
|
mj_deleteVFS(vfs);
|
|
delete vfs;
|
|
});
|
|
mj_defaultVFS(vfs.get());
|
|
for (const auto& asset : assets) {
|
|
const int vfs_error = InterceptMjErrors(mj_makeEmptyFileVFS)(
|
|
vfs.get(), asset.name, asset.content_size);
|
|
if (vfs_error) {
|
|
throw py::value_error("assets dict is too big");
|
|
}
|
|
std::memcpy(vfs->filedata[vfs->nfile - 1],
|
|
asset.content, asset.content_size);
|
|
}
|
|
}
|
|
|
|
raw::MjModel* model = loadfunc(filename.c_str(), vfs.get());
|
|
if (model && !model->buffer) {
|
|
mj_deleteModel(model);
|
|
model = nullptr;
|
|
}
|
|
return model;
|
|
}
|
|
|
|
// Converts a dict with py::bytes value to a vector of standard C++ types.
|
|
// This allows us to release the GIL early. Note that the vector consists only
|
|
// of pointers to existing data so no substantial data copies are being made.
|
|
static std::vector<VfsAsset>
|
|
ConvertAssetsDict(
|
|
const std::optional<std::unordered_map<std::string, py::bytes>>& assets) {
|
|
std::vector<VfsAsset> out;
|
|
if (assets.has_value()) {
|
|
for (const auto& [name, content] : *assets) {
|
|
if (name.length() >= mjMAXVFSNAME) {
|
|
std::ostringstream error;
|
|
error << "Filename length " << name.length() << " exceeds "
|
|
<< mjMAXVFSNAME - 1 << " character limit: " << name;
|
|
throw py::value_error(error.str());
|
|
}
|
|
out.emplace_back(name.c_str(), PYBIND11_BYTES_AS_STRING(content.ptr()),
|
|
py::len(content));
|
|
}
|
|
}
|
|
return out;
|
|
}
|
|
|
|
MjModelWrapper MjModelWrapper::LoadXMLFile(
|
|
const std::string& filename,
|
|
const std::optional<std::unordered_map<std::string, py::bytes>>& assets) {
|
|
const auto converted_assets = ConvertAssetsDict(assets);
|
|
raw::MjModel* model;
|
|
{
|
|
py::gil_scoped_release no_gil;
|
|
char error[1024];
|
|
model = LoadModelFileImpl(
|
|
filename, converted_assets,
|
|
[&error](const char* filename, const mjVFS* vfs) {
|
|
return InterceptMjErrors(mj_loadXML)(
|
|
filename, vfs, error, sizeof(error));
|
|
});
|
|
if (!model) {
|
|
throw py::value_error(error);
|
|
}
|
|
}
|
|
return MjModelWrapper(model);
|
|
}
|
|
|
|
MjModelWrapper MjModelWrapper::LoadBinaryFile(
|
|
const std::string& filename,
|
|
const std::optional<std::unordered_map<std::string, py::bytes>>& assets) {
|
|
const auto converted_assets = ConvertAssetsDict(assets);
|
|
raw::MjModel* model;
|
|
{
|
|
py::gil_scoped_release no_gil;
|
|
model = LoadModelFileImpl(
|
|
filename, converted_assets, InterceptMjErrors(mj_loadModel));
|
|
if (!model) {
|
|
throw py::value_error("mj_loadModel: failed to load from mjb");
|
|
}
|
|
}
|
|
return MjModelWrapper(model);
|
|
}
|
|
|
|
MjModelWrapper MjModelWrapper::LoadXML(
|
|
const std::string& xml,
|
|
const std::optional<std::unordered_map<std::string, py::bytes>>& assets) {
|
|
auto converted_assets = ConvertAssetsDict(assets);
|
|
raw::MjModel* model;
|
|
{
|
|
py::gil_scoped_release no_gil;
|
|
std::string model_filename = "model_.xml";
|
|
if (assets.has_value()) {
|
|
while (assets->find(model_filename) != assets->end()) {
|
|
model_filename =
|
|
model_filename.substr(0, model_filename.size() - 4) + "_.xml";
|
|
}
|
|
}
|
|
converted_assets.emplace_back(
|
|
model_filename.c_str(), xml.c_str(), xml.length());
|
|
char error[1024];
|
|
model = LoadModelFileImpl(
|
|
model_filename, converted_assets,
|
|
[&error](const char* filename, const mjVFS* vfs) {
|
|
return InterceptMjErrors(mj_loadXML)(
|
|
filename, vfs, error, sizeof(error));
|
|
});
|
|
if (!model) {
|
|
throw py::value_error(error);
|
|
}
|
|
}
|
|
return MjModelWrapper(model);
|
|
}
|
|
|
|
namespace {
|
|
// A byte at the start of serialized mjModel structs, which can be incremented
|
|
// when we change the serialization logic to reject pickles from an unsupported
|
|
// future version.
|
|
constexpr static char kSerializationVersion = 1;
|
|
|
|
void CheckInput(const std::istream& input, std::string class_name) {
|
|
if (input.fail()) {
|
|
throw py::value_error("Invalid serialized " + class_name + ".");
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void MjModelWrapper::Serialize(std::ostream& output) const {
|
|
WriteChar(output, kSerializationVersion);
|
|
|
|
int model_size = mj_sizeModel(get());
|
|
WriteInt(output, model_size);
|
|
std::string buffer(model_size, 0);
|
|
mj_saveModel(get(), nullptr, buffer.data(), model_size);
|
|
WriteBytes(output, buffer.data(), model_size);
|
|
}
|
|
|
|
MjModelWrapper MjModelWrapper::Deserialize(std::istream& input) {
|
|
CheckInput(input, "mjModel");
|
|
|
|
char serializationVersion = ReadChar(input);
|
|
CheckInput(input, "mjModel");
|
|
|
|
if (serializationVersion != kSerializationVersion) {
|
|
throw py::value_error("Incompatible serialization version.");
|
|
}
|
|
|
|
int model_size = ReadInt(input);
|
|
CheckInput(input, "mjModel");
|
|
if (model_size < 0) {
|
|
throw py::value_error("Invalid serialized mjModel.");
|
|
}
|
|
std::string model_bytes(model_size, 0);
|
|
ReadBytes(input, model_bytes.data(), model_size);
|
|
CheckInput(input, "mjModel");
|
|
|
|
raw::MjModel* model = LoadModelFileImpl(
|
|
"model.mjb",
|
|
{{"model.mjb", model_bytes.data(), static_cast<std::size_t>(model_size)}},
|
|
InterceptMjErrors(mj_loadModel));
|
|
if (!model) {
|
|
throw py::value_error("Invalid serialized mjModel.");
|
|
}
|
|
return MjModelWrapper(model);
|
|
}
|
|
|
|
// ==================== MJCONTACT ==============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjContactWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjContact{}),
|
|
X(pos),
|
|
X(frame),
|
|
X(friction),
|
|
X(solref),
|
|
X(solimp),
|
|
X(H) {}
|
|
|
|
MjContactWrapper::MjWrapper(raw::MjContact* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner),
|
|
X(pos),
|
|
X(frame),
|
|
X(friction),
|
|
X(solref),
|
|
X(solimp),
|
|
X(H) {}
|
|
#undef X
|
|
|
|
MjContactWrapper::MjWrapper(const MjContactWrapper& other)
|
|
: MjContactWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
MjContactList::MjStructList(raw::MjContact* ptr, int nconmax,
|
|
int* ncon, py::handle owner)
|
|
: StructListBase(ptr, nconmax, owner, /* lazy = */ true),
|
|
ncon_(ncon) {}
|
|
|
|
// Slicing
|
|
MjContactList::MjStructList(MjContactList& other, py::slice slice)
|
|
: StructListBase(other, slice),
|
|
ncon_(other.ncon_) {}
|
|
|
|
// ==================== MJDATA =================================================
|
|
static void MjDataCapsuleDestructor(PyObject* pyobj) {
|
|
mj_deleteData(
|
|
static_cast<raw::MjData*>(PyCapsule_GetPointer(pyobj, nullptr)));
|
|
}
|
|
|
|
absl::flat_hash_map<raw::MjData*, MjDataWrapper*>&
|
|
MjDataRawPointerMap() {
|
|
static auto* hash_map =
|
|
new absl::flat_hash_map<raw::MjData*, MjDataWrapper*>();
|
|
return *hash_map;
|
|
}
|
|
|
|
MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept {
|
|
try {
|
|
auto& map = MjDataRawPointerMap();
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
auto found = map.find(m);
|
|
return found != map.end() ? found->second : nullptr;
|
|
}
|
|
} catch (...) {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
MjDataWrapper::MjWrapper(const MjModelWrapper& model)
|
|
: WrapperBase(InterceptMjErrors(mj_makeData)(model.get()),
|
|
&MjDataCapsuleDestructor),
|
|
#undef MJ_M
|
|
#define MJ_M(x) model.get()->x
|
|
#define X(dtype, var, dim0, dim1) \
|
|
var(InitPyArray(X_ARRAY_SHAPE(model.get()->dim0, dim1), ptr_->var, owner_)),
|
|
MJDATA_POINTERS
|
|
#undef MJ_M
|
|
#define MJ_M(x) (x)
|
|
#undef X
|
|
|
|
contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)),
|
|
|
|
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
|
|
MJDATA_VECTOR
|
|
#undef X
|
|
metadata_(model.get()),
|
|
indexer_(ptr_, &metadata_, owner_) {
|
|
bool is_newly_inserted = false;
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
|
}
|
|
if (!is_newly_inserted) {
|
|
throw UnexpectedError(
|
|
"MjDataRawPointerMap already contains this raw mjData*");
|
|
}
|
|
}
|
|
|
|
MjDataWrapper::MjWrapper(const MjDataWrapper& other)
|
|
: WrapperBase(other.Copy(), &MjDataCapsuleDestructor),
|
|
#undef MJ_M
|
|
#define MJ_M(x) other.metadata_.x
|
|
#define X(dtype, var, dim0, dim1) \
|
|
var(InitPyArray(X_ARRAY_SHAPE(other.metadata_.dim0, dim1), ptr_->var, \
|
|
owner_)),
|
|
MJDATA_POINTERS
|
|
#undef MJ_M
|
|
#define MJ_M(x) (x)
|
|
#undef X
|
|
|
|
contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)),
|
|
|
|
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
|
|
MJDATA_VECTOR
|
|
#undef X
|
|
metadata_(other.metadata_),
|
|
indexer_(ptr_, &metadata_, owner_) {
|
|
bool is_newly_inserted = false;
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
|
}
|
|
if (!is_newly_inserted) {
|
|
throw UnexpectedError(
|
|
"MjDataRawPointerMap already contains this raw mjData*");
|
|
}
|
|
}
|
|
|
|
MjDataWrapper::MjWrapper(MjDataWrapper&& other)
|
|
: WrapperBase(other.ptr_, other.owner_),
|
|
#undef MJ_M
|
|
#define MJ_M(x) other.metadata_.x
|
|
#define X(dtype, var, dim0, dim1) \
|
|
var(InitPyArray(X_ARRAY_SHAPE(other.metadata_.dim0, dim1), ptr_->var, \
|
|
owner_)),
|
|
MJDATA_POINTERS
|
|
#undef MJ_M
|
|
#define MJ_M(x) (x)
|
|
#undef X
|
|
|
|
contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon,owner_)),
|
|
|
|
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
|
|
MJDATA_VECTOR
|
|
#undef X
|
|
metadata_(other.metadata_),
|
|
indexer_(ptr_, &metadata_, owner_) {
|
|
bool is_newly_inserted = false;
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
is_newly_inserted =
|
|
MjDataRawPointerMap().insert_or_assign(ptr_, this).second;
|
|
}
|
|
if (is_newly_inserted) {
|
|
throw UnexpectedError(
|
|
"MjDataRawPointerMap does not contains the moved-from mjData*");
|
|
}
|
|
other.ptr_ = nullptr;
|
|
}
|
|
|
|
MjDataWrapper::MjWrapper(MjDataMetadata&& metadata, raw::MjData* d)
|
|
: WrapperBase(d, &MjDataCapsuleDestructor),
|
|
#undef MJ_M
|
|
#define MJ_M(x) metadata.x
|
|
#define X(dtype, var, dim0, dim1) \
|
|
var(InitPyArray(X_ARRAY_SHAPE(metadata.dim0, dim1), ptr_->var, owner_)),
|
|
MJDATA_POINTERS
|
|
#undef MJ_M
|
|
#define MJ_M(x) (x)
|
|
#undef X
|
|
|
|
contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)),
|
|
|
|
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
|
|
MJDATA_VECTOR
|
|
#undef X
|
|
metadata_(std::move(metadata)),
|
|
indexer_(ptr_, &metadata_, owner_) {
|
|
bool is_newly_inserted = false;
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
|
}
|
|
if (!is_newly_inserted) {
|
|
throw UnexpectedError(
|
|
"MjDataRawPointerMap already contains this raw mjData*");
|
|
}
|
|
}
|
|
|
|
MjDataWrapper::~MjWrapper() {
|
|
if (ptr_) {
|
|
bool erased = false;
|
|
{
|
|
py::gil_scoped_acquire gil;
|
|
erased = MjDataRawPointerMap().erase(ptr_);
|
|
}
|
|
if (!erased) {
|
|
std::cerr << "MjDataRawPointerMap does not contain this raw mjData*"
|
|
<< std::endl;
|
|
std::terminate();
|
|
}
|
|
}
|
|
}
|
|
|
|
void MjDataWrapper::Serialize(std::ostream& output) const {
|
|
// TODO: Replace this custom serialization with a protobuf
|
|
WriteChar(output, kSerializationVersion);
|
|
|
|
// Write all size fields
|
|
#define X(var) WriteInt(output, this->metadata_.var);
|
|
MJMODEL_INTS
|
|
#undef X
|
|
|
|
WriteInt(output, this->metadata_.is_dual);
|
|
|
|
#define X(dtype, var, n) \
|
|
WriteBytes(output, this->metadata_.var.get(), \
|
|
this->metadata_.n * sizeof(dtype));
|
|
|
|
MJDATA_METADATA
|
|
#undef X
|
|
|
|
// Write struct and scalar fields
|
|
#define X(var) WriteBytes(output, &ptr_->var, sizeof(ptr_->var))
|
|
X(parena);
|
|
X(maxuse_stack);
|
|
X(maxuse_arena);
|
|
X(maxuse_con);
|
|
X(maxuse_efc);
|
|
X(solver);
|
|
X(timer);
|
|
X(warning);
|
|
X(ne);
|
|
X(nf);
|
|
X(nnzJ);
|
|
X(nefc);
|
|
X(ncon);
|
|
X(time);
|
|
X(energy);
|
|
#undef X
|
|
|
|
// Write buffer contents
|
|
{
|
|
MJDATA_POINTERS_PREAMBLE((&this->metadata_))
|
|
|
|
#define X(type, name, nr, nc) \
|
|
WriteBytes(output, ptr_->name, sizeof(type)*(this->metadata_.nr)*(nc));
|
|
MJDATA_POINTERS
|
|
#undef X
|
|
|
|
#undef MJ_M
|
|
#define MJ_M(x) this->metadata_.x
|
|
#undef MJ_D
|
|
#define MJ_D(x) this->ptr_->x
|
|
#define X(type, name, nr, nc) \
|
|
if ((nr) * (nc)) { \
|
|
WriteInt(output, PTRDIFF(ptr_->name, ptr_->arena)); \
|
|
WriteBytes(output, ptr_->name, sizeof(type) * (nr) * (nc)); \
|
|
}
|
|
|
|
MJDATA_ARENA_POINTERS_CONTACT
|
|
MJDATA_ARENA_POINTERS_PRIMAL
|
|
|
|
if (this->metadata_.is_dual) {
|
|
MJDATA_ARENA_POINTERS_DUAL
|
|
}
|
|
#undef MJ_M
|
|
#define MJ_M(x) x
|
|
#undef MJ_D
|
|
#define MJ_D(x) x
|
|
#undef X
|
|
}
|
|
}
|
|
|
|
MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
|
|
char serializationVersion = ReadChar(input);
|
|
CheckInput(input, "mjData");
|
|
|
|
if (serializationVersion != kSerializationVersion) {
|
|
throw py::value_error("Incompatible serialization version.");
|
|
}
|
|
|
|
// Read all size and address fields
|
|
MjDataMetadata metadata;
|
|
raw::MjModel m{0};
|
|
|
|
#define X(var) \
|
|
metadata.var = ReadInt(input); \
|
|
CheckInput(input, "mjData"); \
|
|
m.var = metadata.var;
|
|
|
|
MJMODEL_INTS
|
|
#undef X
|
|
|
|
metadata.is_dual = ReadInt(input);
|
|
|
|
#define X(dtype, var, n) \
|
|
metadata.var.reset(new dtype[metadata.n]); \
|
|
ReadBytes(input, metadata.var.get(), metadata.n * sizeof(dtype)); \
|
|
CheckInput(input, "mjData"); \
|
|
m.var = metadata.var.get();
|
|
|
|
MJDATA_METADATA
|
|
#undef X
|
|
|
|
raw::MjData* d = mj_makeData(&m);
|
|
if (!d) {
|
|
throw py::value_error("Failed to create mjData.");
|
|
}
|
|
|
|
// Read structs and scalar fields
|
|
#define X(var) \
|
|
ReadBytes(input, (void*) &d->var, sizeof(d->var)); \
|
|
CheckInput(input, "mjData");
|
|
|
|
X(parena);
|
|
X(maxuse_stack);
|
|
X(maxuse_arena);
|
|
X(maxuse_con);
|
|
X(maxuse_efc);
|
|
X(solver);
|
|
X(timer);
|
|
X(warning);
|
|
X(ne);
|
|
X(nf);
|
|
X(nnzJ);
|
|
X(nefc);
|
|
X(ncon);
|
|
X(time);
|
|
X(energy);
|
|
#undef X
|
|
|
|
// Read buffer contents
|
|
{
|
|
MJDATA_POINTERS_PREAMBLE((&m))
|
|
|
|
#define X(type, name, nr, nc) \
|
|
ReadBytes(input, d->name, sizeof(type)*(m.nr)*(nc));
|
|
MJDATA_POINTERS
|
|
#undef X
|
|
|
|
#undef MJ_M
|
|
#define MJ_M(x) m.x
|
|
#undef MJ_D
|
|
#define MJ_D(x) d->x
|
|
#define X(type, name, nr, nc) \
|
|
if ((nr) * (nc)) { \
|
|
d->name = reinterpret_cast<decltype(d->name)>( \
|
|
static_cast<char*>(d->arena) + ReadInt(input)); \
|
|
ReadBytes(input, d->name, sizeof(type) * (nr) * (nc)); \
|
|
}
|
|
|
|
MJDATA_ARENA_POINTERS_CONTACT
|
|
MJDATA_ARENA_POINTERS_PRIMAL
|
|
|
|
if (metadata.is_dual) {
|
|
MJDATA_ARENA_POINTERS_DUAL
|
|
}
|
|
#undef MJ_M
|
|
#define MJ_M(x) x
|
|
#undef MJ_D
|
|
#define MJ_D(x) x
|
|
#undef X
|
|
}
|
|
CheckInput(input, "mjData");
|
|
|
|
// All bytes should have been used.
|
|
input.ignore(1);
|
|
if (!input.eof()) {
|
|
throw py::value_error("Invalid serialized mjData.");
|
|
}
|
|
|
|
return MjDataWrapper(std::move(metadata), d);
|
|
}
|
|
|
|
raw::MjData* MjDataWrapper::Copy() const {
|
|
raw::MjModel m{0};
|
|
#define X(var) m.var = this->metadata_.var;
|
|
MJMODEL_INTS
|
|
#undef X
|
|
#define X(dtype, var, n) m.var = this->metadata_.var.get();
|
|
MJDATA_METADATA
|
|
#undef X
|
|
return InterceptMjErrors(mj_copyData)(NULL, &m, this->ptr_);
|
|
}
|
|
|
|
// ==================== MJSTATISTIC ============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjStatisticWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjStatistic{}),
|
|
X(center) {}
|
|
|
|
MjStatisticWrapper::MjWrapper(raw::MjStatistic* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner),
|
|
X(center) {}
|
|
#undef X
|
|
|
|
MjStatisticWrapper::MjWrapper(const MjStatisticWrapper& other)
|
|
: MjStatisticWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJWARNINGSTAT ==========================================
|
|
MjWarningStatWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjWarningStat{}) {}
|
|
|
|
MjWarningStatWrapper::MjWrapper(raw::MjWarningStat* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner) {}
|
|
|
|
MjWarningStatWrapper::MjWrapper(const MjWarningStatWrapper& other)
|
|
: MjWarningStatWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
#define X(type, var) \
|
|
var(std::vector<int>{num}, std::vector<int>{sizeof(raw::MjWarningStat)}, \
|
|
&ptr->var, owner)
|
|
MjWarningStatList::MjStructList(raw::MjWarningStat* ptr, int num,
|
|
py::handle owner)
|
|
: StructListBase(ptr, num, owner), X(int, lastinfo), X(int, number) {}
|
|
#undef X
|
|
|
|
// Slicing
|
|
#define X(type, var) var(other.var[slice])
|
|
MjWarningStatList::MjStructList(MjWarningStatList& other, py::slice slice)
|
|
: StructListBase(other, slice),
|
|
X(int, lastinfo),
|
|
X(int, number) {}
|
|
#undef X
|
|
|
|
// ==================== MJTIMERSTAT ============================================
|
|
MjTimerStatWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjTimerStat{}) {}
|
|
|
|
MjTimerStatWrapper::MjWrapper(raw::MjTimerStat* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner) {}
|
|
|
|
MjTimerStatWrapper::MjWrapper(const MjTimerStatWrapper& other)
|
|
: MjTimerStatWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
#define X(type, var) \
|
|
var(std::vector<int>{num}, std::vector<int>{sizeof(raw::MjTimerStat)}, \
|
|
&ptr->var, owner)
|
|
MjTimerStatList::MjStructList(raw::MjTimerStat* ptr, int num, py::handle owner)
|
|
: StructListBase(ptr, num, owner),
|
|
X(mjtNum, duration),
|
|
X(int, number) {}
|
|
#undef X
|
|
|
|
// Slicing
|
|
#define X(type, var) var(other.var[slice])
|
|
MjTimerStatList::MjStructList(MjTimerStatList& other, py::slice slice)
|
|
: StructListBase(other, slice),
|
|
X(mjtNum, duration),
|
|
X(int, number) {}
|
|
#undef X
|
|
|
|
// ==================== MJSOLVERSTAT ===========================================
|
|
MjSolverStatWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjSolverStat{}) {}
|
|
|
|
MjSolverStatWrapper::MjWrapper(raw::MjSolverStat* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner) {}
|
|
|
|
MjSolverStatWrapper::MjWrapper(const MjSolverStatWrapper& other)
|
|
: MjSolverStatWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
#define X(type, var) \
|
|
var(std::vector<int>{num}, std::vector<int>{sizeof(raw::MjSolverStat)}, \
|
|
&ptr->var, owner)
|
|
MjSolverStatList::MjStructList(raw::MjSolverStat* ptr, int num,
|
|
py::handle owner)
|
|
: StructListBase(ptr, num, owner),
|
|
X(mjtNum, improvement),
|
|
X(mjtNum, gradient),
|
|
X(mjtNum, lineslope),
|
|
X(int, nactive),
|
|
X(int, nchange),
|
|
X(int, neval),
|
|
X(int, nupdate) {}
|
|
#undef X
|
|
#undef XN
|
|
|
|
// Slicing
|
|
#define X(type, var) var(other.var[slice])
|
|
MjSolverStatList::MjStructList(MjSolverStatList& other, py::slice slice)
|
|
: StructListBase(other, slice),
|
|
X(mjtNum, improvement),
|
|
X(mjtNum, gradient),
|
|
X(mjtNum, lineslope),
|
|
X(int, nactive),
|
|
X(int, nchange),
|
|
X(int, neval),
|
|
X(int, nupdate) {}
|
|
#undef X
|
|
|
|
// ==================== MJVPERTURB =============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjvPerturbWrapper::MjWrapper()
|
|
: WrapperBase([]() {
|
|
raw::MjvPerturb* const pert = new raw::MjvPerturb;
|
|
mjv_defaultPerturb(pert);
|
|
return pert;
|
|
}()),
|
|
X(refpos),
|
|
X(refquat),
|
|
X(refselpos),
|
|
X(localpos) {}
|
|
#undef X
|
|
|
|
MjvPerturbWrapper::MjWrapper(const MjvPerturbWrapper& other)
|
|
: MjvPerturbWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJVCAMERA ==============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjvCameraWrapper::MjWrapper()
|
|
: WrapperBase([]() {
|
|
raw::MjvCamera* const cam = new raw::MjvCamera;
|
|
mjv_defaultCamera(cam);
|
|
return cam;
|
|
}()),
|
|
X(lookat) {}
|
|
#undef X
|
|
|
|
MjvCameraWrapper::MjWrapper(const MjvCameraWrapper& other)
|
|
: MjvCameraWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJVGLCAMERA ============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjvGLCameraWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjvGLCamera{}),
|
|
X(pos),
|
|
X(forward),
|
|
X(up) {}
|
|
|
|
MjvGLCameraWrapper::MjWrapper(raw::MjvGLCamera* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner),
|
|
X(pos),
|
|
X(forward),
|
|
X(up) {}
|
|
#undef X
|
|
|
|
MjvGLCameraWrapper::MjWrapper(raw::MjvGLCamera&& other)
|
|
: MjvGLCameraWrapper() {
|
|
*this->ptr_ = other;
|
|
}
|
|
|
|
MjvGLCameraWrapper::MjWrapper(const MjvGLCameraWrapper& other)
|
|
: MjvGLCameraWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJVGEOM ================================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjvGeomWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjvGeom{}),
|
|
X(texrepeat),
|
|
X(size),
|
|
X(pos),
|
|
mat([this]() {
|
|
static_assert(sizeof(ptr_->mat) == sizeof(ptr_->mat[0])*9);
|
|
return InitPyArray(std::array{3, 3}, ptr_->mat, owner_);
|
|
}()),
|
|
X(rgba) {}
|
|
|
|
MjvGeomWrapper::MjWrapper(raw::MjvGeom* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner),
|
|
X(texrepeat),
|
|
X(size),
|
|
X(pos),
|
|
mat([this]() {
|
|
static_assert(sizeof(ptr_->mat) == sizeof(ptr_->mat[0])*9);
|
|
return InitPyArray(std::array{3, 3}, ptr_->mat, owner_);
|
|
}()),
|
|
X(rgba) {}
|
|
#undef X
|
|
|
|
MjvGeomWrapper::MjWrapper(const MjvGeomWrapper& other)
|
|
: MjvGeomWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJVLIGHT ===============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjvLightWrapper::MjWrapper()
|
|
: WrapperBase(new raw::MjvLight{}),
|
|
X(pos),
|
|
X(dir),
|
|
X(attenuation),
|
|
X(ambient),
|
|
X(diffuse),
|
|
X(specular) {}
|
|
|
|
MjvLightWrapper::MjWrapper(raw::MjvLight* ptr, py::handle owner)
|
|
: WrapperBase(ptr, owner),
|
|
X(pos),
|
|
X(dir),
|
|
X(attenuation),
|
|
X(ambient),
|
|
X(diffuse),
|
|
X(specular) {}
|
|
#undef X
|
|
|
|
MjvLightWrapper::MjWrapper(const MjvLightWrapper& other)
|
|
: MjvLightWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJVOPTION ==============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjvOptionWrapper::MjWrapper()
|
|
: WrapperBase([]() {
|
|
raw::MjvOption* const opt = new raw::MjvOption;
|
|
mjv_defaultOption(opt);
|
|
return opt;
|
|
}()),
|
|
X(geomgroup),
|
|
X(sitegroup),
|
|
X(jointgroup),
|
|
X(tendongroup),
|
|
X(actuatorgroup),
|
|
X(skingroup),
|
|
X(flags) {}
|
|
#undef X
|
|
|
|
MjvOptionWrapper::MjWrapper(const MjvOptionWrapper& other)
|
|
: MjvOptionWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
// ==================== MJVSCENE ===============================================
|
|
static void MjvSceneCapsuleDestructor(PyObject* pyobj) {
|
|
py::gil_scoped_acquire gil;
|
|
auto* scn = static_cast<raw::MjvScene*>(PyCapsule_GetPointer(pyobj, nullptr));
|
|
if (scn) {
|
|
mjv_freeScene(scn);
|
|
delete scn;
|
|
}
|
|
}
|
|
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
#define XN(var, n) var(InitPyArray(std::array{n}, ptr_->var, owner_))
|
|
MjvSceneWrapper::MjWrapper()
|
|
: WrapperBase([]() {
|
|
raw::MjvScene *const scn = new raw::MjvScene;
|
|
mjv_defaultScene(scn);
|
|
InterceptMjErrors(mjv_makeScene)(nullptr, scn, 0);
|
|
return scn;
|
|
}(), &MjvSceneCapsuleDestructor),
|
|
nskinvert(0),
|
|
XN(geoms, 0),
|
|
XN(geomorder, 0),
|
|
XN(skinfacenum, 0),
|
|
XN(skinvertadr, 0),
|
|
XN(skinvertnum, 0),
|
|
XN(skinvert, 0),
|
|
XN(skinnormal, 0),
|
|
X(lights),
|
|
X(camera),
|
|
X(translate),
|
|
X(rotate),
|
|
X(flags),
|
|
X(framergb) {}
|
|
|
|
#define XN(var, n) var(InitPyArray(std::array{n}, ptr_->var, owner_))
|
|
MjvSceneWrapper::MjWrapper(const MjModelWrapper& model, int maxgeom)
|
|
: WrapperBase(
|
|
[maxgeom](const raw::MjModel* m) {
|
|
raw::MjvScene *const scn = new raw::MjvScene;
|
|
mjv_defaultScene(scn);
|
|
InterceptMjErrors(mjv_makeScene)(m, scn, maxgeom);
|
|
return scn;
|
|
}(model.get()),
|
|
&MjvSceneCapsuleDestructor),
|
|
nskinvert([](const raw::MjModel* m) {
|
|
int nskinvert = 0;
|
|
for (int i = 0; i < m->nskin; ++i) {
|
|
nskinvert += m->skin_vertnum[i];
|
|
}
|
|
return nskinvert;
|
|
}(model.get())),
|
|
XN(geoms, ptr_->maxgeom),
|
|
XN(geomorder, ptr_->maxgeom),
|
|
XN(skinfacenum, ptr_->nskin),
|
|
XN(skinvertadr, ptr_->nskin),
|
|
XN(skinvertnum, ptr_->nskin),
|
|
XN(skinvert, 3*nskinvert),
|
|
XN(skinnormal, 3*nskinvert),
|
|
X(lights),
|
|
X(camera),
|
|
X(translate),
|
|
X(rotate),
|
|
X(flags),
|
|
X(framergb) {}
|
|
#undef X
|
|
#undef XN
|
|
|
|
template <typename T>
|
|
static T* MallocAndCopy(const T* src, int count) {
|
|
if (src) {
|
|
T* out = static_cast<T*>(mju_malloc(count * sizeof(T)));
|
|
std::memcpy(out, src, count * sizeof(T));
|
|
return out;
|
|
} else {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
MjvSceneWrapper::MjWrapper(const MjvSceneWrapper& other)
|
|
: MjvSceneWrapper() {
|
|
mjv_freeScene(ptr_);
|
|
*ptr_ = *other.ptr_;
|
|
|
|
#define XN(var, n) \
|
|
ptr_->var = MallocAndCopy(other.ptr_->var, n); \
|
|
var = InitPyArray(std::array{n}, ptr_->var, owner_);
|
|
|
|
XN(geoms, ptr_->ngeom);
|
|
XN(geomorder, ptr_->ngeom);
|
|
XN(skinfacenum, ptr_->nskin);
|
|
XN(skinvertadr, ptr_->nskin);
|
|
XN(skinvertnum, ptr_->nskin);
|
|
XN(skinvert, 3*nskinvert);
|
|
XN(skinnormal, 3*nskinvert);
|
|
|
|
#undef XN
|
|
}
|
|
|
|
// ==================== MJVFIGURE ==============================================
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
MjvFigureWrapper::MjWrapper()
|
|
: WrapperBase([]() {
|
|
raw::MjvFigure* const fig = new raw::MjvFigure;
|
|
mjv_defaultFigure(fig);
|
|
return fig;
|
|
}()),
|
|
X(flg_ticklabel),
|
|
X(gridsize),
|
|
X(gridrgb),
|
|
X(figurergba),
|
|
X(panergba),
|
|
X(legendrgba),
|
|
X(textrgb),
|
|
X(linergb),
|
|
X(range),
|
|
X(highlight),
|
|
X(linepnt),
|
|
X(linedata),
|
|
X(xaxispixel),
|
|
X(yaxispixel),
|
|
X(xaxisdata),
|
|
X(yaxisdata),
|
|
#undef X
|
|
|
|
linename([](raw::MjvFigure* ptr, py::handle owner) {
|
|
// Use a macro to help us static_assert that the array extents here are kept
|
|
// in sync with mjVisualize.h.
|
|
#define MAKE_STR_ARRAY(N1, N2) \
|
|
static_assert( \
|
|
std::is_same_v<decltype(raw::MjvFigure::linename), char[N1][N2]>); \
|
|
return py::array(py::dtype("|S" #N2), N1, ptr->linename, owner);
|
|
|
|
MAKE_STR_ARRAY(mjMAXLINE, 100);
|
|
|
|
#undef MAKE_STR_ARRAY
|
|
}(ptr_, owner_)) {}
|
|
|
|
MjvFigureWrapper::MjWrapper(const MjvFigureWrapper& other)
|
|
: MjvFigureWrapper() {
|
|
*this->ptr_ = *other.ptr_;
|
|
}
|
|
|
|
PYBIND11_MODULE(_structs, m) {
|
|
py::module_::import("mujoco._enums");
|
|
|
|
// ==================== MJOPTION =============================================
|
|
py::class_<MjOptionWrapper> mjOption(m, "MjOption");
|
|
mjOption.def(py::init<>());
|
|
mjOption.def("__copy__", [](const MjOptionWrapper& other) {
|
|
return MjOptionWrapper(other);
|
|
});
|
|
mjOption.def("__deepcopy__", [](const MjOptionWrapper& other, py::dict) {
|
|
return MjOptionWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjOption);
|
|
|
|
#define X(type, var) \
|
|
mjOption.def_property( \
|
|
#var, [](const MjOptionWrapper& c) { return c.get()->var; }, \
|
|
[](MjOptionWrapper& c, type rhs) { c.get()->var = rhs; });
|
|
MJOPTION_SCALARS
|
|
#undef X
|
|
|
|
#define X(var, dim) DefinePyArray(mjOption, #var, &MjOptionWrapper::var);
|
|
MJOPTION_VECTORS
|
|
#undef X
|
|
|
|
// ==================== MJVISUAL =============================================
|
|
py::class_<MjVisualWrapper> mjVisual(m, "MjVisual");
|
|
mjVisual.def("__copy__", [](const MjVisualWrapper& other) {
|
|
return MjVisualWrapper(other);
|
|
});
|
|
mjVisual.def("__deepcopy__", [](const MjVisualWrapper& other, py::dict) {
|
|
return MjVisualWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjVisual);
|
|
|
|
py::class_<raw::MjVisualGlobal> mjVisualGlobal(mjVisual, "Global");
|
|
mjVisualGlobal.def("__copy__", [](const raw::MjVisualGlobal& other) {
|
|
return raw::MjVisualGlobal(other);
|
|
});
|
|
mjVisualGlobal.def(
|
|
"__deepcopy__", [](const raw::MjVisualGlobal& other, py::dict) {
|
|
return raw::MjVisualGlobal(other);
|
|
});
|
|
DefineStructFunctions(mjVisualGlobal);
|
|
#define X(var) mjVisualGlobal.def_readwrite(#var, &raw::MjVisualGlobal::var)
|
|
X(fovy);
|
|
X(ipd);
|
|
X(azimuth);
|
|
X(elevation);
|
|
X(linewidth);
|
|
X(glow);
|
|
X(realtime);
|
|
X(offwidth);
|
|
X(offheight);
|
|
X(ellipsoidinertia);
|
|
#undef X
|
|
|
|
py::class_<raw::MjVisualQuality> mjVisualQuality(mjVisual, "Quality");
|
|
mjVisualQuality.def("__copy__", [](const raw::MjVisualQuality& other) {
|
|
return raw::MjVisualQuality(other);
|
|
});
|
|
mjVisualQuality.def(
|
|
"__deepcopy__", [](const raw::MjVisualQuality& other, py::dict) {
|
|
return raw::MjVisualQuality(other);
|
|
});
|
|
DefineStructFunctions(mjVisualQuality);
|
|
#define X(var) mjVisualQuality.def_readwrite(#var, &raw::MjVisualQuality::var)
|
|
X(shadowsize);
|
|
X(offsamples);
|
|
X(numslices);
|
|
X(numstacks);
|
|
X(numquads);
|
|
#undef X
|
|
|
|
py::class_<MjVisualHeadlightWrapper> mjVisualHeadlight(mjVisual, "Headlight");
|
|
mjVisualHeadlight.def(
|
|
"__copy__", [](const MjVisualHeadlightWrapper& other) {
|
|
return MjVisualHeadlightWrapper(other);
|
|
});
|
|
mjVisualHeadlight.def(
|
|
"__deepcopy__", [](const MjVisualHeadlightWrapper& other, py::dict) {
|
|
return MjVisualHeadlightWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjVisualHeadlight);
|
|
#define X(var) \
|
|
DefinePyArray(mjVisualHeadlight, #var, &MjVisualHeadlightWrapper::var)
|
|
X(ambient);
|
|
X(diffuse);
|
|
X(specular);
|
|
#undef X
|
|
mjVisualHeadlight.def_property(
|
|
"active",
|
|
[](const MjVisualHeadlightWrapper& c) { return c.get()->active; },
|
|
[](MjVisualHeadlightWrapper& c, int rhs) {
|
|
return c.get()->active = rhs;
|
|
});
|
|
|
|
py::class_<raw::MjVisualMap> mjVisualMap(mjVisual, "Map");
|
|
mjVisualMap.def("__copy__", [](const raw::MjVisualMap& other) {
|
|
return raw::MjVisualMap(other);
|
|
});
|
|
mjVisualMap.def(
|
|
"__deepcopy__", [](const raw::MjVisualMap& other, py::dict) {
|
|
return raw::MjVisualMap(other);
|
|
});
|
|
DefineStructFunctions(mjVisualMap);
|
|
#define X(var) mjVisualMap.def_readwrite(#var, &raw::MjVisualMap::var)
|
|
X(stiffness);
|
|
X(stiffnessrot);
|
|
X(force);
|
|
X(torque);
|
|
X(alpha);
|
|
X(fogstart);
|
|
X(fogend);
|
|
X(znear);
|
|
X(zfar);
|
|
X(haze);
|
|
X(shadowclip);
|
|
X(shadowscale);
|
|
X(actuatortendon);
|
|
#undef X
|
|
|
|
py::class_<raw::MjVisualScale> mjVisualScale(mjVisual, "Scale");
|
|
mjVisualScale.def("__copy__", [](const raw::MjVisualScale& other) {
|
|
return raw::MjVisualScale(other);
|
|
});
|
|
mjVisualScale.def(
|
|
"__deepcopy__", [](const raw::MjVisualScale& other, py::dict) {
|
|
return raw::MjVisualScale(other);
|
|
});
|
|
DefineStructFunctions(mjVisualScale);
|
|
#define X(var) mjVisualScale.def_readwrite(#var, &raw::MjVisualScale::var)
|
|
X(forcewidth);
|
|
X(contactwidth);
|
|
X(contactheight);
|
|
X(connect);
|
|
X(com);
|
|
X(camera);
|
|
X(light);
|
|
X(selectpoint);
|
|
X(jointlength);
|
|
X(jointwidth);
|
|
X(actuatorlength);
|
|
X(actuatorwidth);
|
|
X(framelength);
|
|
X(framewidth);
|
|
X(constraint);
|
|
X(slidercrank);
|
|
#undef X
|
|
|
|
py::class_<MjVisualRgbaWrapper> mjVisualRgba(mjVisual, "Rgba");
|
|
mjVisualRgba.def("__copy__", [](const MjVisualRgbaWrapper& other) {
|
|
return MjVisualRgbaWrapper(other);
|
|
});
|
|
mjVisualRgba.def(
|
|
"__deepcopy__", [](const MjVisualRgbaWrapper& other, py::dict) {
|
|
return MjVisualRgbaWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjVisualRgba);
|
|
#define X(var) DefinePyArray(mjVisualRgba, #var, &MjVisualRgbaWrapper::var)
|
|
X(fog);
|
|
X(haze);
|
|
X(force);
|
|
X(inertia);
|
|
X(joint);
|
|
X(actuator);
|
|
X(actuatornegative);
|
|
X(actuatorpositive);
|
|
X(com);
|
|
X(camera);
|
|
X(light);
|
|
X(selectpoint);
|
|
X(connect);
|
|
X(contactpoint);
|
|
X(contactforce);
|
|
X(contactfriction);
|
|
X(contacttorque);
|
|
X(contactgap);
|
|
X(rangefinder);
|
|
X(constraint);
|
|
X(slidercrank);
|
|
X(crankbroken);
|
|
#undef X
|
|
|
|
#define X(var) \
|
|
mjVisual.def_property_readonly( \
|
|
#var, [](const MjVisualWrapper& c) -> auto& { return c.get()->var; });
|
|
// mjVisual.global is exposed as "global_" to avoid clash with the Python
|
|
// keyword.
|
|
mjVisual.def_property_readonly(
|
|
"global_",
|
|
[](const MjVisualWrapper& c) -> auto& { return c.get()->global; });
|
|
X(quality);
|
|
mjVisual.def_readonly("headlight", &MjVisualWrapper::headlight);
|
|
X(map);
|
|
X(scale);
|
|
mjVisual.def_readonly("rgba", &MjVisualWrapper::rgba);
|
|
#undef X
|
|
|
|
// ==================== MJMODEL ==============================================
|
|
py::class_<MjModelWrapper> mjModel(m, "MjModel");
|
|
mjModel.def_static(
|
|
"from_xml_string", &MjModelWrapper::LoadXML,
|
|
py::arg("xml"), py::arg_v("assets", py::none()),
|
|
py::doc(
|
|
R"(Loads an MjModel from an XML string and an optional assets dictionary.)"));
|
|
mjModel.def_static(
|
|
"from_xml_path", &MjModelWrapper::LoadXMLFile,
|
|
py::arg("filename"), py::arg_v("assets", py::none()),
|
|
py::doc(
|
|
R"(Loads an MjModel from an XML file and an optional assets dictionary.
|
|
|
|
The filename for the XML can also refer to a key in the assets dictionary.
|
|
This is useful for example when the XML is not available as a file on disk.)"));
|
|
mjModel.def_static(
|
|
"from_binary_path", &MjModelWrapper::LoadBinaryFile,
|
|
py::arg("filename"), py::arg_v("assets", py::none()),
|
|
py::doc(
|
|
R"(Loads an MjModel from an MJB file and an optional assets dictionary.
|
|
|
|
The filename for the MJB can also refer to a key in the assets dictionary.
|
|
This is useful for example when the MJB is not available as a file on disk.)"));
|
|
mjModel.def_property_readonly("_address", [](const MjModelWrapper& m) {
|
|
return reinterpret_cast<std::uintptr_t>(m.get());
|
|
});
|
|
mjModel.def("__copy__", [](const MjModelWrapper& other) {
|
|
return MjModelWrapper(other);
|
|
});
|
|
mjModel.def("__deepcopy__", [](const MjModelWrapper& other, py::dict) {
|
|
return MjModelWrapper(other);
|
|
});
|
|
mjModel.def(py::pickle(
|
|
[](const MjModelWrapper& m) { // __getstate__
|
|
std::ostringstream output(std::ios::out | std::ios::binary);
|
|
m.Serialize(output);
|
|
return py::bytes(output.str());
|
|
},
|
|
[](py::bytes b) { // __setstate__
|
|
std::istringstream input(b, std::ios::in | std::ios::binary);
|
|
return MjModelWrapper::Deserialize(input);
|
|
}));
|
|
|
|
mjModel.def_readonly("opt", &MjModelWrapper::opt);
|
|
mjModel.def_readonly("vis", &MjModelWrapper::vis);
|
|
mjModel.def_readonly("stat", &MjModelWrapper::stat);
|
|
|
|
#define X(var) \
|
|
mjModel.def_property_readonly( \
|
|
#var, [](const MjModelWrapper& m) { return m.get()->var; });
|
|
MJMODEL_INTS
|
|
#undef X
|
|
|
|
#define X(dtype, var, dim0, dim1) \
|
|
if constexpr (std::string_view(#var) != "text_data" && \
|
|
std::string_view(#var) != "names") { \
|
|
DefinePyArray(mjModel, #var, &MjModelWrapper::var); \
|
|
}
|
|
MJMODEL_POINTERS
|
|
#undef X
|
|
|
|
mjModel.def_property_readonly(
|
|
"text_data", [](const MjModelWrapper& m) -> const auto& {
|
|
// Return the full bytes array of concatenated text data
|
|
return m.text_data_bytes;
|
|
});
|
|
mjModel.def_property_readonly(
|
|
"names", [](const MjModelWrapper& m) -> const auto& {
|
|
// Return the full bytes array of concatenated names
|
|
return m.names_bytes;
|
|
});
|
|
|
|
#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
|
|
mjModel.def( \
|
|
#field, \
|
|
[](MjModelWrapper& m, int i) -> auto& { return m.indexer().field(i); }, \
|
|
py::return_value_policy::reference_internal); \
|
|
mjModel.def( \
|
|
#field, [](MjModelWrapper& m, std::string_view name) -> auto& { \
|
|
return m.indexer().field##_by_name(name); \
|
|
}, \
|
|
py::return_value_policy::reference_internal, py::arg_v("name", ""));
|
|
|
|
|
|
MJMODEL_VIEW_GROUPS
|
|
#undef XGROUP
|
|
|
|
#define XGROUP(field, altname, FIELD_XMACROS) \
|
|
mjModel.def( \
|
|
#altname, \
|
|
[](MjModelWrapper& m, int i) -> auto& { return m.indexer().field(i); }, \
|
|
py::return_value_policy::reference_internal); \
|
|
mjModel.def( \
|
|
#altname, [](MjModelWrapper& m, std::string_view name) -> auto& { \
|
|
return m.indexer().field##_by_name(name); \
|
|
}, \
|
|
py::return_value_policy::reference_internal, py::arg_v("name", ""));
|
|
|
|
MJMODEL_VIEW_GROUPS_ALTNAMES
|
|
#undef XGROUP
|
|
|
|
#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
|
|
{ \
|
|
using GroupedViews = MjModelGroupedViews; \
|
|
py::class_<MjModelGroupedViews> groupedViews(m, "_" #MjModelGroupedViews); \
|
|
FIELD_XMACROS \
|
|
groupedViews.def("__repr__", MjModelStructRepr<GroupedViews>); \
|
|
groupedViews.def_property_readonly("id", [](GroupedViews& views) { \
|
|
return views.index(); \
|
|
}); \
|
|
groupedViews.def_property_readonly("name", [](GroupedViews& views) { \
|
|
return views.name(); \
|
|
}); \
|
|
}
|
|
#define X(type, prefix, var, dim0, dim1) \
|
|
groupedViews.def_property( \
|
|
#var, &GroupedViews::var, [](GroupedViews& views, py::handle rhs) { \
|
|
(views.var())[py::slice(py::none(), py::none(), py::none())] = rhs; \
|
|
});
|
|
|
|
MJMODEL_VIEW_GROUPS
|
|
#undef X
|
|
#undef XGROUP
|
|
|
|
{
|
|
py::handle builtins(PyEval_GetBuiltins());
|
|
builtins[MjModelWrapper::kFromRawPointer] =
|
|
reinterpret_cast<std::uintptr_t>(reinterpret_cast<void*>(
|
|
&MjModelWrapper::FromRawPointer));
|
|
}
|
|
|
|
// ==================== MJWARNINGSTAT ========================================
|
|
py::class_<MjWarningStatWrapper> mjWarningStat(m, "MjWarningStat");
|
|
mjWarningStat.def(py::init<>());
|
|
mjWarningStat.def("__copy__", [](const MjWarningStatWrapper& other) {
|
|
return MjWarningStatWrapper(other);
|
|
});
|
|
mjWarningStat.def(
|
|
"__deepcopy__", [](const MjWarningStatWrapper& other, py::dict) {
|
|
return MjWarningStatWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjWarningStat);
|
|
#define X(var) \
|
|
mjWarningStat.def_property( \
|
|
#var, [](const MjWarningStatWrapper& d) { return d.get()->var; }, \
|
|
[](MjWarningStatWrapper& d, decltype(raw::MjWarningStat::var) rhs) { \
|
|
d.get()->var = rhs; \
|
|
});
|
|
X(lastinfo);
|
|
X(number);
|
|
#undef X
|
|
|
|
py::class_<MjWarningStatList> mjWarningStatList(m, "_MjWarningStatList");
|
|
mjWarningStatList.def(
|
|
"__getitem__",
|
|
&MjWarningStatList::operator[],
|
|
py::return_value_policy::reference);
|
|
mjWarningStatList.def(
|
|
"__getitem__",
|
|
[](MjWarningStatList& list, ::mjtWarning idx) { return list[idx]; },
|
|
py::return_value_policy::reference);
|
|
mjWarningStatList.def("__getitem__", &MjWarningStatList::Slice);
|
|
mjWarningStatList.def("__len__", &MjWarningStatList::size);
|
|
DefineStructFunctions(mjWarningStatList);
|
|
|
|
#define X(type, var) \
|
|
mjWarningStatList.def_readonly(#var, &MjWarningStatList::var)
|
|
X(int, lastinfo);
|
|
X(int, number);
|
|
#undef X
|
|
|
|
// ==================== MJTIMERSTAT ==========================================
|
|
py::class_<MjTimerStatWrapper> mjTimerStat(m, "MjTimerStat");
|
|
mjTimerStat.def(py::init<>());
|
|
mjTimerStat.def("__copy__", [](const MjTimerStatWrapper& other) {
|
|
return MjTimerStatWrapper(other);
|
|
});
|
|
mjTimerStat.def(
|
|
"__deepcopy__", [](const MjTimerStatWrapper& other, py::dict) {
|
|
return MjTimerStatWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjTimerStat);
|
|
#define X(var) \
|
|
mjTimerStat.def_property( \
|
|
#var, [](const MjTimerStatWrapper& d) { return d.get()->var; }, \
|
|
[](MjTimerStatWrapper& d, decltype(raw::MjTimerStat::var) rhs) { \
|
|
d.get()->var = rhs; \
|
|
});
|
|
X(duration);
|
|
X(number);
|
|
#undef X
|
|
|
|
py::class_<MjTimerStatList> mjTimerStatList(m, "_MjTimerStatList");
|
|
mjTimerStatList.def(
|
|
"__getitem__",
|
|
&MjTimerStatList::operator[],
|
|
py::return_value_policy::reference);
|
|
mjTimerStatList.def(
|
|
"__getitem__",
|
|
[](MjTimerStatList& list, ::mjtTimer idx) { return list[idx]; },
|
|
py::return_value_policy::reference);
|
|
mjTimerStatList.def("__getitem__", &MjTimerStatList::Slice);
|
|
mjTimerStatList.def("__len__", &MjTimerStatList::size);
|
|
DefineStructFunctions(mjTimerStatList);
|
|
|
|
#define X(type, var) \
|
|
mjTimerStatList.def_readonly(#var, &MjTimerStatList::var)
|
|
X(mjtNum, duration);
|
|
X(int, number);
|
|
#undef X
|
|
|
|
// ==================== MJSOLVERSTAT =========================================
|
|
py::class_<MjSolverStatWrapper> mjSolverStat(m, "MjSolverStat");
|
|
mjSolverStat.def(py::init<>());
|
|
mjSolverStat.def("__copy__", [](const MjSolverStatWrapper& other) {
|
|
return MjSolverStatWrapper(other);
|
|
});
|
|
mjSolverStat.def(
|
|
"__deepcopy__", [](const MjSolverStatWrapper& other, py::dict) {
|
|
return MjSolverStatWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjSolverStat);
|
|
#define X(var) \
|
|
mjSolverStat.def_property( \
|
|
#var, [](const MjSolverStatWrapper& d) { return d.get()->var; }, \
|
|
[](MjSolverStatWrapper& d, decltype(raw::MjSolverStat::var) rhs) { \
|
|
d.get()->var = rhs; \
|
|
});
|
|
X(improvement);
|
|
X(gradient);
|
|
X(lineslope);
|
|
X(nactive);
|
|
X(nchange);
|
|
X(neval);
|
|
X(nupdate);
|
|
#undef X
|
|
|
|
py::class_<MjSolverStatList> mjSolverStatList(m, "_MjSolverStatList");
|
|
mjSolverStatList.def("__getitem__", &MjSolverStatList::operator[],
|
|
py::return_value_policy::reference);
|
|
mjSolverStatList.def("__getitem__", &MjSolverStatList::Slice);
|
|
mjSolverStatList.def("__len__", &MjSolverStatList::size);
|
|
DefineStructFunctions(mjSolverStatList);
|
|
|
|
#define X(type, var) \
|
|
mjSolverStatList.def_readonly(#var, &MjSolverStatList::var)
|
|
X(mjtNum, improvement);
|
|
X(mjtNum, gradient);
|
|
X(mjtNum, lineslope);
|
|
X(int, nactive);
|
|
X(int, nchange);
|
|
X(int, neval);
|
|
X(int, nupdate);
|
|
#undef X
|
|
|
|
// ==================== MJCONTACT ============================================
|
|
py::class_<MjContactWrapper> mjContact(m, "MjContact");
|
|
mjContact.def(py::init<>());
|
|
mjContact.def("__copy__", [](const MjContactWrapper& self) {
|
|
return MjContactWrapper(self);
|
|
});
|
|
mjContact.def("__deepcopy__", [](const MjContactWrapper& self, py::dict) {
|
|
return MjContactWrapper(self);
|
|
});
|
|
DefineStructFunctions(mjContact);
|
|
|
|
#define X(var) \
|
|
mjContact.def_property( \
|
|
#var, [](const MjContactWrapper& c) { return c.get()->var; }, \
|
|
[](MjContactWrapper& c, decltype(raw::MjContact::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(dist);
|
|
X(includemargin);
|
|
X(mu);
|
|
X(dim);
|
|
X(geom1);
|
|
X(geom2);
|
|
X(exclude);
|
|
X(efc_address);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjContact, #var, &MjContactWrapper::var)
|
|
X(pos);
|
|
X(frame);
|
|
X(friction);
|
|
X(solref);
|
|
X(solimp);
|
|
X(H);
|
|
#undef X
|
|
|
|
py::class_<MjContactList> mjContactList(m, "_MjContactList");
|
|
mjContactList.def("__getitem__", &MjContactList::operator[],
|
|
py::return_value_policy::reference);
|
|
mjContactList.def("__getitem__", &MjContactList::Slice);
|
|
mjContactList.def("__len__", &MjContactList::size);
|
|
DefineStructFunctions(mjContactList);
|
|
|
|
#define X(type, var) \
|
|
mjContactList.def_property_readonly(#var, [](const MjContactList& c) { \
|
|
return py::array_t<type>(std::vector<int>{c.size()}, \
|
|
std::vector<int>{sizeof(raw::MjContact)}, \
|
|
&c.get()->var, c.owner()); \
|
|
});
|
|
#define XN(type, var) \
|
|
mjContactList.def_property_readonly(#var, [](const MjContactList& c) { \
|
|
return py::array_t<type>( \
|
|
std::vector<int>{c.size(), \
|
|
sizeof(raw::MjContact::var) / sizeof(type)}, \
|
|
std::vector<int>{sizeof(raw::MjContact), sizeof(type)}, \
|
|
&c.get()->var[0], c.owner()); \
|
|
});
|
|
X(mjtNum, dist);
|
|
XN(mjtNum, pos);
|
|
XN(mjtNum, frame);
|
|
X(mjtNum, includemargin);
|
|
XN(mjtNum, friction);
|
|
XN(mjtNum, solref);
|
|
XN(mjtNum, solimp);
|
|
X(mjtNum, mu);
|
|
XN(mjtNum, H);
|
|
X(int, dim);
|
|
X(int, geom1);
|
|
X(int, geom2);
|
|
X(int, exclude);
|
|
X(int, efc_address);
|
|
#undef X
|
|
#undef XN
|
|
|
|
// ==================== MJDATA ===============================================
|
|
py::class_<MjDataWrapper> mjData(m, "MjData");
|
|
mjData.def(py::init<const MjModelWrapper&>());
|
|
mjData.def_property_readonly("_address", [](const MjDataWrapper& d) {
|
|
return reinterpret_cast<std::uintptr_t>(d.get());
|
|
});
|
|
mjData.def("__copy__", [](const MjDataWrapper& other) {
|
|
return MjDataWrapper(other);
|
|
});
|
|
mjData.def("__deepcopy__", [](const MjDataWrapper& other, py::dict) {
|
|
return MjDataWrapper(other);
|
|
});
|
|
mjData.def(py::pickle(
|
|
[](const MjDataWrapper& d) { // __getstate__
|
|
std::ostringstream output(std::ios::out | std::ios::binary);
|
|
d.Serialize(output);
|
|
return py::bytes(output.str());
|
|
},
|
|
[](py::bytes b) { // __setstate__
|
|
std::istringstream input(b, std::ios::in | std::ios::binary);
|
|
return MjDataWrapper::Deserialize(input);
|
|
}));
|
|
|
|
#define X(type, var) \
|
|
mjData.def_property( \
|
|
#var, [](const MjDataWrapper& d) { return d.get()->var; }, \
|
|
[](MjDataWrapper& d, decltype(raw::MjData::var) rhs) { \
|
|
d.get()->var = rhs; \
|
|
});
|
|
MJDATA_SCALAR
|
|
#undef X
|
|
|
|
#define X(dtype, var, dim0, dim1) \
|
|
DefinePyArray(mjData, #var, &MjDataWrapper::var);
|
|
MJDATA_POINTERS
|
|
MJDATA_ARENA_POINTERS_CONTACT
|
|
#undef X
|
|
|
|
#undef MJ_M
|
|
#define MJ_M(x) d.metadata().x
|
|
#undef MJ_D
|
|
#define MJ_D(x) d.get()->x
|
|
#define X(dtype, var, dim0, dim1) \
|
|
mjData.def_property_readonly(#var, [](const MjDataWrapper& d) { \
|
|
return InitPyArray(X_ARRAY_SHAPE(dim0, dim1), d.get()->var, d.owner()); \
|
|
});
|
|
|
|
MJDATA_ARENA_POINTERS_PRIMAL
|
|
MJDATA_ARENA_POINTERS_DUAL
|
|
|
|
#undef MJ_M
|
|
#define MJ_M(x) (x)
|
|
#undef MJ_D
|
|
#define MJ_D(x) (x)
|
|
#undef X
|
|
|
|
#define X(dtype, var, dim0, dim1) \
|
|
DefinePyArray(mjData, #var, &MjDataWrapper::var);
|
|
MJDATA_VECTOR
|
|
#undef X
|
|
|
|
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
|
|
mjData.def( \
|
|
#field, \
|
|
[](MjDataWrapper& d, int i) -> auto& { return d.indexer().field(i); }, \
|
|
py::return_value_policy::reference_internal); \
|
|
mjData.def( \
|
|
#field, [](MjDataWrapper& d, std::string_view name) -> auto& { \
|
|
return d.indexer().field##_by_name(name); \
|
|
}, \
|
|
py::return_value_policy::reference_internal, py::arg_v("name", ""));
|
|
|
|
MJDATA_VIEW_GROUPS
|
|
#undef XGROUP
|
|
|
|
#define XGROUP(field, altname, FIELD_XMACROS) \
|
|
mjData.def( \
|
|
#altname, \
|
|
[](MjDataWrapper& d, int i) -> auto& { return d.indexer().field(i); }, \
|
|
py::return_value_policy::reference_internal); \
|
|
mjData.def( \
|
|
#altname, [](MjDataWrapper& d, std::string_view name) -> auto& { \
|
|
return d.indexer().field##_by_name(name); \
|
|
}, \
|
|
py::return_value_policy::reference_internal, py::arg_v("name", ""));
|
|
|
|
MJDATA_VIEW_GROUPS_ALTNAMES
|
|
#undef XGROUP
|
|
|
|
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
|
|
{ \
|
|
using GroupedViews = MjDataGroupedViews; \
|
|
py::class_<MjDataGroupedViews> groupedViews(m, "_" #MjDataGroupedViews); \
|
|
FIELD_XMACROS \
|
|
groupedViews.def("__repr__", MjDataStructRepr<GroupedViews>); \
|
|
groupedViews.def_property_readonly("id", [](GroupedViews& views) { \
|
|
return views.index(); \
|
|
}); \
|
|
groupedViews.def_property_readonly("name", [](GroupedViews& views) { \
|
|
return views.name(); \
|
|
}); \
|
|
}
|
|
#define X(type, prefix, var, dim0, dim1) \
|
|
groupedViews.def_property( \
|
|
#var, &GroupedViews::var, [](GroupedViews& views, py::handle rhs) { \
|
|
(views.var())[py::slice(py::none(), py::none(), py::none())] = rhs; \
|
|
});
|
|
|
|
MJDATA_VIEW_GROUPS
|
|
#undef X
|
|
#undef XGROUP
|
|
|
|
{
|
|
py::handle builtins(PyEval_GetBuiltins());
|
|
builtins[MjDataWrapper::kFromRawPointer] =
|
|
reinterpret_cast<std::uintptr_t>(reinterpret_cast<void*>(
|
|
&MjDataWrapper::FromRawPointer));
|
|
}
|
|
|
|
// ==================== MJSTATISTIC ==========================================
|
|
py::class_<MjStatisticWrapper> mjStatistic(m, "MjStatistic");
|
|
mjStatistic.def(py::init<>());
|
|
mjStatistic.def("__copy__", [](const MjStatisticWrapper& other) {
|
|
return MjStatisticWrapper(other);
|
|
});
|
|
mjStatistic.def(
|
|
"__deepcopy__", [](const MjStatisticWrapper& other, py::dict) {
|
|
return MjStatisticWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjStatistic);
|
|
|
|
#define X(var) \
|
|
mjStatistic.def_property( \
|
|
#var, [](const MjStatisticWrapper& c) { return c.get()->var; }, \
|
|
[](MjStatisticWrapper& c, decltype(raw::MjStatistic::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(meaninertia);
|
|
X(meanmass);
|
|
X(meansize);
|
|
X(extent);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjStatistic, #var, &MjStatisticWrapper::var)
|
|
X(center);
|
|
#undef X
|
|
|
|
// ==================== MJLROPT ==============================================
|
|
py::class_<raw::MjLROpt> mjLROpt(m, "MjLROpt");
|
|
mjLROpt.def(py::init<>());
|
|
mjLROpt.def("__copy__", [](const raw::MjLROpt& other) {
|
|
return raw::MjLROpt(other);
|
|
});
|
|
mjLROpt.def("__deepcopy__", [](const raw::MjLROpt& other, py::dict) {
|
|
return raw::MjLROpt(other);
|
|
});
|
|
DefineStructFunctions(mjLROpt);
|
|
#define X(var) mjLROpt.def_readwrite(#var, &raw::MjLROpt::var)
|
|
X(mode);
|
|
X(useexisting);
|
|
X(uselimit);
|
|
X(accel);
|
|
X(maxforce);
|
|
X(timeconst);
|
|
X(timestep);
|
|
X(inttotal);
|
|
X(interval);
|
|
X(tolrange);
|
|
#undef X
|
|
|
|
// ==================== MJVPERTURB ===========================================
|
|
py::class_<MjvPerturbWrapper> mjvPerturb(m, "MjvPerturb");
|
|
mjvPerturb.def(py::init<>());
|
|
mjvPerturb.def("__copy__", [](const MjvPerturbWrapper& other) {
|
|
return MjvPerturbWrapper(other);
|
|
});
|
|
mjvPerturb.def("__deepcopy__", [](const MjvPerturbWrapper& other, py::dict) {
|
|
return MjvPerturbWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjvPerturb);
|
|
#define X(var) \
|
|
mjvPerturb.def_property( \
|
|
#var, [](const MjvPerturbWrapper& c) { return c.get()->var; }, \
|
|
[](MjvPerturbWrapper& c, decltype(raw::MjvPerturb::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(select);
|
|
X(skinselect);
|
|
X(active);
|
|
X(active2);
|
|
X(localmass);
|
|
X(scale);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjvPerturb, #var, &MjvPerturbWrapper::var)
|
|
X(refpos);
|
|
X(refquat);
|
|
X(refselpos);
|
|
X(localpos);
|
|
#undef X
|
|
|
|
// ==================== MJVCAMERA ============================================
|
|
py::class_<MjvCameraWrapper> mjvCamera(m, "MjvCamera");
|
|
mjvCamera.def(py::init<>());
|
|
mjvCamera.def("__copy__", [](const MjvCameraWrapper& other) {
|
|
return MjvCameraWrapper(other);
|
|
});
|
|
mjvCamera.def("__deepcopy__", [](const MjvCameraWrapper& other, py::dict) {
|
|
return MjvCameraWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjvCamera);
|
|
#define X(var) \
|
|
mjvCamera.def_property( \
|
|
#var, [](const MjvCameraWrapper& c) { return c.get()->var; }, \
|
|
[](MjvCameraWrapper& c, decltype(raw::MjvCamera::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(type);
|
|
X(fixedcamid);
|
|
X(trackbodyid);
|
|
X(distance);
|
|
X(azimuth);
|
|
X(elevation);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjvCamera, #var, &MjvCameraWrapper::var)
|
|
X(lookat);
|
|
#undef X
|
|
|
|
// ==================== MJVGLCAMERA ==========================================
|
|
py::class_<MjvGLCameraWrapper> mjvGLCamera(m, "MjvGLCamera");
|
|
mjvGLCamera.def(py::init<>());
|
|
mjvGLCamera.def("__copy__", [](const MjvGLCameraWrapper& other) {
|
|
return MjvGLCameraWrapper(other);
|
|
});
|
|
mjvGLCamera.def(
|
|
"__deepcopy__",
|
|
[](const MjvGLCameraWrapper& other, py::dict) {
|
|
return MjvGLCameraWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjvGLCamera);
|
|
#define X(var) \
|
|
mjvGLCamera.def_property( \
|
|
#var, [](const MjvGLCameraWrapper& c) { return c.get()->var; }, \
|
|
[](MjvGLCameraWrapper& c, decltype(raw::MjvGLCamera::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(frustum_center);
|
|
X(frustum_bottom);
|
|
X(frustum_top);
|
|
X(frustum_near);
|
|
X(frustum_far);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjvGLCamera, #var, &MjvGLCameraWrapper::var)
|
|
X(pos);
|
|
X(forward);
|
|
X(up);
|
|
#undef X
|
|
|
|
// ==================== MJVGEOM ==============================================
|
|
py::class_<MjvGeomWrapper> mjvGeom(m, "MjvGeom");
|
|
mjvGeom.def(py::init<>());
|
|
mjvGeom.def("__copy__", [](const MjvGeomWrapper& other) {
|
|
return MjvGeomWrapper(other);
|
|
});
|
|
mjvGeom.def("__deepcopy__", [](const MjvGeomWrapper& other, py::dict) {
|
|
return MjvGeomWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjvGeom);
|
|
#define X(var) \
|
|
mjvGeom.def_property( \
|
|
#var, [](const MjvGeomWrapper& c) { return c.get()->var; }, \
|
|
[](MjvGeomWrapper& c, decltype(raw::MjvGeom::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(type);
|
|
X(dataid);
|
|
X(objtype);
|
|
X(objid);
|
|
X(category);
|
|
X(texid);
|
|
X(texuniform);
|
|
X(texcoord);
|
|
X(segid);
|
|
X(emission);
|
|
X(specular);
|
|
X(shininess);
|
|
X(reflectance);
|
|
X(camdist);
|
|
X(modelrbound);
|
|
X(transparent);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjvGeom, #var, &MjvGeomWrapper::var)
|
|
X(texrepeat);
|
|
X(size);
|
|
X(pos);
|
|
X(mat);
|
|
X(rgba);
|
|
#undef X
|
|
|
|
DefinePyStr(mjvGeom, "label", &raw::MjvGeom::label);
|
|
|
|
// ==================== MJVLIGHT =============================================
|
|
py::class_<MjvLightWrapper> mjvLight(m, "MjvLight");
|
|
mjvLight.def(py::init<>());
|
|
mjvLight.def("__copy__", [](const MjvLightWrapper& other) {
|
|
return MjvLightWrapper(other);
|
|
});
|
|
mjvLight.def("__deepcopy__", [](const MjvLightWrapper& other, py::dict) {
|
|
return MjvLightWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjvLight);
|
|
#define X(var) \
|
|
mjvLight.def_property( \
|
|
#var, [](const MjvLightWrapper& c) { return c.get()->var; }, \
|
|
[](MjvLightWrapper& c, decltype(raw::MjvLight::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(cutoff);
|
|
X(exponent);
|
|
X(headlight);
|
|
X(directional);
|
|
X(castshadow);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjvLight, #var, &MjvLightWrapper::var)
|
|
X(pos);
|
|
X(dir);
|
|
X(attenuation);
|
|
X(ambient);
|
|
X(diffuse);
|
|
X(specular);
|
|
#undef X
|
|
|
|
// ==================== MJVOPTION ============================================
|
|
py::class_<MjvOptionWrapper> mjvOption(m, "MjvOption");
|
|
mjvOption.def(py::init<>());
|
|
mjvOption.def("__copy__", [](const MjvOptionWrapper& other) {
|
|
return MjvOptionWrapper(other);
|
|
});
|
|
mjvOption.def("__deepcopy__", [](const MjvOptionWrapper& other, py::dict) {
|
|
return MjvOptionWrapper(other);
|
|
});
|
|
DefineStructFunctions(mjvOption);
|
|
#define X(var) \
|
|
mjvOption.def_property( \
|
|
#var, [](const MjvOptionWrapper& c) { return c.get()->var; }, \
|
|
[](MjvOptionWrapper& c, decltype(raw::MjvOption::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(label);
|
|
X(frame);
|
|
X(bvh_depth);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjvOption, #var, &MjvOptionWrapper::var)
|
|
X(geomgroup);
|
|
X(sitegroup);
|
|
X(jointgroup);
|
|
X(tendongroup);
|
|
X(actuatorgroup);
|
|
X(skingroup);
|
|
X(flags);
|
|
#undef X
|
|
|
|
// ==================== MJVSCENE =============================================
|
|
py::class_<MjvSceneWrapper> mjvScene(m, "MjvScene");
|
|
mjvScene.def(py::init<>());
|
|
mjvScene.def(py::init<const MjModelWrapper&, int>(),
|
|
py::arg("model"), py::arg("maxgeom"));
|
|
mjvScene.def("__copy__", [](const MjvSceneWrapper& other) {
|
|
return MjvSceneWrapper(other);
|
|
});
|
|
mjvScene.def("__deepcopy__", [](const MjvSceneWrapper& other, py::dict) {
|
|
return MjvSceneWrapper(other);
|
|
});
|
|
#define X(var) \
|
|
mjvScene.def_property( \
|
|
#var, [](const MjvSceneWrapper& c) { return c.get()->var; }, \
|
|
[](MjvSceneWrapper& c, decltype(raw::MjvScene::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(maxgeom);
|
|
X(ngeom);
|
|
X(nskin);
|
|
X(nlight);
|
|
X(enabletransform);
|
|
X(scale);
|
|
X(stereo);
|
|
X(framewidth);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjvScene, #var, &MjvSceneWrapper::var)
|
|
X(geoms);
|
|
X(geomorder);
|
|
X(skinfacenum);
|
|
X(skinvertadr);
|
|
X(skinvertnum);
|
|
X(skinvert);
|
|
X(skinnormal);
|
|
X(lights);
|
|
X(camera);
|
|
X(translate);
|
|
X(rotate);
|
|
X(flags);
|
|
X(framergb);
|
|
#undef X
|
|
|
|
// ==================== MJVFIGURE ============================================
|
|
py::class_<MjvFigureWrapper> mjvFigure(m, "MjvFigure");
|
|
mjvFigure.def(py::init<>());
|
|
mjvFigure.def("__copy__", [](const MjvFigureWrapper& other) {
|
|
return MjvFigureWrapper(other);
|
|
});
|
|
mjvFigure.def("__deepcopy__", [](const MjvFigureWrapper& other, py::dict) {
|
|
return MjvFigureWrapper(other);
|
|
});
|
|
#define X(var) \
|
|
mjvFigure.def_property( \
|
|
#var, [](const MjvFigureWrapper& c) { return c.get()->var; }, \
|
|
[](MjvFigureWrapper& c, decltype(raw::MjvFigure::var) rhs) { \
|
|
c.get()->var = rhs; \
|
|
})
|
|
X(flg_legend);
|
|
X(flg_extend);
|
|
X(flg_barplot);
|
|
X(flg_selection);
|
|
X(flg_symmetric);
|
|
X(linewidth);
|
|
X(gridwidth);
|
|
X(legendoffset);
|
|
X(subplot);
|
|
X(highlightid);
|
|
X(selection);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjvFigure, #var, &MjvFigureWrapper::var)
|
|
X(flg_ticklabel);
|
|
X(gridsize);
|
|
X(gridrgb);
|
|
X(figurergba);
|
|
X(panergba);
|
|
X(legendrgba);
|
|
X(textrgb);
|
|
X(linergb);
|
|
X(range);
|
|
X(highlight);
|
|
X(linepnt);
|
|
X(linedata);
|
|
X(xaxispixel);
|
|
X(yaxispixel);
|
|
X(xaxisdata);
|
|
X(yaxisdata);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyStr(mjvFigure, #var, &raw::MjvFigure::var);
|
|
X(xformat);
|
|
X(yformat);
|
|
X(minwidth);
|
|
X(title);
|
|
X(xlabel);
|
|
#undef X
|
|
|
|
mjvFigure.def_readonly("linename", &MjvFigureWrapper::linename);
|
|
|
|
// mjv_averageCamera returns an mjvGLCamera and we need to call the wrapper's
|
|
// constructor on the return value. Defining the binding for this function
|
|
// in this file to avoid symbol dependency across modules.
|
|
m.def(
|
|
"mjv_averageCamera",
|
|
[](const MjvGLCameraWrapper& cam1, const MjvGLCameraWrapper& cam2) {
|
|
return MjvGLCameraWrapper([&cam1, &cam2]() {
|
|
py::gil_scoped_release no_gil;
|
|
return InterceptMjErrors(mjv_averageCamera)(cam1.get(), cam2.get());
|
|
}());
|
|
},
|
|
py::arg("cam1"), py::arg("cam2"),
|
|
py::doc(python_traits::mjv_averageCamera::doc));
|
|
} // PYBIND11_MODULE NOLINT(readability/fn_size)
|
|
} // namespace mujoco::python::_impl
|