Dynamically allocate contact and efc_ arrays on a new memory arena.
- Add private function `mj_arenaAlloc`. This is used internally to allocate memory from the arena. - Add private function `mj_nefc` to count constraints. This function returns a tight upper bound on `d->nefc`. The number of counted constraints can be slightly bigger than exact `d->nefc` in the case of constraints with empty Jacobian, as when placing a frictional tendon between two world sites. - Add new `memory` attribute to the `size` XML element for specification of arena memory size. This attribute is mutually exclusive with `nstack` and `njmax` specifications, which are now deprecated (but left around for the time being for legacy compatibility). - Move `d->stack` to the end of the new arena space. The stack now grows in reverse from the end. PiperOrigin-RevId: 479341539 Change-Id: Ie019c202e0908577ffc6f833a37920858116f667
This commit is contained in:
committed by
Copybara-Service
parent
4d85a464cc
commit
58fd72f53d
@@ -457,18 +457,22 @@ class MuJoCoBindingsTest(parameterized.TestCase):
|
||||
|
||||
# Grab a reference to the contacts upfront so that we know that they're
|
||||
# a view into mjData rather than a copy.
|
||||
contact = self.data.contact[:4]
|
||||
contact = self.data.contact
|
||||
|
||||
self.model.opt.timestep = 2**-9 # 0.001953125; allows exact comparisons
|
||||
self.assertEqual(self.data.time, 0)
|
||||
while self.data.time < expected_contact_time:
|
||||
self.assertEqual(self.data.ncon, 0)
|
||||
self.assertEmpty(self.data.efc_type)
|
||||
self.assertTrue(self.data.efc_type.flags['OWNDATA'])
|
||||
prev_time = self.data.time
|
||||
mujoco.mj_step(self.model, self.data)
|
||||
self.assertEqual(self.data.time, prev_time + self.model.opt.timestep)
|
||||
|
||||
mujoco.mj_forward(self.model, self.data)
|
||||
self.assertEqual(self.data.ncon, 4)
|
||||
self.assertLen(self.data.efc_type, 16)
|
||||
self.assertFalse(self.data.efc_type.flags['OWNDATA'])
|
||||
|
||||
# Sort contacts in anticlockwise order
|
||||
sorted_contact = sorted(
|
||||
@@ -478,6 +482,11 @@ class MuJoCoBindingsTest(parameterized.TestCase):
|
||||
np.testing.assert_allclose(sorted_contact[2].pos[:2], [0.1, 0.1])
|
||||
np.testing.assert_allclose(sorted_contact[3].pos[:2], [-0.1, 0.1])
|
||||
|
||||
mujoco.mj_resetData(self.model, self.data)
|
||||
self.assertEqual(self.data.ncon, 0)
|
||||
self.assertEmpty(self.data.efc_type)
|
||||
self.assertTrue(self.data.efc_type.flags['OWNDATA'])
|
||||
|
||||
def test_mj_step_multiple(self):
|
||||
self.model.opt.timestep = 2**-9 # 0.001953125; allows exact comparisons
|
||||
self.assertEqual(self.data.time, 0)
|
||||
@@ -488,24 +497,31 @@ class MuJoCoBindingsTest(parameterized.TestCase):
|
||||
self.assertIn('Optionally, repeat nstep times.', mujoco.mj_step.__doc__)
|
||||
|
||||
def test_mj_contact_list(self):
|
||||
self.assertLen(self.data.contact, self.model.nconmax)
|
||||
self.assertEmpty(self.data.contact)
|
||||
|
||||
expected_ncon = 1234
|
||||
self.data.ncon = expected_ncon
|
||||
self.assertLen(self.data.contact, expected_ncon)
|
||||
|
||||
expected_pos = []
|
||||
for contact in self.data.contact:
|
||||
expected_pos.append(np.random.uniform(size=3))
|
||||
contact.pos = expected_pos[-1]
|
||||
self.assertLen(expected_pos, expected_ncon)
|
||||
np.testing.assert_array_equal(self.data.contact.pos, expected_pos)
|
||||
|
||||
expected_friction = []
|
||||
for contact in self.data.contact:
|
||||
expected_friction.append(np.random.uniform(size=5))
|
||||
contact.friction = expected_friction[-1]
|
||||
self.assertLen(expected_friction, expected_ncon)
|
||||
np.testing.assert_array_equal(self.data.contact.friction, expected_friction)
|
||||
|
||||
expected_H = [] # pylint: disable=invalid-name
|
||||
for contact in self.data.contact:
|
||||
expected_H.append(np.random.uniform(size=36))
|
||||
contact.H = expected_H[-1]
|
||||
self.assertLen(expected_H, expected_ncon)
|
||||
np.testing.assert_array_equal(self.data.contact.H, expected_H)
|
||||
|
||||
def test_mj_struct_list_equality(self):
|
||||
@@ -516,16 +532,16 @@ class MuJoCoBindingsTest(parameterized.TestCase):
|
||||
self.assertEqual(self.data.ncon, 4)
|
||||
mujoco.mj_forward(model2, data2)
|
||||
self.assertEqual(data2.ncon, 4)
|
||||
self.assertEqual(data2.contact[:4], self.data.contact[:4])
|
||||
self.assertEqual(data2.contact, self.data.contact)
|
||||
|
||||
self.data.qpos[3:7] = [np.cos(np.pi/8), np.sin(np.pi/8), 0, 0]
|
||||
self.data.qpos[2] *= (np.sqrt(2) - 1) * 0.1 - 1e-6
|
||||
mujoco.mj_forward(self.model, self.data)
|
||||
self.assertEqual(self.data.ncon, 2)
|
||||
self.assertNotEqual(data2.contact[:2], self.data.contact[:2])
|
||||
self.assertNotEqual(data2.contact, self.data.contact)
|
||||
|
||||
# Check that we can compare slices of different lengths
|
||||
self.assertNotEqual(data2.contact[:2], self.data.contact[:4])
|
||||
self.assertNotEqual(data2.contact, self.data.contact)
|
||||
|
||||
# Check that comparing things of different types do not raise an error
|
||||
self.assertNotEqual(self.data.contact, self.data.warning)
|
||||
@@ -856,7 +872,7 @@ Euler integrator, semi-implicit in velocity.
|
||||
|
||||
def test_can_raise_error(self):
|
||||
self.data.pstack = self.data.nstack
|
||||
with self.assertRaisesWithLiteralMatch(mujoco.FatalError, 'Stack overflow'):
|
||||
with self.assertRaisesRegex(mujoco.FatalError, r'\Astack overflow'):
|
||||
mujoco.mj_forward(self.model, self.data)
|
||||
|
||||
def test_mjcb_time(self):
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
#ifndef MUJOCO_PYTHON_MJDATA_META_H_
|
||||
#define MUJOCO_PYTHON_MJDATA_META_H_
|
||||
|
||||
#include <mujoco/mujoco.h>
|
||||
#include <mujoco/mjxmacro.h>
|
||||
#include "raw.h"
|
||||
#include "util/crossplatform.h"
|
||||
@@ -74,6 +75,8 @@ struct MjDataMetadata {
|
||||
MJDATA_METADATA
|
||||
#undef X
|
||||
|
||||
bool is_dual;
|
||||
|
||||
private:
|
||||
MjDataMetadata() = default;
|
||||
MjDataMetadata(const MjDataMetadata& other) = default;
|
||||
@@ -94,10 +97,7 @@ struct MjDataMetadata {
|
||||
|
||||
MJDATA_METADATA
|
||||
#undef X
|
||||
dummy_() {}
|
||||
|
||||
// Dummy variable to terminate X macro sequences.
|
||||
MUJOCO_MAYBE_UNUSED bool dummy_;
|
||||
is_dual(mj_isDual(m)) {}
|
||||
};
|
||||
|
||||
} // namespace mujoco::python
|
||||
|
||||
@@ -41,25 +41,26 @@ inline char ReadChar(std::istream& input) {
|
||||
return c;
|
||||
}
|
||||
|
||||
inline void WriteInt(std::ostream& output, int i) {
|
||||
output.write(reinterpret_cast<char*>(&i), sizeof(int));
|
||||
inline void WriteInt(std::ostream& output, std::size_t i) {
|
||||
output.write(reinterpret_cast<char*>(&i), sizeof(std::size_t));
|
||||
}
|
||||
|
||||
inline int ReadInt(std::istream& input) {
|
||||
int i = 0;
|
||||
input.read(reinterpret_cast<char*>(&i), sizeof(int));
|
||||
inline std::size_t ReadInt(std::istream& input) {
|
||||
std::size_t i = 0;
|
||||
input.read(reinterpret_cast<char*>(&i), sizeof(std::size_t));
|
||||
return i;
|
||||
}
|
||||
|
||||
inline void WriteBytes(std::ostream& output, const void* src, size_t nbytes) {
|
||||
inline void WriteBytes(std::ostream& output, const void* src,
|
||||
std::size_t nbytes) {
|
||||
// Start by writing nbytes itself, so it can be validated at the time of
|
||||
// reading.
|
||||
WriteInt(output, nbytes);
|
||||
output.write(reinterpret_cast<const char*>(src), nbytes);
|
||||
}
|
||||
|
||||
inline void ReadBytes(std::istream& input, void* dest, size_t nbytes) {
|
||||
size_t actual_nbytes = ReadInt(input);
|
||||
inline void ReadBytes(std::istream& input, void* dest, std::size_t nbytes) {
|
||||
std::size_t actual_nbytes = ReadInt(input);
|
||||
if (actual_nbytes != nbytes) {
|
||||
input.setstate(input.rdstate() | std::ios_base::failbit);
|
||||
return;
|
||||
|
||||
+134
-52
@@ -52,6 +52,9 @@ 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
|
||||
@@ -74,6 +77,9 @@ constexpr auto XArrayShapeImpl(const std::string_view dim1_str) {
|
||||
}
|
||||
}
|
||||
|
||||
inline std::size_t NConMax(const mjData* d) {
|
||||
return d->nstack * sizeof(mjtNum) / sizeof(mjContact);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// ==================== MJOPTION ===============================================
|
||||
@@ -498,51 +504,15 @@ MjContactWrapper::MjWrapper(const MjContactWrapper& other)
|
||||
*this->ptr_ = *other.ptr_;
|
||||
}
|
||||
|
||||
#define X(type, var) \
|
||||
var(std::vector<int>{num}, std::vector<int>{sizeof(raw::MjContact)}, \
|
||||
&ptr->var, owner)
|
||||
#define XN(type, var) \
|
||||
var(std::vector<int>{num, sizeof(raw::MjContact::var) / sizeof(type)}, \
|
||||
std::vector<int>{sizeof(raw::MjContact), sizeof(type)}, &ptr->var[0], \
|
||||
owner)
|
||||
MjContactList::MjStructList(raw::MjContact* ptr, int num, py::handle owner)
|
||||
: StructListBase(ptr, num, 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
|
||||
MjContactList::MjStructList(raw::MjContact* ptr, int nconmax,
|
||||
int* ncon, py::handle owner)
|
||||
: StructListBase(ptr, nconmax, owner, /* lazy = */ true),
|
||||
ncon_(ncon) {}
|
||||
|
||||
// Slicing
|
||||
#define X(type, var) var(other.var[slice])
|
||||
MjContactList::MjStructList(MjContactList& other, py::slice slice)
|
||||
: StructListBase(other, slice),
|
||||
X(mjtNum, dist),
|
||||
X(mjtNum, pos),
|
||||
X(mjtNum, frame),
|
||||
X(mjtNum, includemargin),
|
||||
X(mjtNum, friction),
|
||||
X(mjtNum, solref),
|
||||
X(mjtNum, solimp),
|
||||
X(mjtNum, mu),
|
||||
X(mjtNum, H),
|
||||
X(int, dim),
|
||||
X(int, geom1),
|
||||
X(int, geom2),
|
||||
X(int, exclude),
|
||||
X(int, efc_address) {}
|
||||
#undef X
|
||||
ncon_(other.ncon_) {}
|
||||
|
||||
// ==================== MJDATA =================================================
|
||||
static void MjDataCapsuleDestructor(PyObject* pyobj) {
|
||||
@@ -579,9 +549,11 @@ MjDataWrapper::MjWrapper(const MjModelWrapper& model)
|
||||
var(InitPyArray(X_ARRAY_SHAPE(model.get()->dim0, dim1), ptr_->var, owner_)),
|
||||
MJDATA_POINTERS
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) x
|
||||
#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
|
||||
@@ -607,9 +579,11 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other)
|
||||
owner_)),
|
||||
MJDATA_POINTERS
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) x
|
||||
#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
|
||||
@@ -635,9 +609,11 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other)
|
||||
owner_)),
|
||||
MJDATA_POINTERS
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) x
|
||||
#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
|
||||
@@ -664,9 +640,11 @@ MjDataWrapper::MjWrapper(MjDataMetadata&& metadata, raw::MjData* d)
|
||||
var(InitPyArray(X_ARRAY_SHAPE(metadata.dim0, dim1), ptr_->var, owner_)),
|
||||
MJDATA_POINTERS
|
||||
#undef MJ_M
|
||||
#define MJ_M(x) x
|
||||
#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
|
||||
@@ -707,6 +685,8 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
|
||||
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));
|
||||
@@ -716,9 +696,17 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
|
||||
|
||||
// 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(nefc);
|
||||
X(ncon);
|
||||
X(time);
|
||||
X(energy);
|
||||
@@ -727,9 +715,32 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
|
||||
// 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
|
||||
}
|
||||
}
|
||||
@@ -754,6 +765,8 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
|
||||
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)); \
|
||||
@@ -773,9 +786,17 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
|
||||
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(nefc);
|
||||
X(ncon);
|
||||
X(time);
|
||||
X(energy);
|
||||
@@ -784,9 +805,33 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
|
||||
// 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");
|
||||
@@ -1697,22 +1742,36 @@ This is useful for example when the MJB is not available as a file on disk.)"));
|
||||
mjContactList.def("__len__", &MjContactList::size);
|
||||
DefineStructFunctions(mjContactList);
|
||||
|
||||
#define X(type, var) mjContactList.def_readonly(#var, &MjContactList::var)
|
||||
#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);
|
||||
X(mjtNum, pos);
|
||||
X(mjtNum, frame);
|
||||
XN(mjtNum, pos);
|
||||
XN(mjtNum, frame);
|
||||
X(mjtNum, includemargin);
|
||||
X(mjtNum, friction);
|
||||
X(mjtNum, solref);
|
||||
X(mjtNum, solimp);
|
||||
XN(mjtNum, friction);
|
||||
XN(mjtNum, solref);
|
||||
XN(mjtNum, solimp);
|
||||
X(mjtNum, mu);
|
||||
X(mjtNum, H);
|
||||
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");
|
||||
@@ -1749,6 +1808,29 @@ This is useful for example when the MJB is not available as a file on disk.)"));
|
||||
#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
|
||||
|
||||
|
||||
+53
-27
@@ -92,38 +92,59 @@ class MjWrapper {};
|
||||
template <typename T>
|
||||
class StructListBase {
|
||||
public:
|
||||
StructListBase(T* ptr, int num, pybind11::handle owner) : ptr_(ptr) {
|
||||
for (int i = 0; i < num; ++i) {
|
||||
wrappers_.push_back(std::make_shared<MjWrapper<T>>(&ptr[i], owner));
|
||||
StructListBase(T* ptr, int num, pybind11::handle owner, bool lazy = false)
|
||||
: ptr_(ptr), num_(num), owner_(owner) {
|
||||
if (!lazy) {
|
||||
PopulateUpTo(size());
|
||||
}
|
||||
}
|
||||
|
||||
StructListBase(const StructListBase& other) = delete;
|
||||
StructListBase(StructListBase&& other) = default;
|
||||
|
||||
virtual ~StructListBase() = default;
|
||||
|
||||
MjWrapper<T>& operator[](int i) {
|
||||
if (i < 0 || i >= wrappers_.size()) {
|
||||
if (i < 0 || i >= size()) {
|
||||
throw pybind11::index_error();
|
||||
}
|
||||
PopulateUpTo(i);
|
||||
return *wrappers_[i];
|
||||
}
|
||||
|
||||
int size() const {
|
||||
return wrappers_.size();
|
||||
virtual int size() const {
|
||||
return num_;
|
||||
}
|
||||
|
||||
T* get() const { return ptr_; }
|
||||
pybind11::handle owner() const { return owner_; }
|
||||
|
||||
protected:
|
||||
void PopulateUpTo(int n) {
|
||||
while (wrappers_.size() <= n) {
|
||||
wrappers_.push_back(
|
||||
std::make_shared<MjWrapper<T>>(&ptr_[wrappers_.size()], owner_));
|
||||
}
|
||||
}
|
||||
|
||||
// Slicing
|
||||
StructListBase(StructListBase& other, pybind11::slice slice) {
|
||||
StructListBase(StructListBase& other, pybind11::slice slice)
|
||||
: owner_(other.owner_) {
|
||||
pybind11::size_t start, stop, step, slicelength;
|
||||
slice.compute(other.size(), &start, &stop, &step, &slicelength);
|
||||
if (!slice.compute(other.size(), &start, &stop, &step, &slicelength)) {
|
||||
throw pybind11::index_error();
|
||||
}
|
||||
other.PopulateUpTo(stop);
|
||||
ptr_ = &other.ptr_[start];
|
||||
for (int i = start; i < stop; i += step) {
|
||||
wrappers_.push_back(other.wrappers_[i]);
|
||||
}
|
||||
num_ = wrappers_.size();
|
||||
}
|
||||
|
||||
T* ptr_;
|
||||
int num_;
|
||||
pybind11::handle owner_;
|
||||
|
||||
// Using shared_ptr here so that we get identical Python objects when slicing.
|
||||
std::vector<std::shared_ptr<MjWrapper<T>>> wrappers_;
|
||||
@@ -277,6 +298,8 @@ class MjStructList<raw::MjWarningStat>
|
||||
: public StructListBase<raw::MjWarningStat> {
|
||||
public:
|
||||
MjStructList(raw::MjWarningStat* ptr, int num, pybind11::handle owner);
|
||||
MjStructList(MjStructList&&) = default;
|
||||
~MjStructList() override = default;
|
||||
|
||||
using StructListBase::operator[];
|
||||
using StructListBase::size;
|
||||
@@ -325,6 +348,8 @@ template <>
|
||||
class MjStructList<raw::MjTimerStat> : public StructListBase<raw::MjTimerStat> {
|
||||
public:
|
||||
MjStructList(raw::MjTimerStat* ptr, int num, pybind11::handle owner);
|
||||
MjStructList(MjStructList&&) = default;
|
||||
~MjStructList() override = default;
|
||||
|
||||
using StructListBase::operator[];
|
||||
using StructListBase::size;
|
||||
@@ -374,6 +399,8 @@ class MjStructList<raw::MjSolverStat>
|
||||
: public StructListBase<raw::MjSolverStat> {
|
||||
public:
|
||||
MjStructList(raw::MjSolverStat* ptr, int num, pybind11::handle owner);
|
||||
MjStructList(MjStructList&&) = default;
|
||||
~MjStructList() override = default;
|
||||
|
||||
using StructListBase::operator[];
|
||||
using StructListBase::size;
|
||||
@@ -489,33 +516,28 @@ struct enable_if_mj_struct<raw::MjContact> { using type = void; };
|
||||
template <>
|
||||
class MjStructList<raw::MjContact> : public StructListBase<raw::MjContact> {
|
||||
public:
|
||||
MjStructList(raw::MjContact* ptr, int num, pybind11::handle owner);
|
||||
MjStructList(raw::MjContact* ptr, int nconmax,
|
||||
int* ncon, pybind11::handle owner);
|
||||
MjStructList(MjStructList&&) = default;
|
||||
~MjStructList() override = default;
|
||||
|
||||
using StructListBase::operator[];
|
||||
using StructListBase::size;
|
||||
|
||||
int size() const override {
|
||||
if (ncon_) {
|
||||
return *ncon_;
|
||||
} else {
|
||||
return StructListBase::size();
|
||||
}
|
||||
}
|
||||
|
||||
MjStructList Slice(pybind11::slice slice) {
|
||||
return MjStructList(*this, slice);
|
||||
}
|
||||
|
||||
#define X(type, var) pybind11::array_t<type> var
|
||||
X(mjtNum, dist);
|
||||
X(mjtNum, pos);
|
||||
X(mjtNum, frame);
|
||||
X(mjtNum, includemargin);
|
||||
X(mjtNum, friction);
|
||||
X(mjtNum, solref);
|
||||
X(mjtNum, solimp);
|
||||
X(mjtNum, mu);
|
||||
X(mjtNum, H);
|
||||
X(int, dim);
|
||||
X(int, geom1);
|
||||
X(int, geom2);
|
||||
X(int, exclude);
|
||||
X(int, efc_address);
|
||||
#undef X
|
||||
|
||||
protected:
|
||||
MjStructList(MjStructList& other, pybind11::slice slice);
|
||||
int* ncon_ = nullptr;
|
||||
};
|
||||
|
||||
using MjContactList = MjStructList<raw::MjContact>;
|
||||
@@ -539,6 +561,7 @@ class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
|
||||
MjWrapper(MjWrapper&&);
|
||||
~MjWrapper();
|
||||
|
||||
const MjDataMetadata& metadata() const { return metadata_; }
|
||||
MjDataIndexer& indexer() { return indexer_; }
|
||||
|
||||
void Serialize(std::ostream& output) const;
|
||||
@@ -548,10 +571,13 @@ class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
|
||||
"__MUJOCO_STRUCTS_MJDATAWRAPPER_LOOKUP";
|
||||
static MjWrapper* FromRawPointer(raw::MjData* m) noexcept;
|
||||
|
||||
|
||||
#define X(dtype, var, dim0, dim1) py_array_or_tuple_t<dtype> var;
|
||||
MJDATA_POINTERS
|
||||
#undef X
|
||||
|
||||
py_array_or_tuple_t<mjContact> contact;
|
||||
|
||||
py_array_or_tuple_t<raw::MjWarningStat> warning;
|
||||
py_array_or_tuple_t<raw::MjTimerStat> timer;
|
||||
py_array_or_tuple_t<raw::MjSolverStat> solver;
|
||||
|
||||
Reference in New Issue
Block a user