Saran Tunyasuvunakool
abf6d41b7f
The mismatch between the `py::array_t` template argument and the type of the MuJoCo struct member being wrapped caused the `py::array_t` to be constructed via an unintended constructor overload, which in turn triggered `use-after-poison` under asan. This bug affected the `geom`, `flex`, `elem`, and `vert` arrays in `mjContact`, and all array members in `mjrContext`. PiperOrigin-RevId: 608632860 Change-Id: Ic227f691041b004a0fcf0cb6fa3d62c44861d9aa
304 lines
8.7 KiB
C++
304 lines
8.7 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 <array>
|
|
#include <cstdint>
|
|
#include <optional>
|
|
#include <type_traits>
|
|
|
|
#include <Eigen/Core>
|
|
#include <mujoco/mjrender.h>
|
|
#include <mujoco/mujoco.h>
|
|
#include "errors.h"
|
|
#include "function_traits.h"
|
|
#include "functions.h"
|
|
#include "raw.h"
|
|
#include "structs.h"
|
|
#include <pybind11/eigen.h>
|
|
#include <pybind11/pybind11.h>
|
|
#include <pybind11/stl.h>
|
|
|
|
namespace mujoco::python {
|
|
namespace _impl {
|
|
template <>
|
|
class MjWrapper<raw::MjrContext> : public WrapperBase<raw::MjrContext> {
|
|
public:
|
|
MjWrapper();
|
|
MjWrapper(const MjModelWrapper& model, int fontscale);
|
|
MjWrapper(const MjWrapper&) = delete;
|
|
MjWrapper(MjWrapper&&) = default;
|
|
~MjWrapper() = default;
|
|
|
|
void Free();
|
|
|
|
#define X(var) \
|
|
py_array_or_tuple_t< \
|
|
std::remove_all_extents_t<decltype(raw::MjrContext::var)>> \
|
|
var
|
|
X(fogRGBA);
|
|
X(auxWidth);
|
|
X(auxHeight);
|
|
X(auxSamples);
|
|
X(auxFBO);
|
|
X(auxFBO_r);
|
|
X(auxColor);
|
|
X(auxColor_r);
|
|
X(textureType);
|
|
X(texture);
|
|
X(skinvertVBO);
|
|
X(skinnormalVBO);
|
|
X(skintexcoordVBO);
|
|
X(skinfaceVBO);
|
|
X(charWidth);
|
|
X(charWidthBig);
|
|
#undef X
|
|
};
|
|
|
|
using MjrContextWrapper = MjWrapper<raw::MjrContext>;
|
|
|
|
template <>
|
|
struct enable_if_mj_struct<raw::MjrContext> {
|
|
using type = void;
|
|
};
|
|
|
|
static void MjrContextCapsuleDestructor(PyObject* pyobj) {
|
|
auto* ptr =
|
|
static_cast<raw::MjrContext*>(PyCapsule_GetPointer(pyobj, nullptr));
|
|
mjr_freeContext(ptr);
|
|
delete ptr;
|
|
}
|
|
|
|
#define X(var) var(InitPyArray(ptr_->var, owner_))
|
|
#define X_SKIN(var) var(InitPyArray(std::array{ptr_->nskin}, ptr_->var, owner_))
|
|
MjrContextWrapper::MjWrapper()
|
|
: WrapperBase([]() {
|
|
raw::MjrContext *const ctx = new raw::MjrContext;
|
|
mjr_defaultContext(ctx);
|
|
return ctx;
|
|
}()),
|
|
X(fogRGBA),
|
|
X(auxWidth),
|
|
X(auxHeight),
|
|
X(auxSamples),
|
|
X(auxFBO),
|
|
X(auxFBO_r),
|
|
X(auxColor),
|
|
X(auxColor_r),
|
|
X(textureType),
|
|
X(texture),
|
|
X_SKIN(skinvertVBO),
|
|
X_SKIN(skinnormalVBO),
|
|
X_SKIN(skintexcoordVBO),
|
|
X_SKIN(skinfaceVBO),
|
|
X(charWidth),
|
|
X(charWidthBig) {}
|
|
|
|
MjrContextWrapper::MjWrapper(const MjModelWrapper& model, int fontscale)
|
|
: WrapperBase([fontscale](const raw::MjModel* m) {
|
|
raw::MjrContext *const ctx = new raw::MjrContext;
|
|
mjr_defaultContext(ctx);
|
|
InterceptMjErrors(mjr_makeContext)(m, ctx, fontscale);
|
|
return ctx;
|
|
}(model.get()), &MjrContextCapsuleDestructor),
|
|
X(fogRGBA),
|
|
X(auxWidth),
|
|
X(auxHeight),
|
|
X(auxSamples),
|
|
X(auxFBO),
|
|
X(auxFBO_r),
|
|
X(auxColor),
|
|
X(auxColor_r),
|
|
X(textureType),
|
|
X(texture),
|
|
X_SKIN(skinvertVBO),
|
|
X_SKIN(skinnormalVBO),
|
|
X_SKIN(skintexcoordVBO),
|
|
X_SKIN(skinfaceVBO),
|
|
X(charWidth),
|
|
X(charWidthBig) {}
|
|
#undef X_SKIN
|
|
#undef X
|
|
|
|
void MjrContextWrapper::Free() {
|
|
// mjr_freeContext is safe to call multiple times.
|
|
InterceptMjErrors(mjr_freeContext)(ptr_);
|
|
}
|
|
|
|
} // namespace _impl
|
|
|
|
namespace {
|
|
PYBIND11_MODULE(_render, pymodule) {
|
|
namespace py = ::pybind11;
|
|
namespace traits = python_traits;
|
|
|
|
using _impl::MjModelWrapper;
|
|
using _impl::MjrContextWrapper;
|
|
|
|
// Import the _structs module so that pybind11 knows about Python bindings
|
|
// for MjWrapper types and therefore generates prettier docstrings.
|
|
py::module::import("mujoco._structs");
|
|
|
|
py::class_<raw::MjrRect> mjrRect(pymodule, "MjrRect");
|
|
mjrRect.def(py::init([](int left, int bottom, int width, int height) {
|
|
return raw::MjrRect{left, bottom, width, height};
|
|
}),
|
|
py::arg("left"), py::arg("bottom"), py::arg("width"),
|
|
py::arg("height"));
|
|
mjrRect.def("__copy__",
|
|
[](const raw::MjrRect& other) { return raw::MjrRect(other); });
|
|
mjrRect.def("__deepcopy__", [](const raw::MjrRect& other, py::dict) {
|
|
return raw::MjrRect(other);
|
|
});
|
|
DefineStructFunctions(mjrRect);
|
|
#define X(var) mjrRect.def_readwrite(#var, &raw::MjrRect::var)
|
|
X(left);
|
|
X(bottom);
|
|
X(width);
|
|
X(height);
|
|
#undef X
|
|
|
|
py::class_<MjrContextWrapper> mjrContext(pymodule, "MjrContext");
|
|
mjrContext.def(py::init<>());
|
|
mjrContext.def(py::init<const MjModelWrapper&, int>());
|
|
mjrContext.def(
|
|
"free", [](MjrContextWrapper& self) { self.Free(); },
|
|
py::doc("Frees resources in current active OpenGL context, sets struct "
|
|
"to default."));
|
|
#define X(var) \
|
|
mjrContext.def_property( \
|
|
#var, [](const MjrContextWrapper& c) { return c.get()->var; }, \
|
|
[](MjrContextWrapper& c, mjtNum rhs) { c.get()->var = rhs; })
|
|
X(lineWidth);
|
|
X(shadowClip);
|
|
X(shadowScale);
|
|
X(fogStart);
|
|
X(fogEnd);
|
|
X(shadowSize);
|
|
X(offWidth);
|
|
X(offHeight);
|
|
X(offSamples);
|
|
X(fontScale);
|
|
X(offFBO);
|
|
X(offFBO_r);
|
|
X(offColor);
|
|
X(offColor_r);
|
|
X(offDepthStencil);
|
|
X(offDepthStencil_r);
|
|
X(shadowFBO);
|
|
X(shadowTex);
|
|
X(ntexture);
|
|
X(basePlane);
|
|
X(baseMesh);
|
|
X(baseHField);
|
|
X(baseBuiltin);
|
|
X(baseFontNormal);
|
|
X(baseFontShadow);
|
|
X(baseFontBig);
|
|
X(rangePlane);
|
|
X(rangeMesh);
|
|
X(rangeHField);
|
|
X(rangeBuiltin);
|
|
X(rangeFont);
|
|
X(nskin);
|
|
X(charHeight);
|
|
X(charHeightBig);
|
|
X(glInitialized);
|
|
X(windowAvailable);
|
|
X(windowSamples);
|
|
X(windowStereo);
|
|
X(windowDoublebuffer);
|
|
X(currentBuffer);
|
|
X(readPixelFormat);
|
|
X(readDepthMap);
|
|
#undef X
|
|
|
|
#define X(var) \
|
|
mjrContext.def_property_readonly( \
|
|
#var, [](const MjrContextWrapper& c) { return c.get()->var; })
|
|
X(nskin);
|
|
#undef X
|
|
|
|
#define X(var) DefinePyArray(mjrContext, #var, &MjrContextWrapper::var)
|
|
X(fogRGBA);
|
|
X(auxWidth);
|
|
X(auxHeight);
|
|
X(auxSamples);
|
|
X(auxFBO);
|
|
X(auxFBO_r);
|
|
X(auxColor);
|
|
X(auxColor_r);
|
|
X(textureType);
|
|
X(texture);
|
|
X(skinvertVBO);
|
|
X(skinnormalVBO);
|
|
X(skintexcoordVBO);
|
|
X(skinfaceVBO);
|
|
X(charWidth);
|
|
X(charWidthBig);
|
|
#undef X
|
|
|
|
using EigenUnsignedCharVectorX = Eigen::Vector<unsigned char, Eigen::Dynamic>;
|
|
using EigenFloatVectorX = Eigen::Vector<float, Eigen::Dynamic>;
|
|
|
|
// Skipped: mjr_defaultContext (have MjrContext.__init__)
|
|
// Skipped: mjr_makeContext (have MjrContext.__init__)
|
|
Def<traits::mjr_changeFont>(pymodule);
|
|
Def<traits::mjr_addAux>(pymodule);
|
|
// Skipped: mjr_freeContext (have MjrContext.__del__)
|
|
Def<traits::mjr_resizeOffscreen>(pymodule);
|
|
Def<traits::mjr_uploadTexture>(pymodule);
|
|
Def<traits::mjr_uploadMesh>(pymodule);
|
|
Def<traits::mjr_uploadHField>(pymodule);
|
|
Def<traits::mjr_restoreBuffer>(pymodule);
|
|
Def<traits::mjr_setBuffer>(pymodule);
|
|
DefWithGil<traits::mjr_readPixels>(
|
|
pymodule, [](std::optional<py::array_t<std::uint8_t>> rgb,
|
|
std::optional<py::array_t<float>> depth,
|
|
const raw::MjrRect* viewport, const raw::MjrContext* con) {
|
|
std::uint8_t* const rgb_data =
|
|
rgb.has_value() ? rgb->mutable_data() : nullptr;
|
|
float* const depth_data =
|
|
depth.has_value() ? depth->mutable_data() : nullptr;
|
|
{
|
|
py::gil_scoped_release no_gil;
|
|
return InterceptMjErrors(::mjr_readPixels)(rgb_data, depth_data,
|
|
*viewport, con);
|
|
}
|
|
});
|
|
Def<traits::mjr_drawPixels>(
|
|
pymodule,
|
|
[](std::optional<Eigen::Ref<const EigenUnsignedCharVectorX>> rgb,
|
|
std::optional<Eigen::Ref<const EigenFloatVectorX>> depth,
|
|
const raw::MjrRect* viewport, const raw::MjrContext* con) {
|
|
return InterceptMjErrors(::mjr_drawPixels)(
|
|
rgb.has_value() ? rgb->data() : nullptr,
|
|
depth.has_value() ? depth->data() : nullptr, *viewport, con);
|
|
});
|
|
Def<traits::mjr_blitBuffer>(pymodule);
|
|
Def<traits::mjr_setAux>(pymodule);
|
|
Def<traits::mjr_blitAux>(pymodule);
|
|
Def<traits::mjr_text>(pymodule);
|
|
Def<traits::mjr_overlay>(pymodule);
|
|
Def<traits::mjr_maxViewport>(pymodule);
|
|
Def<traits::mjr_rectangle>(pymodule);
|
|
Def<traits::mjr_label>(pymodule);
|
|
Def<traits::mjr_figure>(pymodule);
|
|
Def<traits::mjr_render>(pymodule);
|
|
Def<traits::mjr_finish>(pymodule);
|
|
Def<traits::mjr_getError>(pymodule);
|
|
Def<traits::mjr_findRect>(pymodule);
|
|
} // PYBIND11_MODULE
|
|
} // namespace
|
|
} // namespace mujoco::python
|