mujoco

// Copyright 2021 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.

#ifndef MUJOCO_SIMULATE_SIMULATE_H_
#define MUJOCO_SIMULATE_SIMULATE_H_

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <memory>
#include <mutex>
#include <optional>
#include <ratio>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include <mujoco/mjui.h>
#include <mujoco/mujoco.h>
#include "platform_ui_adapter.h"

namespace mujoco {

// 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 {
 public:
  using Clock = std::chrono::steady_clock;
  static_assert(std::ratio_less_equal_v<Clock::period, std::milli>);

  static constexpr int kMaxGeom = 100000;

  // create object and initialize the simulate ui
  Simulate(std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
           mjvCamera*                         cam,
           mjvOption*                         opt,
           mjvPerturb*                        pert,
           bool                               is_passive);

  // Synchronize state with UI inputs, and update visualization.  If state_only
  // is false mjData and mjModel will be updated, otherwise only the subset of
  // mjData corresponding to mjSTATE_INTEGRATION will be synced.
  void Sync(bool state_only = false);

  void UpdateHField(int hfieldid);
  void UpdateMesh(int meshid);
  void UpdateTexture(int texid);

  // Request that the Simulate UI display a "loading" message
  // Called prior to Load or LoadMessageClear
  void LoadMessage(const char* displayed_filename);

  // Request that the Simulate UI thread render a new model
  void Load(mjModel* m, mjData* d, const char* displayed_filename);

  // Clear the loading message
  // Can be called instead of Load to clear the message without
  // requesting the UI load a model
  void LoadMessageClear(void);

  // functions below are used by the renderthread
  // load mjb or xml model that has been requested by load()
  void LoadOnRenderThread();

  // render the ui to the window
  void Render();

  // loop to render the UI (must be called from main thread because of MacOS)
  void RenderLoop();

  // add state to history buffer
  void AddToHistory();

  // inject control noise
  void InjectNoise(int key);

  // constants
  static constexpr int kMaxFilenameLength = 1000;

  // whether the viewer is operating in passive mode, where it cannot assume
  // that it has exclusive access to mjModel, mjData, and various mjv objects
  bool is_passive_ = false;

  // model and data to be visualized
  mjModel* mnew_ = nullptr;
  mjData*  dnew_ = nullptr;

  mjModel* m_ = nullptr;
  mjData*  d_ = nullptr;

  int ncam_           = 0;
  int nkey_           = 0;
  int state_size_     = 0;  // number of mjtNums in a history buffer state
  int nhistory_       = 0;  // number of states saved in history buffer
  int history_cursor_ = 0;  // cursor pointing at last saved state

  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_;


  std::vector<std::string> equality_names_;

  std::vector<mjtNum> history_;  // history buffer (nhistory x state_size)

  // 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_;
  std::vector<mjtByte> eq_active_;
  std::vector<mjtByte> eq_active_prev_;

  // in passive mode the user owns m_ and d_, these "passive" instances are
  // owned by Simulate, updated from the user by the Sync() method
  mjModel* m_passive_ = nullptr;
  mjData*  d_passive_ = nullptr;

  std::vector<mjvGeom> user_scn_geoms_;

