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:
Saran Tunyasuvunakool
2022-10-06 10:02:35 -07:00
committed by Copybara-Service
parent 4d85a464cc
commit 58fd72f53d
29 changed files with 1281 additions and 433 deletions
+22 -6
View File
@@ -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):
+4 -4
View File
@@ -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
+9 -8
View File
@@ -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
View File
@@ -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
View File
@@ -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;