Python bindings: Store a reference to MjModel in MjDataWrapper.
Before this change, MjDataWrapper contained a custom data structure called MjDataMetadata, which contained the sizes of various arrays needed to create MjData. This was used for serialization and deserialization, as well as copying MjData instances. This worked fine for native MuJoCo models, but as soon as plugins were used, there was information in the model that was needed and not available in MjDataMetadata. Since plugins are so general, keeping MjDataMetadata was untenable. PiperOrigin-RevId: 578205355 Change-Id: I11b9ee797d1da5aaf1fecbbe170b3b0f1f1a3e6c
This commit is contained in:
Nimrod Gileadi
committed by
Copybara-Service
Copybara-Service
parent
64a59bb2dc
commit
084facc9ab
@@ -265,7 +265,6 @@ target_sources(
|
||||
structs_header
|
||||
INTERFACE indexer_xmacro.h
|
||||
indexers.h
|
||||
mjdata_meta.h
|
||||
structs.h
|
||||
)
|
||||
set_target_properties(structs_header PROPERTIES PUBLIC_HEADER structs.h)
|
||||
|
||||
@@ -74,10 +74,21 @@ TEST_XML_SENSOR = r"""
|
||||
"""
|
||||
|
||||
TEST_XML_PLUGIN = r"""
|
||||
<mujoco model="test">
|
||||
<mujoco>
|
||||
<option gravity="0 0 0"/>
|
||||
<extension>
|
||||
<plugin plugin="mujoco.elasticity.cable"/>
|
||||
</extension>
|
||||
<worldbody>
|
||||
<composite type="cable" curve="s" count="41 1 1" size="1" offset="0 0 1" initial="none">
|
||||
<plugin plugin="mujoco.elasticity.cable">
|
||||
<config key="twist" value="1e6"/>
|
||||
<config key="bend" value="1e9"/>
|
||||
</plugin>
|
||||
<joint kind="main" damping="2"/>
|
||||
<geom type="capsule" size=".005" density="1"/>
|
||||
</composite>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
@@ -1048,10 +1059,24 @@ Euler integrator, semi-implicit in velocity.
|
||||
while data_instances:
|
||||
d = data_instances.pop()
|
||||
self.assertEqual(sys.getrefcount(d), 2)
|
||||
del d
|
||||
while model_instances:
|
||||
m = model_instances.pop()
|
||||
self.assertEqual(sys.getrefcount(m), 2)
|
||||
|
||||
# This test is disabled on PyPy as it uses sys.getrefcount
|
||||
# However PyPy is not officially supported by MuJoCo
|
||||
@absltest.skipIf(sys.implementation.name == 'pypy',
|
||||
reason='requires sys.getrefcount')
|
||||
def test_mjdata_holds_ref_to_model(self):
|
||||
data = mujoco.MjData(mujoco.MjModel.from_xml_string('<mujoco/>'))
|
||||
model = data.model
|
||||
# references: one in `data.model, one in `model`, one in the temporary ref
|
||||
# passed to getrefcount.
|
||||
self.assertEqual(sys.getrefcount(data.model), 3)
|
||||
del data
|
||||
self.assertEqual(sys.getrefcount(model), 2)
|
||||
|
||||
def test_can_initialize_mjv_structs(self):
|
||||
self.assertIsInstance(mujoco.MjvScene(), mujoco.MjvScene)
|
||||
self.assertIsInstance(mujoco.MjvCamera(), mujoco.MjvCamera)
|
||||
@@ -1287,6 +1312,34 @@ Euler integrator, semi-implicit in velocity.
|
||||
self.assertEqual(data.geom(3).xpos[2], 4)
|
||||
self.assertEqual(data.geom(4).xpos[2], 5)
|
||||
|
||||
def test_load_plugin(self):
|
||||
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
|
||||
data = mujoco.MjData(model)
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
def test_copy_mjdata_with_plugin(self):
|
||||
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
|
||||
data1 = mujoco.MjData(model)
|
||||
self.assertIs(data1.model, model)
|
||||
mujoco.mj_step(model, data1)
|
||||
data2 = copy.copy(data1)
|
||||
mujoco.mj_step(model, data1)
|
||||
mujoco.mj_step(model, data2)
|
||||
np.testing.assert_array_equal(data1.qpos, data2.qpos)
|
||||
self.assertIs(data1.model, data2.model)
|
||||
|
||||
def test_deepcopy_mjdata_with_plugin(self):
|
||||
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
|
||||
data1 = mujoco.MjData(model)
|
||||
self.assertIs(data1.model, model)
|
||||
mujoco.mj_step(model, data1)
|
||||
data2 = copy.deepcopy(data1)
|
||||
mujoco.mj_step(model, data1)
|
||||
mujoco.mj_step(model, data2)
|
||||
np.testing.assert_array_equal(data1.qpos, data2.qpos)
|
||||
self.assertIsNot(data1.model, data2.model)
|
||||
self.assertNotEqual(data1.model._address, data2.model._address)
|
||||
|
||||
def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
|
||||
for name in attr_to_compare:
|
||||
actual_value = getattr(actual_obj, name)
|
||||
@@ -1300,9 +1353,6 @@ Euler integrator, semi-implicit in velocity.
|
||||
self.fail("Attribute '{}' differs from expected value: {}".format(
|
||||
name, str(e)))
|
||||
|
||||
def test_load_plugin(self):
|
||||
mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
|
||||
@@ -23,6 +23,7 @@
|
||||
|
||||
#include <Eigen/Core>
|
||||
#include <mujoco/mjxmacro.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "function_traits.h"
|
||||
#include "functions.h"
|
||||
#include "private.h"
|
||||
@@ -1367,7 +1368,7 @@ PYBIND11_MODULE(_functions, pymodule) {
|
||||
}
|
||||
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) d.metadata().x
|
||||
#define MJ_M(x) d.model().get()->x
|
||||
#undef MJ_D
|
||||
#define MJ_D(x) data->x
|
||||
#define X(type, name, nr, nc) \
|
||||
@@ -1379,7 +1380,7 @@ PYBIND11_MODULE(_functions, pymodule) {
|
||||
}
|
||||
|
||||
MJDATA_ARENA_POINTERS_PRIMAL
|
||||
if (d.metadata().is_dual) {
|
||||
if (mj_isDual(d.model().get())) {
|
||||
MJDATA_ARENA_POINTERS_DUAL
|
||||
}
|
||||
#undef X
|
||||
|
||||
+13
-21
@@ -13,16 +13,13 @@
|
||||
// limitations under the License.
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <utility>
|
||||
#include <variant>
|
||||
#include <vector>
|
||||
|
||||
#include "indexers.h"
|
||||
#include "mjdata_meta.h"
|
||||
#include "raw.h"
|
||||
#include "util/crossplatform.h"
|
||||
|
||||
@@ -76,7 +73,6 @@ IDToName MakeIDToName(int count, IntPtr name_offsets, CharPtr names) {
|
||||
// the MuJoCo category of the entity itself, e.g. nbody indicates that the
|
||||
// field belongs to a body.
|
||||
// T: Scalar data type of the field.
|
||||
// M: Either raw::MjModel or MjDataMetadata.
|
||||
//
|
||||
// Args:
|
||||
// base_ptr: Pointer to the first entry in the entire field.
|
||||
@@ -87,35 +83,35 @@ IDToName MakeIDToName(int count, IntPtr name_offsets, CharPtr names) {
|
||||
// additional dimension of the size len(qvel) of the particular joint.
|
||||
// m: Used for dereferencing MjSize.
|
||||
// owner: The base object whose lifetime is tied to the returned array.
|
||||
template <auto MjSize, typename T, typename M>
|
||||
template <auto MjSize, typename T>
|
||||
py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
|
||||
const M& m, py::handle owner) {
|
||||
const raw::MjModel& m, py::handle owner) {
|
||||
int offset;
|
||||
if (MjSize == &M::nq) {
|
||||
if (MjSize == &raw::MjModel::nq) {
|
||||
offset = m.jnt_qposadr[index];
|
||||
shape.insert(
|
||||
shape.begin(),
|
||||
((index < m.njnt-1) ? m.jnt_qposadr[index+1] : m.nq) - offset);
|
||||
} else if (MjSize == &M::nv) {
|
||||
} else if (MjSize == &raw::MjModel::nv) {
|
||||
offset = m.jnt_dofadr[index];
|
||||
shape.insert(
|
||||
shape.begin(),
|
||||
((index < m.njnt-1) ? m.jnt_dofadr[index+1] : m.nv) - offset);
|
||||
} else if (MjSize == &M::nhfielddata) {
|
||||
} else if (MjSize == &raw::MjModel::nhfielddata) {
|
||||
offset = m.hfield_adr[index];
|
||||
shape.insert(shape.begin(), m.hfield_ncol[index]);
|
||||
shape.insert(shape.begin(), m.hfield_nrow[index]);
|
||||
} else if (MjSize == &M::ntexdata) {
|
||||
} else if (MjSize == &raw::MjModel::ntexdata) {
|
||||
offset = m.tex_adr[index];
|
||||
shape.insert(shape.begin(), m.tex_width[index]);
|
||||
shape.insert(shape.begin(), m.tex_height[index]);
|
||||
} else if (MjSize == &M::nsensordata) {
|
||||
} else if (MjSize == &raw::MjModel::nsensordata) {
|
||||
offset = m.sensor_adr[index];
|
||||
shape.insert(shape.begin(), m.sensor_dim[index]);
|
||||
} else if (MjSize == &M::nnumericdata) {
|
||||
} else if (MjSize == &raw::MjModel::nnumericdata) {
|
||||
offset = m.numeric_adr[index];
|
||||
shape.insert(shape.begin(), m.numeric_size[index]);
|
||||
} else if (MjSize == &M::ntupledata) {
|
||||
} else if (MjSize == &raw::MjModel::ntupledata) {
|
||||
offset = m.tuple_adr[index];
|
||||
shape.insert(shape.begin(), m.tuple_size[index]);
|
||||
} else {
|
||||
@@ -135,9 +131,7 @@ py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// M is either a raw::MjModel or MjDataMetadata.
|
||||
template <typename M>
|
||||
NameToIDMappings::NameToIDMappings(const M& m)
|
||||
NameToIDMappings::NameToIDMappings(const raw::MjModel& m)
|
||||
: body(MakeNameToID(m.nbody, m.name_bodyadr, m.names)),
|
||||
jnt(MakeNameToID(m.njnt, m.name_jntadr, m.names)),
|
||||
geom(MakeNameToID(m.ngeom, m.name_geomadr, m.names)),
|
||||
@@ -160,9 +154,7 @@ NameToIDMappings::NameToIDMappings(const M& m)
|
||||
tuple(MakeNameToID(m.ntuple, m.name_tupleadr, m.names)),
|
||||
key(MakeNameToID(m.nkey, m.name_keyadr, m.names)) {}
|
||||
|
||||
// M is either a raw::MjModel or MjDataMetadata.
|
||||
template <typename M>
|
||||
IDToNameMappings::IDToNameMappings(const M& m)
|
||||
IDToNameMappings::IDToNameMappings(const raw::MjModel& m)
|
||||
: body(MakeIDToName(m.nbody, m.name_bodyadr, m.names)),
|
||||
jnt(MakeIDToName(m.njnt, m.name_jntadr, m.names)),
|
||||
geom(MakeIDToName(m.ngeom, m.name_geomadr, m.names)),
|
||||
@@ -223,7 +215,7 @@ MJMODEL_VIEW_GROUPS
|
||||
MJMODEL_VIEW_GROUPS
|
||||
#undef XGROUP
|
||||
|
||||
MjDataIndexer::MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
|
||||
MjDataIndexer::MjDataIndexer(raw::MjData* d, const raw::MjModel* m,
|
||||
py::handle owner)
|
||||
: d_(d),
|
||||
m_(m),
|
||||
@@ -369,7 +361,7 @@ MJMODEL_KEYFRAME
|
||||
#define X(type, prefix, var, dim0, dim1) \
|
||||
py::array_t<type> XGROUP::var() { \
|
||||
if (!var##_.has_value()) { \
|
||||
var##_.emplace(MakeArray<&MjDataMetadata::dim0>( \
|
||||
var##_.emplace(MakeArray<&raw::MjModel::dim0>( \
|
||||
d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
|
||||
} \
|
||||
return *var##_; \
|
||||
|
||||
@@ -25,7 +25,6 @@
|
||||
#include <absl/container/flat_hash_map.h>
|
||||
#include <mujoco/mjxmacro.h>
|
||||
#include "indexer_xmacro.h"
|
||||
#include "mjdata_meta.h"
|
||||
#include "raw.h"
|
||||
#include <pybind11/numpy.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
@@ -35,9 +34,7 @@ using NameToID = absl::flat_hash_map<std::string, int>;
|
||||
using IDToName = std::vector<std::string>;
|
||||
|
||||
struct NameToIDMappings {
|
||||
// M is either a raw::MjModel or MjDataMetadata.
|
||||
template <typename M>
|
||||
explicit NameToIDMappings(const M& m);
|
||||
explicit NameToIDMappings(const raw::MjModel& m);
|
||||
|
||||
NameToID body;
|
||||
NameToID jnt;
|
||||
@@ -63,9 +60,7 @@ struct NameToIDMappings {
|
||||
};
|
||||
|
||||
struct IDToNameMappings {
|
||||
// M is either a raw::MjModel or MjDataMetadata.
|
||||
template <typename M>
|
||||
explicit IDToNameMappings(const M& m);
|
||||
explicit IDToNameMappings(const raw::MjModel& m);
|
||||
|
||||
IDToName body;
|
||||
IDToName jnt;
|
||||
@@ -163,7 +158,7 @@ class MjModelIndexer {
|
||||
class MjDataGroupedViewsBase {
|
||||
public:
|
||||
MjDataGroupedViewsBase(int index, std::string_view name, raw::MjData* d,
|
||||
const MjDataMetadata* m,
|
||||
const raw::MjModel* m,
|
||||
pybind11::handle owner)
|
||||
: index_(index), name_(name), d_(d), m_(m), owner_(owner) {}
|
||||
|
||||
@@ -175,7 +170,7 @@ class MjDataGroupedViewsBase {
|
||||
int index_;
|
||||
std::string name_;
|
||||
raw::MjData* d_;
|
||||
const MjDataMetadata* m_;
|
||||
const raw::MjModel* m_;
|
||||
pybind11::handle owner_;
|
||||
};
|
||||
|
||||
@@ -201,7 +196,7 @@ MJDATA_VIEW_GROUPS
|
||||
// (e.g. a particular geom or joint) either by name or by ID.
|
||||
class MjDataIndexer {
|
||||
public:
|
||||
MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
|
||||
MjDataIndexer(raw::MjData* d, const raw::MjModel* m,
|
||||
pybind11::handle owner);
|
||||
|
||||
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
|
||||
@@ -213,7 +208,7 @@ class MjDataIndexer {
|
||||
|
||||
private:
|
||||
raw::MjData* d_;
|
||||
const MjDataMetadata* m_;
|
||||
const raw::MjModel* m_;
|
||||
pybind11::handle owner_;
|
||||
NameToIDMappings name_to_id_;
|
||||
IDToNameMappings id_to_name_;
|
||||
|
||||
@@ -1,108 +0,0 @@
|
||||
// Copyright 2021 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.
|
||||
|
||||
#ifndef MUJOCO_PYTHON_MJDATA_META_H_
|
||||
#define MUJOCO_PYTHON_MJDATA_META_H_
|
||||
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
|
||||
#include <mujoco/mujoco.h>
|
||||
#include <mujoco/mjxmacro.h>
|
||||
#include "raw.h"
|
||||
#include "util/crossplatform.h"
|
||||
|
||||
namespace mujoco::python {
|
||||
namespace _impl {
|
||||
template <typename T> class MjWrapper;
|
||||
} // namespace _impl
|
||||
|
||||
#define MJDATA_METADATA \
|
||||
X( int, jnt_qposadr, njnt ) \
|
||||
X( int, jnt_dofadr, njnt ) \
|
||||
X( int, hfield_nrow, nhfield ) \
|
||||
X( int, hfield_ncol, nhfield ) \
|
||||
X( int, hfield_adr, nhfield ) \
|
||||
X( int, tex_height, ntex ) \
|
||||
X( int, tex_width, ntex ) \
|
||||
X( int, tex_adr, ntex ) \
|
||||
X( int, sensor_dim, nsensor ) \
|
||||
X( int, sensor_adr, nsensor ) \
|
||||
X( int, numeric_adr, nnumeric ) \
|
||||
X( int, numeric_size, nnumeric ) \
|
||||
X( int, tuple_adr, ntuple ) \
|
||||
X( int, tuple_size, ntuple ) \
|
||||
X( int, name_bodyadr, nbody ) \
|
||||
X( int, name_jntadr, njnt ) \
|
||||
X( int, name_geomadr, ngeom ) \
|
||||
X( int, name_siteadr, nsite ) \
|
||||
X( int, name_camadr, ncam ) \
|
||||
X( int, name_lightadr, nlight ) \
|
||||
X( int, name_meshadr, nmesh ) \
|
||||
X( int, name_skinadr, nskin ) \
|
||||
X( int, name_hfieldadr, nhfield ) \
|
||||
X( int, name_texadr, ntex ) \
|
||||
X( int, name_matadr, nmat ) \
|
||||
X( int, name_pairadr, npair ) \
|
||||
X( int, name_excludeadr, nexclude ) \
|
||||
X( int, name_eqadr, neq ) \
|
||||
X( int, name_tendonadr, ntendon ) \
|
||||
X( int, name_actuatoradr, nu ) \
|
||||
X( int, name_sensoradr, nsensor ) \
|
||||
X( int, name_numericadr, nnumeric ) \
|
||||
X( int, name_textadr, ntext ) \
|
||||
X( int, name_tupleadr, ntuple ) \
|
||||
X( int, name_keyadr, nkey ) \
|
||||
X( char, names, nnames )
|
||||
|
||||
// A subset of mjModel fields that are required to reconstruct an MjDataWrapper.
|
||||
struct MjDataMetadata {
|
||||
public:
|
||||
friend class _impl::MjWrapper<raw::MjData>;
|
||||
explicit MjDataMetadata(const raw::MjModel* m)
|
||||
:
|
||||
#define X(var) var(m->var),
|
||||
MJMODEL_INTS
|
||||
#undef X
|
||||
|
||||
#define X(dtype, var, n) \
|
||||
var([](dtype* src, int len) { \
|
||||
dtype* dst = new dtype[len]; \
|
||||
std::memcpy(dst, src, len * sizeof(dtype)); \
|
||||
return dst; \
|
||||
}(m->var, m->n)),
|
||||
|
||||
MJDATA_METADATA
|
||||
#undef X
|
||||
is_dual(mj_isDual(m)) {
|
||||
}
|
||||
|
||||
#define X(var) decltype(raw::MjModel::var) var;
|
||||
MJMODEL_INTS
|
||||
#undef X
|
||||
#define X(type, var, n) std::shared_ptr<type[]> var;
|
||||
MJDATA_METADATA
|
||||
#undef X
|
||||
|
||||
bool is_dual;
|
||||
|
||||
private:
|
||||
MjDataMetadata() = default;
|
||||
MjDataMetadata(const MjDataMetadata& other) = default;
|
||||
MjDataMetadata(MjDataMetadata&& other) = default;
|
||||
};
|
||||
|
||||
} // namespace mujoco::python
|
||||
|
||||
#endif // MUJOCO_PYTHON_MJDATA_META_H_
|
||||
+82
-75
@@ -37,10 +37,11 @@
|
||||
#include "errors.h"
|
||||
#include "function_traits.h"
|
||||
#include "indexers.h"
|
||||
#include "mjdata_meta.h"
|
||||
#include "private.h"
|
||||
#include "raw.h"
|
||||
#include "serialization.h"
|
||||
#include <pybind11/cast.h>
|
||||
#include <pybind11/detail/common.h>
|
||||
#include <pybind11/numpy.h>
|
||||
#include <pybind11/operators.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
@@ -453,7 +454,8 @@ void MjModelWrapper::Serialize(std::ostream& output) const {
|
||||
WriteBytes(output, buffer.data(), model_size);
|
||||
}
|
||||
|
||||
MjModelWrapper MjModelWrapper::Deserialize(std::istream& input) {
|
||||
std::unique_ptr<MjModelWrapper> MjModelWrapper::Deserialize(
|
||||
std::istream& input) {
|
||||
CheckInput(input, "mjModel");
|
||||
|
||||
char serializationVersion = ReadChar(input);
|
||||
@@ -479,7 +481,7 @@ MjModelWrapper MjModelWrapper::Deserialize(std::istream& input) {
|
||||
if (!model) {
|
||||
throw py::value_error("Invalid serialized mjModel.");
|
||||
}
|
||||
return MjModelWrapper(model);
|
||||
return std::unique_ptr<MjModelWrapper>(new MjModelWrapper(model));
|
||||
}
|
||||
|
||||
// ==================== MJCONTACT ==============================================
|
||||
@@ -554,13 +556,13 @@ MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept {
|
||||
}
|
||||
}
|
||||
|
||||
MjDataWrapper::MjWrapper(const MjModelWrapper& model)
|
||||
: WrapperBase(InterceptMjErrors(mj_makeData)(model.get()),
|
||||
MjDataWrapper::MjWrapper(MjModelWrapper* model)
|
||||
: WrapperBase(InterceptMjErrors(mj_makeData)(model->get()),
|
||||
&MjDataCapsuleDestructor),
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) model.get()->x
|
||||
#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_)),
|
||||
var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)),
|
||||
MJDATA_POINTERS
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) (x)
|
||||
@@ -571,8 +573,9 @@ MjDataWrapper::MjWrapper(const MjModelWrapper& model)
|
||||
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
|
||||
MJDATA_VECTOR
|
||||
#undef X
|
||||
metadata_(model.get()),
|
||||
indexer_(ptr_, &metadata_, owner_) {
|
||||
model_(model),
|
||||
model_ref_(py::cast(model_)),
|
||||
indexer_(ptr_, model_->get(), owner_) {
|
||||
bool is_newly_inserted = false;
|
||||
{
|
||||
py::gil_scoped_acquire gil;
|
||||
@@ -587,9 +590,9 @@ MjDataWrapper::MjWrapper(const MjModelWrapper& model)
|
||||
MjDataWrapper::MjWrapper(const MjDataWrapper& other)
|
||||
: WrapperBase(other.Copy(), &MjDataCapsuleDestructor),
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) other.metadata_.x
|
||||
#define MJ_M(x) other.model_->get()->x
|
||||
#define X(dtype, var, dim0, dim1) \
|
||||
var(InitPyArray(X_ARRAY_SHAPE(other.metadata_.dim0, dim1), ptr_->var, \
|
||||
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
|
||||
owner_)),
|
||||
MJDATA_POINTERS
|
||||
#undef MJ_M
|
||||
@@ -601,8 +604,9 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other)
|
||||
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
|
||||
MJDATA_VECTOR
|
||||
#undef X
|
||||
metadata_(other.metadata_),
|
||||
indexer_(ptr_, &metadata_, owner_) {
|
||||
model_(other.model_),
|
||||
model_ref_(other.model_ref_),
|
||||
indexer_(ptr_, model_->get(), owner_) {
|
||||
bool is_newly_inserted = false;
|
||||
{
|
||||
py::gil_scoped_acquire gil;
|
||||
@@ -617,9 +621,9 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other)
|
||||
MjDataWrapper::MjWrapper(MjDataWrapper&& other)
|
||||
: WrapperBase(other.ptr_, other.owner_),
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) other.metadata_.x
|
||||
#define MJ_M(x) other.model_->get()->x
|
||||
#define X(dtype, var, dim0, dim1) \
|
||||
var(InitPyArray(X_ARRAY_SHAPE(other.metadata_.dim0, dim1), ptr_->var, \
|
||||
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
|
||||
owner_)),
|
||||
MJDATA_POINTERS
|
||||
#undef MJ_M
|
||||
@@ -631,8 +635,9 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other)
|
||||
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
|
||||
MJDATA_VECTOR
|
||||
#undef X
|
||||
metadata_(other.metadata_),
|
||||
indexer_(ptr_, &metadata_, owner_) {
|
||||
model_(other.model_),
|
||||
model_ref_(std::move(other.model_ref_)),
|
||||
indexer_(ptr_, model_->get(), owner_) {
|
||||
bool is_newly_inserted = false;
|
||||
{
|
||||
py::gil_scoped_acquire gil;
|
||||
@@ -646,12 +651,13 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other)
|
||||
other.ptr_ = nullptr;
|
||||
}
|
||||
|
||||
MjDataWrapper::MjWrapper(MjDataMetadata&& metadata, raw::MjData* d)
|
||||
: WrapperBase(d, &MjDataCapsuleDestructor),
|
||||
MjDataWrapper::MjWrapper(const MjDataWrapper& other, MjModelWrapper* model)
|
||||
: WrapperBase(other.Copy(), &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_)),
|
||||
#define MJ_M(x) other.model_->get()->x
|
||||
#define X(dtype, var, dim0, dim1) \
|
||||
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
|
||||
owner_)),
|
||||
MJDATA_POINTERS
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) (x)
|
||||
@@ -662,8 +668,39 @@ MjDataWrapper::MjWrapper(MjDataMetadata&& metadata, raw::MjData* d)
|
||||
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
|
||||
MJDATA_VECTOR
|
||||
#undef X
|
||||
metadata_(std::move(metadata)),
|
||||
indexer_(ptr_, &metadata_, owner_) {
|
||||
model_(model),
|
||||
model_ref_(py::cast(model_)),
|
||||
indexer_(ptr_, model_->get(), 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(MjModelWrapper* model, raw::MjData* d)
|
||||
: WrapperBase(d, &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
|
||||
model_(model),
|
||||
model_ref_(py::cast(model_)),
|
||||
indexer_(ptr_, model_->get(), owner_) {
|
||||
bool is_newly_inserted = false;
|
||||
{
|
||||
py::gil_scoped_acquire gil;
|
||||
@@ -694,19 +731,7 @@ 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
|
||||
model_->Serialize(output);
|
||||
|
||||
// Write struct and scalar fields
|
||||
#define X(var) WriteBytes(output, &ptr_->var, sizeof(ptr_->var))
|
||||
@@ -729,15 +754,15 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
|
||||
|
||||
// Write buffer contents
|
||||
{
|
||||
MJDATA_POINTERS_PREAMBLE((&this->metadata_))
|
||||
MJDATA_POINTERS_PREAMBLE((this->model_->get()))
|
||||
|
||||
#define X(type, name, nr, nc) \
|
||||
WriteBytes(output, ptr_->name, sizeof(type)*(this->metadata_.nr)*(nc));
|
||||
WriteBytes(output, ptr_->name, sizeof(type)*(this->model_->get()->nr)*(nc));
|
||||
MJDATA_POINTERS
|
||||
#undef X
|
||||
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) this->metadata_.x
|
||||
#define MJ_M(x) this->model_->get()->x
|
||||
#undef MJ_D
|
||||
#define MJ_D(x) this->ptr_->x
|
||||
#define X(type, name, nr, nc) \
|
||||
@@ -747,7 +772,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
|
||||
|
||||
MJDATA_ARENA_POINTERS_CONTACT
|
||||
MJDATA_ARENA_POINTERS_PRIMAL
|
||||
if (this->metadata_.is_dual) {
|
||||
if (mj_isDual(this->model_->get())) {
|
||||
MJDATA_ARENA_POINTERS_DUAL
|
||||
}
|
||||
if (this->ptr_->nisland) {
|
||||
@@ -769,28 +794,12 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
|
||||
throw py::value_error("Incompatible serialization version.");
|
||||
}
|
||||
|
||||
// Read all size and address fields
|
||||
MjDataMetadata metadata;
|
||||
raw::MjModel m{0};
|
||||
// Read the model that was used to create the mjData.
|
||||
std::unique_ptr<MjModelWrapper> m_wrapper =
|
||||
MjModelWrapper::Deserialize(input);
|
||||
raw::MjModel& m = *m_wrapper->get();
|
||||
|
||||
#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
|
||||
bool is_dual = mj_isDual(&m);
|
||||
|
||||
raw::MjData* d = mj_makeData(&m);
|
||||
if (!d) {
|
||||
@@ -841,7 +850,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
|
||||
|
||||
MJDATA_ARENA_POINTERS_CONTACT
|
||||
MJDATA_ARENA_POINTERS_PRIMAL
|
||||
if (metadata.is_dual) {
|
||||
if (is_dual) {
|
||||
MJDATA_ARENA_POINTERS_DUAL
|
||||
}
|
||||
if (d->nisland) {
|
||||
@@ -861,18 +870,12 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
|
||||
throw py::value_error("Invalid serialized mjData.");
|
||||
}
|
||||
|
||||
return MjDataWrapper(std::move(metadata), d);
|
||||
return MjDataWrapper(m_wrapper.release(), 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_);
|
||||
const raw::MjModel* m = model_->get();
|
||||
return InterceptMjErrors(mj_copyData)(NULL, m, this->ptr_);
|
||||
}
|
||||
|
||||
// ==================== MJSTATISTIC ============================================
|
||||
@@ -1890,15 +1893,19 @@ This is useful for example when the MJB is not available as a file on disk.)"));
|
||||
|
||||
// ==================== MJDATA ===============================================
|
||||
py::class_<MjDataWrapper> mjData(m, "MjData");
|
||||
mjData.def(py::init<const MjModelWrapper&>());
|
||||
mjData.def(py::init<MjModelWrapper*>());
|
||||
mjData.def_property_readonly("_address", [](const MjDataWrapper& d) {
|
||||
return reinterpret_cast<std::uintptr_t>(d.get());
|
||||
});
|
||||
mjData.def_property_readonly("model", [](const MjDataWrapper& d) {
|
||||
return &d.model();
|
||||
});
|
||||
mjData.def("__copy__", [](const MjDataWrapper& other) {
|
||||
return MjDataWrapper(other);
|
||||
});
|
||||
mjData.def("__deepcopy__", [](const MjDataWrapper& other, py::dict) {
|
||||
return MjDataWrapper(other);
|
||||
MjModelWrapper* model_copy = new MjModelWrapper(other.model());
|
||||
return MjDataWrapper(other, model_copy);
|
||||
});
|
||||
mjData.def(py::pickle(
|
||||
[](const MjDataWrapper& d) { // __getstate__
|
||||
@@ -1927,7 +1934,7 @@ This is useful for example when the MJB is not available as a file on disk.)"));
|
||||
#undef X
|
||||
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) d.metadata().x
|
||||
#define MJ_M(x) d.model().get()->x
|
||||
#undef MJ_D
|
||||
#define MJ_D(x) d.get()->x
|
||||
#define X(dtype, var, dim0, dim1) \
|
||||
|
||||
+13
-6
@@ -29,10 +29,10 @@
|
||||
#include <mujoco/mujoco.h>
|
||||
#include <mujoco/mjxmacro.h>
|
||||
#include "indexers.h"
|
||||
#include "mjdata_meta.h"
|
||||
#include "raw.h"
|
||||
#include <pybind11/numpy.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/pytypes.h>
|
||||
|
||||
namespace mujoco::python {
|
||||
namespace _impl {
|
||||
@@ -447,7 +447,8 @@ class MjWrapper<raw::MjModel> : public WrapperBase<raw::MjModel> {
|
||||
MjModelIndexer& indexer() { return indexer_; }
|
||||
|
||||
void Serialize(std::ostream& output) const;
|
||||
static MjWrapper<raw::MjModel> Deserialize(std::istream& input);
|
||||
static std::unique_ptr<MjWrapper<raw::MjModel>> Deserialize(
|
||||
std::istream& input);
|
||||
|
||||
static MjWrapper LoadXMLFile(
|
||||
const std::string& filename,
|
||||
@@ -564,12 +565,14 @@ struct is_mj_struct_list<raw::MjContact> {
|
||||
template <>
|
||||
class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
|
||||
public:
|
||||
explicit MjWrapper(const MjModelWrapper& model);
|
||||
explicit MjWrapper(MjModelWrapper* model);
|
||||
MjWrapper(const MjWrapper& other);
|
||||
MjWrapper(MjWrapper&&);
|
||||
// Used for deepcopy
|
||||
MjWrapper(const MjWrapper& other, MjModelWrapper* model);
|
||||
~MjWrapper();
|
||||
|
||||
const MjDataMetadata& metadata() const { return metadata_; }
|
||||
const MjModelWrapper& model() const { return *model_; }
|
||||
MjDataIndexer& indexer() { return indexer_; }
|
||||
|
||||
void Serialize(std::ostream& output) const;
|
||||
@@ -598,10 +601,14 @@ class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
|
||||
protected:
|
||||
// Internal constructor which takes ownership of given mjData pointer.
|
||||
// Used for deserialization.
|
||||
explicit MjWrapper(MjDataMetadata&& metadata, raw::MjData* d);
|
||||
explicit MjWrapper(MjModelWrapper* model, raw::MjData* d);
|
||||
raw::MjData* Copy() const;
|
||||
|
||||
MjDataMetadata metadata_;
|
||||
// A reference to the model that was used to create this mjData.
|
||||
MjModelWrapper* model_;
|
||||
// A py::object pointing to the same model as model_, to make sure Python
|
||||
// doesn't doesn't garbage collect it until this mjData is released.
|
||||
pybind11::object model_ref_;
|
||||
MjDataIndexer indexer_;
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user