From 3e12f0d50c15c5d6fb3da45a1db1c29bfcfd77b0 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 4 Oct 2023 06:35:14 -0700 Subject: [PATCH] Add history scrubbing mechanism to `simulate`. [screen capture](https://www.youtube.com/watch?v=YSvWn_poqWs) PiperOrigin-RevId: 570676968 Change-Id: Id9a1fccc9418b851e007ddcf3fa89acd4edff3e9 --- doc/changelog.rst | 18 ++++- python/mujoco/simulate.cc | 3 +- python/mujoco/viewer.py | 8 +++ simulate/main.cc | 15 ++-- simulate/simulate.cc | 143 +++++++++++++++++++++++++++++++++----- simulate/simulate.h | 16 ++++- 6 files changed, 177 insertions(+), 26 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index c549ee74..05c88c99 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -129,13 +129,27 @@ Python bindings 22. Fixed `#870 `__ where calling ``update_scene`` with an invalid camera name used the default camera. +Simulate +^^^^^^^^ + +.. youtube:: YSvWn_poqWs + :align: right + :width: 240px + +23. Added **state history** mechanism to :ref:`simulate` and the managed + :ref:`Python viewer`. State history can be viewed by scrubbing the History slider and (more + precisely) with the left and right arrow keys. See screen capture: + +24. The ``LOADING...`` label is now shown correctly. + `Contribution `__ by + `Levi Burner `__. + Bug fixes ^^^^^^^^^ -23. Fixed a bug that was causing :ref:`geom margin` to be ignored during the construction of +25. Fixed a bug that was causing :ref:`geom margin` to be ignored during the construction of midphase collision trees. - Version 2.3.7 (July 20, 2023) ----------------------------- diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 707a8f86..42a70211 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -210,7 +210,8 @@ PYBIND11_MODULE(_simulate, pymodule) { py::call_guard()) .def("sync", CallIfNotNull(&mujoco::Simulate::Sync), py::call_guard()) - + .def("add_to_history", CallIfNotNull(&mujoco::Simulate::AddToHistory), + py::call_guard()) .def("render_loop", CallIfNotNull(&mujoco::Simulate::RenderLoop), py::call_guard()) .def("lock", GetIfNotNull(&mujoco::Simulate::mtx), diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 61520240..2a0be3b9 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -237,6 +237,7 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]): if m is not None: assert d is not None if simulate.run: + stepped = False # Record CPU time at start of iteration. startcpu = time.time() @@ -275,6 +276,7 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]): # Run single step, let next iteration deal with timing. mujoco.mj_step(m, d) + stepped = True # In-sync: step until ahead of cpu. else: @@ -292,10 +294,16 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]): # Call mj_step. mujoco.mj_step(m, d) + stepped = True # Break if reset. if d.time < prevsim: break + + # save current state to history buffer + if (stepped): + simulate.add_to_history() + else: # simulate.run is False: GUI is paused. # Run mj_forward, to update rendering and joint sliders. diff --git a/simulate/main.cc b/simulate/main.cc index 71aad17d..a9fde2dd 100644 --- a/simulate/main.cc +++ b/simulate/main.cc @@ -279,8 +279,7 @@ void PhysicsLoop(mj::Simulate& sim) { free(ctrlnoise); ctrlnoise = (mjtNum*) malloc(sizeof(mjtNum)*m->nu); mju_zero(ctrlnoise, m->nu); - } - else { + } else { sim.LoadMessageClear(); } } @@ -305,8 +304,7 @@ void PhysicsLoop(mj::Simulate& sim) { free(ctrlnoise); ctrlnoise = static_cast(malloc(sizeof(mjtNum)*m->nu)); mju_zero(ctrlnoise, m->nu); - } - else { + } else { sim.LoadMessageClear(); } } @@ -327,6 +325,8 @@ void PhysicsLoop(mj::Simulate& sim) { if (m) { // running if (sim.run) { + bool stepped = false; + // record cpu time at start of iteration const auto startCPU = mj::Simulate::Clock::now(); @@ -366,6 +366,7 @@ void PhysicsLoop(mj::Simulate& sim) { // run single step, let next iteration deal with timing mj_step(m, d); + stepped = true; } // in-sync: step until ahead of cpu @@ -387,6 +388,7 @@ void PhysicsLoop(mj::Simulate& sim) { // call mj_step mj_step(m, d); + stepped = true; // break if reset if (d->time < prevSim) { @@ -394,6 +396,11 @@ void PhysicsLoop(mj::Simulate& sim) { } } } + + // save current state to history buffer + if (stepped) { + sim.AddToHistory(); + } } // paused diff --git a/simulate/simulate.cc b/simulate/simulate.cc index 3d2c96d4..6b7f4acb 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -131,6 +131,7 @@ const char help_content[] = "Space\n" "+ -\n" "Right arrow\n" + "Left arrow\n" "[ ]\n" "Esc\n" "Double-click\n" @@ -153,7 +154,8 @@ const char help_content[] = const char help_title[] = "Play / Pause\n" "Speed up / down\n" - "Step\n" + "Step forward\n" + "Step back\n" "Cycle cameras\n" "Free camera\n" "Select\n" @@ -539,6 +541,25 @@ void ShowSensor(mj::Simulate* sim, mjrRect rect) { mjr_figure(viewport, &sim->figsensor, &sim->platform_ui->mjr_context()); } +// load state from history buffer +static void LoadScrubState(mj::Simulate* sim) { + // get index into circular buffer + int i = (sim->scrub_index + sim->history_cursor_) % sim->nhistory_; + i = (i + sim->nhistory_) % sim->nhistory_; + + // load state + mjtNum* state = &sim->history_[i * sim->state_size_]; + mj_setState(sim->m_, sim->d_, state, mjSTATE_INTEGRATION); + + // call forward dynamics + mj_forward(sim->m_, sim->d_); +} + +// update an entire section of ui0 +static void mjui0_update_section(mj::Simulate* sim, int section) { + mjui_update(section, -1, &sim->ui0, &sim->uistate, &sim->platform_ui->mjr_context()); +} + // prepare info text void UpdateInfoText(mj::Simulate* sim, const mjModel* m, const mjData* d, char (&title)[mj::Simulate::kMaxFilenameLength], @@ -1369,6 +1390,12 @@ void UiEvent(mjuiState* state) { case 7: // Save key sim->pending_.save_key = true; break; + + case 11: // History scrubber + sim->run = 0; + sim->pending_.load_from_history = true; + mjui0_update_section(sim, SECT_SIMULATION); + break; } } @@ -1404,8 +1431,7 @@ void UiEvent(mjuiState* state) { } else { sim->cam.type = mjCAMERA_FREE; sim->camera = 0; - mjui_update(SECT_RENDERING, -1, &sim->ui0, &sim->uistate, - &sim->platform_ui->mjr_context()); + mjui0_update_section(sim, SECT_RENDERING); } } else { sim->cam.type = mjCAMERA_FIXED; @@ -1477,20 +1503,51 @@ void UiEvent(mjuiState* state) { if (sim->fully_managed_ && sim->m_) { sim->run = 1 - sim->run; sim->pert.active = 0; - mjui_update(-1, -1, &sim->ui0, state, &sim->platform_ui->mjr_context()); + + if (sim->run) sim->scrub_index = 0; // reset scrubber + + mjui0_update_section(sim, -1); } break; case mjKEY_RIGHT: // step forward if (sim->fully_managed_ && sim->m_ && !sim->run) { ClearTimers(sim->d_); - mj_step(sim->m_, sim->d_); + + // currently in scrubber: increment scrub, load state, update slider UI + if (sim->scrub_index < 0) { + sim->scrub_index++; + sim->pending_.load_from_history = true; + mjui0_update_section(sim, SECT_SIMULATION); + } + + // not in scrubber: step + else { + mj_step(sim->m_, sim->d_); + } + UpdateProfiler(sim, sim->m_, sim->d_); UpdateSensor(sim, sim->m_, sim->d_); UpdateSettings(sim, sim->m_); } break; + case mjKEY_LEFT: // step backward + if (sim->fully_managed_ && sim->m_) { + sim->run = 0; + ClearTimers(sim->d_); + + // decrement scrub, load state + sim->scrub_index = mjMAX(sim->scrub_index - 1, 1 - sim->nhistory_); + sim->pending_.load_from_history = true; + + // update slider UI, profiler, sensor + mjui0_update_section(sim, SECT_SIMULATION); + UpdateProfiler(sim, sim->m_, sim->d_); + UpdateSensor(sim, sim->m_, sim->d_); + } + break; + case mjKEY_PAGE_UP: // select parent body if ((sim->m_ || !sim->fully_managed_) && sim->pert.select > 0) { sim->pert.select = sim->body_parentid_[sim->pert.select]; @@ -1514,7 +1571,7 @@ void UiEvent(mjuiState* state) { sim->camera += 1; } sim->cam.fixedcamid = sim->camera - 2; - mjui_update(SECT_RENDERING, -1, &sim->ui0, &sim->uistate, &sim->platform_ui->mjr_context()); + mjui0_update_section(sim, SECT_RENDERING); } break; @@ -1528,28 +1585,28 @@ void UiEvent(mjuiState* state) { sim->camera -= 1; } sim->cam.fixedcamid = sim->camera - 2; - mjui_update(SECT_RENDERING, -1, &sim->ui0, &sim->uistate, &sim->platform_ui->mjr_context()); + mjui0_update_section(sim, SECT_RENDERING); } break; case mjKEY_F6: // cycle frame visualisation if (sim->m_ || !sim->fully_managed_) { sim->opt.frame = (sim->opt.frame + 1) % mjNFRAME; - mjui_update(SECT_RENDERING, -1, &sim->ui0, &sim->uistate, &sim->platform_ui->mjr_context()); + mjui0_update_section(sim, SECT_RENDERING); } break; case mjKEY_F7: // cycle label visualisation if (sim->m_ || !sim->fully_managed_) { sim->opt.label = (sim->opt.label + 1) % mjNLABEL; - mjui_update(SECT_RENDERING, -1, &sim->ui0, &sim->uistate, &sim->platform_ui->mjr_context()); + mjui0_update_section(sim, SECT_RENDERING); } break; case mjKEY_ESCAPE: // free camera sim->cam.type = mjCAMERA_FREE; sim->camera = 0; - mjui_update(SECT_RENDERING, -1, &sim->ui0, &sim->uistate, &sim->platform_ui->mjr_context()); + mjui0_update_section(sim, SECT_RENDERING); break; case '-': // slow down @@ -1693,6 +1750,8 @@ Simulate::Simulate(std::unique_ptr platform_ui, mjv_defaultSceneState(&scnstate_); } +// synchronize model and data +// operations which require holding the mutex, prevents racing with physics thread void Simulate::Sync() { MutexLock lock(this->mtx); @@ -1826,6 +1885,8 @@ void Simulate::Sync() { mj_forward(m_, d_); update_profiler = true; update_sensor = true; + scrub_index = 0; + mjui0_update_section(this, SECT_SIMULATION); pending_.reset = false; } @@ -1839,6 +1900,13 @@ void Simulate::Sync() { pending_.copy_pose = false; } + if (pending_.load_from_history) { + LoadScrubState(this); + update_profiler = true; + update_sensor = true; + pending_.load_from_history = false; + } + if (pending_.load_key) { int i = this->key; d_->time = m_->key_time[i]; @@ -1846,8 +1914,7 @@ void Simulate::Sync() { mju_copy(d_->qvel, m_->key_qvel + i*m_->nv, m_->nv); mju_copy(d_->act, m_->key_act + i*m_->na, m_->na); mju_copy(d_->mocap_pos, m_->key_mpos + i*3*m_->nmocap, 3*m_->nmocap); - mju_copy(d_->mocap_quat, m_->key_mquat + i*4*m_->nmocap, - 4*m_->nmocap); + mju_copy(d_->mocap_quat, m_->key_mquat + i*4*m_->nmocap, 4*m_->nmocap); mju_copy(d_->ctrl, m_->key_ctrl + i*m_->nu, m_->nu); mj_forward(m_, d_); update_profiler = true; @@ -2080,6 +2147,30 @@ void Simulate::LoadOnRenderThread() { std::memcpy(ctrl_.data(), this->d_->ctrl, sizeof(this->d_->ctrl[0]) * this->m_->nu); ctrl_prev_ = ctrl_; + // allocate history buffer: smaller of {2000 states, 100 MB} + if (this->fully_managed_) { + constexpr int kHistoryLength = 2000; + constexpr int kMaxHistoryBytes = 1e8; + + // get state size, size of history buffer + state_size_ = mj_stateSize(this->m_, mjSTATE_INTEGRATION); + int state_bytes = state_size_ * sizeof(mjtNum); + int history_bytes = mjMIN(state_bytes * kHistoryLength, kMaxHistoryBytes); + nhistory_ = history_bytes / state_bytes; + + // allocate history buffer, reset cursor and UI slider + history_.clear(); + history_.resize(nhistory_ * state_size_); + history_cursor_ = 0; + scrub_index = 0; + + // fill buffer with initial state + mj_getState(this->m_, this->d_, history_.data(), mjSTATE_INTEGRATION); + for (int i = 1; i < nhistory_; ++i) { + mju_copy(&history_[i * state_size_], history_.data(), state_size_); + } + } + // re-create scene and context if (this->fully_managed_) { mjv_makeScene(this->m_, &this->scn, kMaxGeom); @@ -2132,6 +2223,10 @@ void Simulate::LoadOnRenderThread() { this->ui0.sect[SECT_SIMULATION].item[5].slider.range[1] = mjMAX(0, this->m_->nkey - 1); this->ui0.sect[SECT_SIMULATION].item[5].slider.divisions = mjMAX(1, this->m_->nkey - 1); + // set scrubber range and divisions + this->ui0.sect[SECT_SIMULATION].item[11].slider.range[0] = 1 - nhistory_; + this->ui0.sect[SECT_SIMULATION].item[11].slider.divisions = nhistory_; + // rebuild UI sections MakeUiSections(this, this->m_, this->d_); @@ -2217,12 +2312,12 @@ void Simulate::Render() { // update UI sections from last sync if (this->ui0_enable && this->ui0.sect[SECT_WATCH].state) { - mjui_update(SECT_WATCH, -1, &this->ui0, &this->uistate, &this->platform_ui->mjr_context()); + mjui0_update_section(this, SECT_WATCH); } if (pending_.ui_update_physics) { if (this->ui0_enable && this->ui0.sect[SECT_PHYSICS].state) { - mjui_update(SECT_PHYSICS, -1, &this->ui0, &this->uistate, &this->platform_ui->mjr_context()); + mjui0_update_section(this, SECT_PHYSICS); } pending_.ui_update_physics = false; } @@ -2244,8 +2339,7 @@ void Simulate::Render() { IsDifferent(opt_prev_.tendongroup, opt.tendongroup) || IsDifferent(opt_prev_.actuatorgroup, opt.actuatorgroup) || IsDifferent(opt_prev_.skingroup, opt.skingroup))) { - mjui_update(SECT_GROUP, -1, &this->ui0, &this->uistate, - &this->platform_ui->mjr_context()); + mjui0_update_section(this, SECT_GROUP); } opt_prev_ = opt; @@ -2254,8 +2348,7 @@ void Simulate::Render() { if (pending_.ui_update_rendering) { if (this->ui0_enable && this->ui0.sect[SECT_RENDERING].state) { - mjui_update(SECT_RENDERING, -1, &this->ui0, &this->uistate, - &this->platform_ui->mjr_context()); + mjui0_update_section(this, SECT_RENDERING); } pending_.ui_update_rendering = false; } @@ -2525,6 +2618,20 @@ void Simulate::RenderLoop() { this->exitrequest.store(2); } +// add state to history buffer +void Simulate::AddToHistory() { + if (history_.empty()) { + return; + } + + // circular increment of cursor + history_cursor_ = (history_cursor_ + 1) % nhistory_; + + // add state at cursor + mjtNum* state = &history_[state_size_ * history_cursor_]; + mj_getState(m_, d_, state, mjSTATE_INTEGRATION); +} + void Simulate::UpdateHField(int hfieldid) { MutexLock lock(this->mtx); if (!m_ || hfieldid < 0 || hfieldid >= m_->nhfield) { diff --git a/simulate/simulate.h b/simulate/simulate.h index a1567c0d..afcb0dda 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -84,6 +84,9 @@ class Simulate { // loop to render the UI (must be called from main thread because of MacOS) void RenderLoop(); + // add state to history buffer + void AddToHistory(); + // constants static constexpr int kMaxFilenameLength = 1000; @@ -100,6 +103,9 @@ class Simulate { 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 body_parentid_; @@ -113,6 +119,8 @@ class Simulate { std::vector>> actuator_ctrlrange_; std::vector actuator_names_; + std::vector history_; // history buffer (nhistory x state_size) + // mjModel and mjData fields that can be modified by the user through the GUI std::vector qpos_; std::vector qpos_prev_; @@ -134,6 +142,7 @@ class Simulate { bool reset; bool align; bool copy_pose; + bool load_from_history; bool load_key; bool save_key; bool zero_ctrl; @@ -171,6 +180,9 @@ class Simulate { // keyframe index int key = 0; + // index of history-scrubber slider + int scrub_index = 0; + // simulation int run = 1; @@ -266,7 +278,7 @@ class Simulate { // simulation section of UI - const mjuiDef def_simulation[12] = { + const mjuiDef def_simulation[14] = { {mjITEM_SECTION, "Simulation", 1, nullptr, "AS"}, {mjITEM_RADIO, "", 5, &this->run, "Pause\nRun"}, {mjITEM_BUTTON, "Reset", 2, nullptr, " #259"}, @@ -278,6 +290,8 @@ class Simulate { {mjITEM_BUTTON, "Save key", 3}, {mjITEM_SLIDERNUM, "Noise scale", 5, &this->ctrl_noise_std, "0 2"}, {mjITEM_SLIDERNUM, "Noise rate", 5, &this->ctrl_noise_rate, "0 2"}, + {mjITEM_SEPARATOR, "History", 1}, + {mjITEM_SLIDERINT, "", 5, &this->scrub_index, "0 0"}, {mjITEM_END} };