chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Allow user-defined visualization geoms to be added to the passive viewer.

Browse Source 
Fixes #1023, #1082

PiperOrigin-RevId: 572303635
Change-Id: Ia8399f816f311b38a7b137f5bd4f359f77a1da1c
This commit is contained in:
Saran Tunyasuvunakool
2023-10-10 10:51:29 -07:00
committed by Copybara-Service
parent 1bd6c94e6d
commit 2e15574b58
13 changed files with 112 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+5 -3
View File
@@ -131,6 +131,8 @@ Python bindings
23. Fixed `#870 <https://github.com/google-deepmind/mujoco/issues/870>`__ where calling ``update_scene`` with an invalid
camera name used the default camera.
24. Added ``user_scn`` to the :ref:`passive viewer<PyViewerPassive>` handle, which allows users to add custom
visualization geoms (`#1023 <https://github.com/google-deepmind/mujoco/issues/870>`__).
Simulate
^^^^^^^^
@@ -139,18 +141,18 @@ Simulate
:align: right
:width: 240px
24. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
25. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
:ref:`Python viewer<PyViewerManaged>`. State history can be viewed by scrubbing the History slider and (more
precisely) with the left and right arrow keys. See screen capture:
25. The ``LOADING...`` label is now shown correctly.
26. The ``LOADING...`` label is now shown correctly.
`Contribution <https://github.com/google-deepmind/mujoco/pull/1070>`__ by
`Levi Burner <https://github.com/aftersomemath>`__.
Bug fixes
^^^^^^^^^
26. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
27. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
midphase collision trees.
27. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints.
doc/includes/references.h
+1 -1
View File
@@ -1989,7 +1989,7 @@ struct mjvSceneState_ {
int nbuffer; // size of the buffer in bytes
void* buffer; // heap-allocated memory for all arrays in this struct
int maxgeom; // maximum number of mjvGeom supported by this state object
mjvScene plugincache; // scratch space for vis geoms inserted by plugins
mjvScene scratch; // scratch space for vis geoms inserted by the user and plugins
// fields in mjModel that are necessary to re-render a scene
struct {
doc/python.rst
+35 -5
View File
@@ -96,13 +96,13 @@ perturbations will not work unless the user explicitly synchronizes incoming eve
The ``launch_passive`` function returns a handle which can be used to interact with the viewer. It has the following
attributes:
- ``scn``, ``cam``, ``opt``, and ``pert`` properties: correspond to :ref:`mjvScene`, :ref:`mjvCamera`,
:ref:`mjvOption`, and :ref:`mjvPerturb` structs, respectively.
- ``cam``, ``opt``, and ``pert`` properties: correspond to :ref:`mjvCamera`, :ref:`mjvOption`, and :ref:`mjvPerturb`
structs, respectively.
- ``lock()``: provides a mutex lock for the viewer as a context manager. Since the viewer operates its own
thread, user code must ensure that it is holding the viewer lock before modifying any physics or visualization
state. These include the ``mjModel`` and ``mjData`` instance passed to ``launch_passive``, and also the ``scn``,
``cam``, ``opt``, and ``pert`` properties of the viewer handle.
state. These include the ``mjModel`` and ``mjData`` instance passed to ``launch_passive``, and also the ``cam``,
``opt``, and ``pert`` properties of the viewer handle.
- ``sync()``: synchronizes state between ``mjModel``, ``mjData``, and GUI user inputs since the previous call to
``sync``. In order to allow user scripts to make arbitrary modifications to ``mjModel`` and ``mjData`` without
@@ -124,6 +124,37 @@ attributes:
- ``is_running()``: returns ``True`` if the viewer window is running and ``False`` if it is closed.
This method can be safely called without locking.
- ``user_scn``: an :ref:`mjvScene` object that allows users to add custom visualization geoms to the rendered scene.
This is separate from the ``mjvScene`` that the viewer uses internally to render the final scene, and is entirely
under the user's control. User scripts can call e.g. :ref:`mjv_initGeom` or :ref:`mjv_makeConnector` to add
visualization geoms to ``user_scn``, and upon the next call to ``sync()``, the viewer will incorporate
these geoms to future rendered images. For example:
.. code-block:: python
with mujoco.viewer.launch_passive(m, d, key_callback=key_callback) as viewer:
while viewer.is_running():
...
# Step the physics.
mujoco.mj_step(m, d)
# Add a 3x3x3 grid of variously colored spheres to the middle of the scene.
viewer.user_scn.ngeom = 0
i = 0
for x, y, z in itertools.product(*((range(-1, 2),) * 3)):
mujoco.mjv_initGeom(
viewer.user_scn.geoms[i],
type=mujoco.mjtGeom.mjGEOM_SPHERE,
size=[0.02, 0, 0],
pos=0.1*np.array([x, y, z]),
mat=np.eye(3).flatten(),
rgba=0.5*np.array([x + 1, y + 1, z + 1, 2])
)
i += 1
viewer.user_scn.ngeom = i
viewer.sync()
...
The viewer handle can also be used as a context manager which calls ``close()`` automatically upon exit. A minimal
example of a user script that uses ``launch_passive`` might look like the following. (Note that example is a simple
illustrative example that does **not** necessarily keep the physics ticking at the correct wallclock rate.)
@@ -183,7 +214,6 @@ pause or resume the run loop when the spacebar is pressed.
viewer.sync()
...
.. _PyUsage:
Basic usage
include/mujoco/mjvisualize.h
+1 -1
View File
@@ -375,7 +375,7 @@ struct mjvSceneState_ {
int nbuffer; // size of the buffer in bytes
void* buffer; // heap-allocated memory for all arrays in this struct
int maxgeom; // maximum number of mjvGeom supported by this state object
mjvScene plugincache; // scratch space for vis geoms inserted by plugins
mjvScene scratch; // scratch space for vis geoms inserted by the user and plugins
// fields in mjModel that are necessary to re-render a scene
struct {
introspect/structs.py
+2 -2
View File
@@ -5476,9 +5476,9 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='maximum number of mjvGeom supported by this state object',
),
StructFieldDecl(
name='plugincache',
name='scratch',
type=ValueType(name='mjvScene'),
doc='scratch space for vis geoms inserted by plugins',
doc='scratch space for vis geoms inserted by the user and plugins', # pylint: disable=line-too-long
),
StructFieldDecl(
name='model',
python/mujoco/simulate.cc
+12 -7
View File
@@ -69,19 +69,23 @@ class UIAdapterWithPyCallback : public Adapter {
class SimulateWrapper {
public:
SimulateWrapper(std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
py::object scn, py::object cam, py::object opt,
py::object pert, bool is_passive)
py::object cam, py::object opt,
py::object pert, py::object user_scn, bool is_passive)
: simulate_(new mujoco::Simulate(
std::move(platform_ui_adapter), scn.cast<MjvSceneWrapper&>().get(),
std::move(platform_ui_adapter),
cam.cast<MjvCameraWrapper&>().get(),
opt.cast<MjvOptionWrapper&>().get(),
pert.cast<MjvPerturbWrapper&>().get(), is_passive)),
m_(py::none()),
d_(py::none()),
scn_(scn),
cam_(cam),
opt_(opt),
pert_(pert) {}
pert_(pert),
user_scn_(user_scn) {
if (!user_scn.is_none()) {
simulate_->user_scn = user_scn_.cast<MjvSceneWrapper&>().get();
}
}
~SimulateWrapper() { Destroy(); }
@@ -127,10 +131,10 @@ class SimulateWrapper {
// simulate object.
py::object m_;
py::object d_;
py::object scn_;
py::object cam_;
py::object opt_;
py::object pert_;
py::object user_scn_;
mjModel* m_raw_ = nullptr;
mjData* d_raw_ = nullptr;
@@ -207,7 +211,8 @@ PYBIND11_MODULE(_simulate, pymodule) {
.def("load_message", CallIfNotNull(&mujoco::Simulate::LoadMessage),
py::call_guard<py::gil_scoped_release>())
.def("load", &SimulateWrapper::Load)
.def("load_message_clear", CallIfNotNull(&mujoco::Simulate::LoadMessageClear),
.def("load_message_clear",
CallIfNotNull(&mujoco::Simulate::LoadMessageClear),
py::call_guard<py::gil_scoped_release>())
.def("sync", CallIfNotNull(&mujoco::Simulate::Sync),
py::call_guard<py::gil_scoped_release>())
python/mujoco/structs.cc
+3 -1
View File
@@ -1050,7 +1050,9 @@ MjvGeomWrapper::MjWrapper()
static_assert(sizeof(ptr_->mat) == sizeof(ptr_->mat[0])*9);
return InitPyArray(std::array{3, 3}, ptr_->mat, owner_);
}()),
X(rgba) {}
X(rgba) {
mjv_initGeom(ptr_, mjGEOM_NONE, nullptr, nullptr, nullptr, nullptr);
}
MjvGeomWrapper::MjWrapper(raw::MjvGeom* ptr, py::handle owner)
: WrapperBase(ptr, owner),
python/mujoco/viewer.py
+16 -13
View File
@@ -67,20 +67,16 @@ class Handle:
def __init__(
self,
sim: _Simulate,
scn: mujoco.MjvScene,
cam: mujoco.MjvCamera,
opt: mujoco.MjvOption,
pert: mujoco.MjvPerturb,
user_scn: Optional[mujoco.MjvScene],
):
self._sim = weakref.ref(sim)
self._scn = scn
self._cam = cam
self._opt = opt
self._pert = pert
@property
def scn(self):
return self._scn
self._user_scn = user_scn
@property
def cam(self):
@@ -94,6 +90,10 @@ class Handle:
def perturb(self):
return self._pert
@property
def user_scn(self):
return self._user_scn
def close(self):
sim = self._sim()
if sim is not None:
@@ -340,14 +340,16 @@ def _launch_internal(
loader = _loader
if model and not run_physics_thread:
scn = mujoco.MjvScene(model, _Simulate.MAX_GEOM)
else:
scn = mujoco.MjvScene()
cam = mujoco.MjvCamera()
opt = mujoco.MjvOption()
pert = mujoco.MjvPerturb()
simulate = _Simulate(scn, cam, opt, pert, run_physics_thread, key_callback)
if model and not run_physics_thread:
user_scn = mujoco.MjvScene(model, _Simulate.MAX_GEOM)
else:
user_scn = None
simulate = _Simulate(
cam, opt, pert, user_scn, run_physics_thread, key_callback
)
# Initialize GLFW if not using mjpython.
if _MJPYTHON is None:
@@ -357,8 +359,9 @@ def _launch_internal(
notify_loaded = None
if handle_return:
notify_loaded = (
lambda: handle_return.put_nowait(Handle(simulate, scn, cam, opt, pert)))
notify_loaded = lambda: handle_return.put_nowait(
Handle(simulate, cam, opt, pert, user_scn)
)
if run_physics_thread:
side_thread = threading.Thread(
simulate/main.cc
+1 -4
View File
@@ -474,9 +474,6 @@ int main(int argc, char** argv) {
// scan for libraries in the plugin directory to load additional plugins
scanPluginLibraries();
mjvScene scn;
mjv_defaultScene(&scn);
mjvCamera cam;
mjv_defaultCamera(&cam);
@@ -489,7 +486,7 @@ int main(int argc, char** argv) {
// simulate object encapsulates the UI
auto sim = std::make_unique<mj::Simulate>(
std::make_unique<mj::GlfwAdapter>(),
&scn, &cam, &opt, &pert, /* is_passive = */ false
&cam, &opt, &pert, /* is_passive = */ false
);
const char* filename = nullptr;
simulate/simulate.cc
+21 -6
View File
@@ -1737,16 +1737,15 @@ namespace mujoco {
namespace mju = ::mujoco::sample_util;
Simulate::Simulate(std::unique_ptr<PlatformUIAdapter> platform_ui,
mjvScene* scn, mjvCamera* cam,
mjvOption* opt, mjvPerturb* pert,
mjvCamera* cam, mjvOption* opt, mjvPerturb* pert,
bool is_passive)
: is_passive_(is_passive),
scn(*scn),
cam(*cam),
opt(*opt),
pert(*pert),
platform_ui(std::move(platform_ui)),
uistate(this->platform_ui->state()) {
mjv_defaultScene(&scn);
mjv_defaultSceneState(&scnstate_);
}
@@ -2011,6 +2010,23 @@ void Simulate::Sync() {
mjv_updateScene(m_, d_, &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn);
} else {
mjv_updateSceneState(m_, d_, &this->opt, &scnstate_);
// append geoms from user_scn to scnstate_ scratch space
if (user_scn) {
int ngeom = user_scn->ngeom;
int maxgeom = scnstate_.scratch.maxgeom - scnstate_.scratch.ngeom;
if (ngeom > maxgeom) {
mj_warning(d_, mjWARN_VGEOMFULL, scnstate_.scratch.maxgeom);
ngeom = maxgeom;
}
if (ngeom > 0) {
std::memcpy(scnstate_.scratch.geoms + scnstate_.scratch.ngeom,
user_scn->geoms,
sizeof(mjvGeom) * ngeom);
scnstate_.scratch.ngeom += ngeom;
}
}
mjopt_prev_ = scnstate_.model.opt;
warn_vgeomfull_prev_ = scnstate_.data.warning[mjWARN_VGEOMFULL].number;
}
@@ -2172,9 +2188,8 @@ void Simulate::LoadOnRenderThread() {
}
// re-create scene and context
if (!this->is_passive_) {
mjv_makeScene(this->m_, &this->scn, kMaxGeom);
} else {
mjv_makeScene(this->m_, &this->scn, kMaxGeom);
if (this->is_passive_) {
mjopt_prev_ = m_->opt;
opt_prev_ = opt;
cam_prev_ = cam;
simulate/simulate.h
+5 -3
View File
@@ -51,8 +51,7 @@ class Simulate {
// create object and initialize the simulate ui
Simulate(
std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
mjvScene* scn, mjvCamera* cam,
mjvOption* opt, mjvPerturb* pert, bool is_passive);
mjvCamera* cam, mjvOption* opt, mjvPerturb* pert, bool is_passive);
// Synchronize mjModel and mjData state with UI inputs, and update
// visualization.
@@ -233,7 +232,7 @@ class Simulate {
int camera = 0;
// abstract visualization
mjvScene& scn;
mjvScene scn;
mjvCamera& cam;
mjvOption& opt;
mjvPerturb& pert;
@@ -243,6 +242,9 @@ class Simulate {
mjvFigure figsize = {};
mjvFigure figsensor = {};
// additional user-defined visualization geoms (used in passive mode)
mjvScene* user_scn = nullptr;
// OpenGL rendering and UI
int refresh_rate = 60;
int window_pos[2] = {0};
src/engine/engine_vis_state.c
+9 -9
View File
@@ -48,19 +48,19 @@ static inline size_t roundUpToCacheLine(size_t n) {
// set default scene
void mjv_defaultSceneState(mjvSceneState* scnstate) {
memset(scnstate, 0, sizeof(mjvSceneState));
mjv_defaultScene(&scnstate->plugincache);
mjv_defaultScene(&scnstate->scratch);
}
// allocate and init scene state
void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnstate, int maxgeom) {
mjv_freeScene(&scnstate->plugincache);
mjv_freeScene(&scnstate->scratch);
mju_free(scnstate->buffer);
#ifdef MEMORY_SANITIZER
__msan_allocated_memory(scnstate, sizeof(mjvSceneState));
mjv_defaultScene(&scnstate->plugincache);
mjv_defaultScene(&scnstate->scratch);
#endif
scnstate->nbuffer = 0;
@@ -144,14 +144,14 @@ void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnsta
mjERROR("mjvSceneState buffer is not fully used");
}
mjv_makeScene(m, &scnstate->plugincache, maxgeom);
mjv_makeScene(m, &scnstate->scratch, maxgeom);
}
// free scene state
void mjv_freeSceneState(mjvSceneState* scnstate) {
mjv_freeScene(&scnstate->plugincache);
mjv_freeScene(&scnstate->scratch);
mju_free(scnstate->buffer);
mjv_defaultSceneState(scnstate);
}
@@ -229,14 +229,14 @@ int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt
int warning_start = d.warning[mjWARN_VGEOMFULL].number;
// copy mjvGeoms added by plugins
int nplugingeom = scnstate->plugincache.ngeom;
int nplugingeom = scnstate->scratch.ngeom;
if (nplugingeom > scn->maxgeom) {
mj_warning(&d, mjWARN_VGEOMFULL, scn->maxgeom);
scn->ngeom = scn->maxgeom;
} else {
scn->ngeom = nplugingeom;
}
memcpy(scn->geoms, scnstate->plugincache.geoms, sizeof(mjvGeom) * scn->ngeom);
memcpy(scn->geoms, scnstate->scratch.geoms, sizeof(mjvGeom) * scn->ngeom);
// add all categories
mjv_addGeoms(&m, &d, opt, pert, catmask, scn);
@@ -272,7 +272,7 @@ void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
#undef X
// Update plugin visualization cache.
scnstate->plugincache.ngeom = 0;
scnstate->scratch.ngeom = 0;
if (m->nplugin) {
const int nslot = mjp_pluginCount();
// iterate over plugins, call visualize if defined
@@ -283,7 +283,7 @@ void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
mjERROR("invalid plugin slot: %d", slot);
}
if (plugin->visualize) {
plugin->visualize(m, d, opt, &scnstate->plugincache, i);
plugin->visualize(m, d, opt, &scnstate->scratch, i);
}
}
}
unity/Runtime/Bindings/MjBindings.cs
+1 -1
View File
@@ -6148,7 +6148,7 @@ public unsafe struct mjvSceneState_ {
public int nbuffer;
public void* buffer;
public int maxgeom;
public mjvScene_ plugincache;
public mjvScene_ scratch;
public model model;
public data data;
}public struct mjuiItem_ {}public struct mjfItemEnable {}