Expose a handle for the Python viewer.

This change also requires user scripts to explicitly synchronize changes to physics state to the viewer. The Simulate class was reconfigured so that certain UI events are handled during this sync operation, outside of the render loop on the main thread. These correspond to operations that require access to the full mjModel/mjData.

To support other, more interactive operations (e.g. camera movements), a new mjvSceneState struct is introduced which captures only the portion of the physics state required for scene re-rendering. The mjvSceneState is updated from mjModel/mjData during the viewer sync operation, and is significantly cheaper than a full mj_copyModel and mj_copyData.

Fixes https://github.com/deepmind/mujoco/issues/796

PiperOrigin-RevId: 525723636
Change-Id: Id08d0210a2c067d5afe85e2bf104f276aeddd75e

simulate/main.cc

@@ -313,7 +313,7 @@ void PhysicsLoop(mj::Simulate& sim) {
 
     {
       // lock the sim mutex
-      const std::lock_guard<std::mutex> lock(sim.mtx);
+      const std::unique_lock<std::recursive_mutex> lock(sim.mtx);
 
       // run only if model is present
       if (m) {
@@ -356,11 +356,6 @@ void PhysicsLoop(mj::Simulate& sim) {
             syncSim = d->time;
             sim.speed_changed = false;
 
-            // clear old perturbations, apply new
-            mju_zero(d->xfrc_applied, 6*m->nbody);
-            sim.ApplyPosePerturbations(0);  // move mocap bodies only
-            sim.ApplyForcePerturbations();
-
             // run single step, let next iteration deal with timing
             mj_step(m, d);
           }
@@ -382,11 +377,6 @@ void PhysicsLoop(mj::Simulate& sim) {
                 measured = true;
               }
 
-              // clear old perturbations, apply new
-              mju_zero(d->xfrc_applied, 6*m->nbody);
-              sim.ApplyPosePerturbations(0);  // move mocap bodies only
-              sim.ApplyForcePerturbations();
-
               // call mj_step
               mj_step(m, d);
 
@@ -400,9 +390,6 @@ void PhysicsLoop(mj::Simulate& sim) {
 
         // paused
         else {
-          // apply pose perturbation
-          sim.ApplyPosePerturbations(1);  // move mocap and dynamic bodies
-
           // run mj_forward, to update rendering and joint sliders
           mj_forward(m, d);
         }
@@ -468,9 +455,23 @@ int main(int argc, const char** argv) {
   // scan for libraries in the plugin directory to load additional plugins
   scanPluginLibraries();
 
+  mjvScene scn;
+  mjv_defaultScene(&scn);
+
+  mjvCamera cam;
+  mjv_defaultCamera(&cam);
+
+  mjvOption opt;
+  mjv_defaultOption(&opt);
+
+  mjvPerturb pert;
+  mjv_defaultPerturb(&pert);
+
   // simulate object encapsulates the UI
   auto sim = std::make_unique<mj::Simulate>(
-      std::make_unique<mj::GlfwAdapter>());
+      std::make_unique<mj::GlfwAdapter>(),
+      &scn, &cam, &opt, &pert, /* fully_managed = */ true
+  );
 
   const char* filename = nullptr;
   if (argc >  1) {

simulate/platform_ui_adapter.cc

@@ -35,6 +35,10 @@ bool PlatformUIAdapter::RefreshMjrContext(const mjModel* m, int fontscale) {
   return false;
 }
 
+bool PlatformUIAdapter::EnsureContextSize() {
+  return false;
+}
+
 void PlatformUIAdapter::OnFilesDrop(int count, const char** paths) {
   state_.type = mjEVENT_FILESDROP;
   state_.dropcount = count;

simulate/platform_ui_adapter.h

@@ -41,6 +41,8 @@ class PlatformUIAdapter {
   // Optionally overrideable function to (re)create an mjrContext for an mjModel
   virtual bool RefreshMjrContext(const mjModel* m, int fontscale);
 
+  virtual bool EnsureContextSize();
+
   // Pure virtual functions to be implemented by individual adapters
   virtual std::pair<double, double> GetCursorPosition() const = 0;
   virtual double GetDisplayPixelsPerInch() const = 0;

simulate/simulate.h

@@ -20,15 +20,24 @@
 #include <condition_variable>
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <ratio>
-#include <thread>
+#include <utility>
+#include <vector>
 
+#include <mujoco/mjui.h>
 #include <mujoco/mujoco.h>
 #include "platform_ui_adapter.h"
 
 namespace mujoco {
 
-//-------------------------------- global -----------------------------------------------
+// The viewer itself doesn't require a reentrant mutex, however we use it in
+// order to provide a Python sync API that doesn't require separate locking
+// (since sync is by far the most common operation), but that also won't
+// deadlock if called when a lock is already held by the user script on the
+// same thread.
+class SimulateMutex : public std::recursive_mutex {};
+using MutexLock = std::unique_lock<std::recursive_mutex>;
 
 // Simulate states not contained in MuJoCo structures
 class Simulate {
@@ -36,14 +45,17 @@ class Simulate {
   using Clock = std::chrono::steady_clock;
   static_assert(std::ratio_less_equal_v<Clock::period, std::milli>);
 
+  static constexpr int kMaxGeom = 20000;
+
   // create object and initialize the simulate ui
-  Simulate(std::unique_ptr<PlatformUIAdapter> platform_ui_adapter);
+  Simulate(
+      std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
+      mjvScene* scn, mjvCamera* cam,
+      mjvOption* opt, mjvPerturb* pert, bool fully_managed);
 
-  // Apply UI pose perturbations to model and data
-  void ApplyPosePerturbations(int flg_paused);
-
-  // Apply UI force perturbations to model and data
-  void ApplyForcePerturbations();
+  // Synchronize mjModel and mjData state with UI inputs, and update
+  // visualization.
+  void Sync();
 
   // Request that the Simulate UI thread render a new model
   // optionally delete the old model and data when done
@@ -53,9 +65,6 @@ class Simulate {
   // load mjb or xml model that has been requested by load()
   void LoadOnRenderThread();
 
-  // prepare to render
-  void PrepareScene();
-
   // render the ui to the window
   void Render();
 
@@ -65,14 +74,71 @@ class Simulate {
   // constants
   static constexpr int kMaxFilenameLength = 1000;
 
-  // model and data to be visualized
-  mjModel* mnew = nullptr;
-  mjData* dnew = nullptr;
+  // whether the viewer is operating in fully managed mode, where it can assume
+  // that it has exclusive access to mjModel, mjData, and various mjv objects
+  bool fully_managed_ = true;
 
-  mjModel* m = nullptr;
-  mjData* d = nullptr;
-  std::mutex mtx;
-  std::condition_variable cond_loadrequest;
+  // model and data to be visualized
+  mjModel* mnew_ = nullptr;
+  mjData* dnew_ = nullptr;
+
+  mjModel* m_ = nullptr;
+  mjData* d_ = nullptr;
+
+  int ncam_ = 0;
+  int nkey_ = 0;
+
+  std::vector<int> body_parentid_;
+
+  std::vector<int> jnt_type_;
+  std::vector<int> jnt_group_;
+  std::vector<int> jnt_qposadr_;
+  std::vector<std::optional<std::pair<mjtNum, mjtNum>>> jnt_range_;
+  std::vector<std::string> jnt_names_;
+
+  std::vector<int> actuator_group_;
+  std::vector<std::optional<std::pair<mjtNum, mjtNum>>> actuator_ctrlrange_;
+  std::vector<std::string> actuator_names_;
+
+  // mjModel and mjData fields that can be modified by the user through the GUI
+  std::vector<mjtNum> qpos_;
+  std::vector<mjtNum> qpos_prev_;
+  std::vector<mjtNum> ctrl_;
+  std::vector<mjtNum> ctrl_prev_;
+
+  mjvSceneState scnstate_;
+  mjOption mjopt_prev_;
+  mjvOption opt_prev_;
+  mjvCamera cam_prev_;
+  int warn_vgeomfull_prev_;
+
+  // pending GUI-driven actions, to be applied at the next call to Sync
+  struct {
+    std::optional<std::string> save_xml;
+    std::optional<std::string> save_mjb;
+    std::optional<std::string> print_model;
+    std::optional<std::string> print_data;
+    bool reset;
+    bool align;
+    bool copy_pose;
+    bool load_key;
+    bool save_key;
+    bool zero_ctrl;
+    int newperturb;
+    bool select;
+    mjuiState select_state;
+    bool full_ui_update;
+    bool ui_update_physics;
+    bool ui_update_joint;
+    bool ui_update_ctrl;
+  } pending_ = {};
+
+  SimulateMutex mtx;
+  std::condition_variable_any cond_loadrequest;
+
+  int frames_ = 0;
+  std::chrono::time_point<Clock> last_fps_update_;
+  double fps_ = 0;
 
   // options
   int spacing = 0;
@@ -140,10 +206,10 @@ class Simulate {
   int camera = 0;
 
   // abstract visualization
-  mjvScene scn = {};
-  mjvCamera cam = {};
-  mjvOption opt = {};
-  mjvPerturb pert = {};
+  mjvScene& scn;
+  mjvCamera& cam;
+  mjvOption& opt;
+  mjvPerturb& pert;
   mjvFigure figconstraint = {};
   mjvFigure figcost = {};
   mjvFigure figtimer = {};
@@ -186,16 +252,16 @@ class Simulate {
   // simulation section of UI
   const mjuiDef def_simulation[12] = {
     {mjITEM_SECTION,   "Simulation",    1, nullptr,              "AS"},
-    {mjITEM_RADIO,     "",              2, &this->run,           "Pause\nRun"},
+    {mjITEM_RADIO,     "",              5, &this->run,           "Pause\nRun"},
     {mjITEM_BUTTON,    "Reset",         2, nullptr,              " #259"},
-    {mjITEM_BUTTON,    "Reload",        2, nullptr,              "CL"},
+    {mjITEM_BUTTON,    "Reload",        5, nullptr,              "CL"},
     {mjITEM_BUTTON,    "Align",         2, nullptr,              "CA"},
     {mjITEM_BUTTON,    "Copy pose",     2, nullptr,              "CC"},
     {mjITEM_SLIDERINT, "Key",           3, &this->key,           "0 0"},
     {mjITEM_BUTTON,    "Load key",      3},
     {mjITEM_BUTTON,    "Save key",      3},
-    {mjITEM_SLIDERNUM, "Noise scale",   2, &this->ctrl_noise_std,  "0 2"},
-    {mjITEM_SLIDERNUM, "Noise rate",    2, &this->ctrl_noise_rate, "0 2"},
+    {mjITEM_SLIDERNUM, "Noise scale",   5, &this->ctrl_noise_std,  "0 2"},
+    {mjITEM_SLIDERNUM, "Noise rate",    5, &this->ctrl_noise_rate, "0 2"},
     {mjITEM_END}
   };