  mjOption    mjopt_prev_;
  mjVisual    mjvis_prev_;
  mjStatistic mjstat_prev_;
  mjvOption   opt_prev_;
  mjvCamera   cam_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      update_threadpool;
    bool      reset;
    bool      align;
    bool      copy_key;
    bool      copy_key_full_precision;
    bool      load_from_history;
    bool      load_key;
    bool      save_key;
    bool      zero_ctrl;
    int       newperturb;
    bool      select;
    mjuiState select_state;
    bool      ui_update_simulation;
    bool      ui_update_physics;
    bool      ui_update_rendering;
    bool      ui_update_visualization;
    bool      ui_update_joint;
    bool      ui_update_ctrl;
    bool      ui_update_equality;
    bool      ui_update_logging;
    bool      ui_remake_ctrl;
  } pending_ = {};

  SimulateMutex               mtx;
  std::condition_variable_any cond_loadrequest;

  int    frames_ = 0;
  double fps_    = 0;

  std::chrono::time_point<Clock> last_fps_update_;

  // options
  int spacing      = 0;
  int color        = 0;
  int font         = 0;
  int ui0_enable   = 1;
  int ui1_enable   = 1;
  int help         = 0;
  int info         = 0;
  int profiler     = 0;
  int sensor       = 0;
  int pause_update = 0;
  int fullscreen   = 0;
  int vsync        = 1;
  int busywait     = 0;

  // keyframe index
  int key = -1;

  // index of history-scrubber slider
  int scrub_index = 0;

  // simulation
  int run = 1;


  // number of workers in threadpool
  int nthread = 0;

  // atomics for cross-thread messages
  std::atomic_int exitrequest       = 0;
  std::atomic_int droploadrequest   = 0;
  std::atomic_int screenshotrequest = 0;
  std::atomic_int uiloadrequest     = 0;
  std::atomic_int newfigurerequest  = 0;
  std::atomic_int newtextrequest    = 0;
  std::atomic_int newimagerequest   = 0;

  // loadrequest
  //   3: display a loading message
  //   2: render thread asked to update its model
  //   1: showing "loading" label, about to load
  //   0: model loaded or no load requested.
  int loadrequest = 0;

  // strings
  char load_error[kMaxFilenameLength]        = "";
  char dropfilename[kMaxFilenameLength]      = "";
  char filename[kMaxFilenameLength]          = "";
  char previous_filename[kMaxFilenameLength] = "";

  // time synchronization
  int   real_time_index   = 0;
  bool  speed_changed     = true;
  float measured_slowdown = 1.0;

  // logarithmically spaced real-time slow-down coefficients (percent)
  // clang-format off
  static constexpr float percentRealTime[] = {
      100, 80, 63,  50,  40, 32,  25,  20, 16,  13,
      10,  8,  6.3, 5.0, 4,  3.2, 2.5, 2,  1.6, 1.3,
      1,  .8, .63, .5,  .4, .32, .25, .2, .16, .13,
     .1
  };
  // clang-format on

  // control noise
  double ctrl_noise_std  = 0.0;
  double ctrl_noise_rate = 0.0;

  // watch
  char field[mjMAXUITEXT] = "qpos";
  int  index              = 0;

  // physics: need sync
  int disable[mjNDISABLE]      = {0};
  int enable[mjNENABLE]        = {0};
  int enableactuator[mjNGROUP] = {0};

  // logging: need sync
  mjtByte     log_console           = 0;
  mjtByte     log_file              = 0;
  mjtByte     log_topics[mjNTOPIC]  = {0};
  mjTimerStat timer_prev_[mjNTIMER] = {};

  // rendering: need sync
  int camera = 0;

  // abstract visualization
  mjvScene    scn;
  mjvCamera&  cam;
  mjvOption&  opt;
  mjvPerturb& pert;
  mjvFigure   figconstraint = {};
  mjvFigure   figcost       = {};
  mjvFigure   figtimer      = {};
  mjvFigure   figsize       = {};
  mjvFigure   figsensor     = {};

  // Image sensor visualization - displays pre-rendered images from mjSENS_USER
  int image_sensor_count        = 0;
  int selected_image_sensor     = -1;  // -1 = show bar chart
  int image_sensor_ui_selection = 0;   // UI dropdown index (0=All, 1+=sensor)

  std::vector<int>         image_sensor_indices;
  std::vector<std::string> image_sensor_names;

  std::unique_ptr<unsigned char[]> sensor_image;

  int sensor_image_width    = 0;
  int sensor_image_height   = 0;
  int sensor_image_last_seq = -1;  // Last seq read from sensordata

  // additional user-defined visualization
  mjvScene* user_scn = nullptr;
  mjtByte   user_scn_flags_prev_[mjNRNDFLAG];

  std::vector<std::pair<mjrRect, mjvFigure>> user_figures_;
  std::vector<std::pair<mjrRect, mjvFigure>> user_figures_new_;

  std::vector<std::tuple<int, int, std::string, std::string>> user_texts_;
  std::vector<std::tuple<int, int, std::string, std::string>> user_texts_new_;

  std::vector<std::tuple<mjrRect, std::unique_ptr<unsigned char[]>>> user_images_;
  std::vector<std::tuple<mjrRect, std::unique_ptr<unsigned char[]>>> user_images_new_;

  // OpenGL rendering and UI
  int refresh_rate   = 60;
  int window_pos[2]  = {0};
  int window_size[2] = {0};

  std::unique_ptr<PlatformUIAdapter> platform_ui;

  mjuiState& uistate;
  mjUI       ui0 = {};
  mjUI       ui1 = {};


  // Constant arrays needed for the option section of UI and the UI interface
  // clang-format off
  // TODO setting the size here is not ideal
  const mjuiDef def_option[13] = {
    {mjITEM_SECTION,  "Option",        mjPRESERVE, nullptr,  "AO"},
    {mjITEM_CHECKINT, "Help",          2, &this->help,       " #290"},
    {mjITEM_CHECKINT, "Info",          2, &this->info,       " #291"},
    {mjITEM_CHECKINT, "Profiler",      2, &this->profiler,   " #292"},
    {mjITEM_CHECKINT, "Sensor",        2, &this->sensor,     " #293"},
    {mjITEM_CHECKINT, "Pause update",  2, &this->pause_update,    ""},
#ifdef __APPLE__
    {mjITEM_CHECKINT, "Fullscreen",    0, &this->fullscreen, " #294"},
#else
    {mjITEM_CHECKINT, "Fullscreen",    1, &this->fullscreen, " #294"},
#endif
    {mjITEM_CHECKINT, "Vertical Sync", 1, &this->vsync,      ""},
    {mjITEM_CHECKINT, "Busy Wait",     1, &this->busywait,   ""},
    {mjITEM_SELECT,   "Spacing",       1, &this->spacing,    "Tight\nWide"},
    {mjITEM_SELECT,   "Color",         1, &this->color,      "Default\nOrange\nWhite\nBlack"},
    {mjITEM_SELECT,   "Font",          1, &this->font,       "50 %\n100 %\n150 %\n200 %\n250 %\n300 %"},
    {mjITEM_END}
  };


  // simulation section of UI
  const mjuiDef def_simulation[15] = {
    {mjITEM_SECTION,   "Simulation",    mjPRESERVE, nullptr,     "AS"},
    {mjITEM_RADIO,     "",              5, &this->run,           "Pause\nRun"},
    {mjITEM_SLIDERINT, "Num threads",   5, &this->nthread,       "0 10"},
    {mjITEM_BUTTON,    "Reset",         2, nullptr,              " #259"},
    {mjITEM_BUTTON,    "Reload",        5, nullptr,              "CL"},
    {mjITEM_BUTTON,    "Align",         2, nullptr,              "CA"},
    {mjITEM_BUTTON,    "Copy state",    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",   5, &this->ctrl_noise_std,  "0 1"},
    {mjITEM_SLIDERNUM, "Noise rate",    5, &this->ctrl_noise_rate, "0 4"},
    {mjITEM_SEPARATOR, "History",       1},
    {mjITEM_SLIDERINT, "",              5, &this->scrub_index,     "0 0"},
    {mjITEM_END}
  };


  // watch section of UI
  const mjuiDef def_watch[5] = {
    {mjITEM_SECTION,   "Watch",         mjPRESERVE, nullptr,     "AW"},
    {mjITEM_EDITTXT,   "Field",         2, this->field,          "qpos"},
    {mjITEM_EDITINT,   "Index",         2, &this->index,         "1"},
    {mjITEM_STATIC,    "Value",         2, nullptr,              " "},
    {mjITEM_END}
  };
  // clang-format on

  // info strings
  char info_title[Simulate::kMaxFilenameLength]   = {0};
  char info_content[Simulate::kMaxFilenameLength] = {0};

  // pending uploads
  std::condition_variable_any cond_upload_;

  int texture_upload_ = -1;
  int mesh_upload_    = -1;
  int hfield_upload_  = -1;
};
}  // namespace mujoco

#endif