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

Add history scrubbing mechanism to simulate.

Browse Source 
[screen capture](https://www.youtube.com/watch?v=YSvWn_poqWs)

PiperOrigin-RevId: 570676968
Change-Id: Id9a1fccc9418b851e007ddcf3fa89acd4edff3e9
This commit is contained in:
Yuval Tassa
2023-10-04 06:35:14 -07:00
committed by Copybara-Service
parent 1e94c12cdd
commit 3e12f0d50c
6 changed files with 177 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+16 -2
View File
@@ -129,13 +129,27 @@ Python bindings
22. Fixed `#870 <https://github.com/google-deepmind/mujoco/issues/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<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:
24. 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
^^^^^^^^^
23. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
25. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
midphase collision trees.
Version 2.3.7 (July 20, 2023)
-----------------------------
python/mujoco/simulate.cc
+2 -1
View File
@@ -210,7 +210,8 @@ PYBIND11_MODULE(_simulate, pymodule) {
py::call_guard<py::gil_scoped_release>())
.def("sync", CallIfNotNull(&mujoco::Simulate::Sync),
py::call_guard<py::gil_scoped_release>())
.def("add_to_history", CallIfNotNull(&mujoco::Simulate::AddToHistory),
py::call_guard<py::gil_scoped_release>())
.def("render_loop", CallIfNotNull(&mujoco::Simulate::RenderLoop),
py::call_guard<py::gil_scoped_release>())
.def("lock", GetIfNotNull(&mujoco::Simulate::mtx),
python/mujoco/viewer.py
+8
View File
@@ -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.
simulate/main.cc
+11 -4
View File
@@ -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<mjtNum*>(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
simulate/simulate.cc
+125 -18
View File
@@ -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<PlatformUIAdapter> 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) {
simulate/simulate.h
+15 -1
View File
@@ -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<int> body_parentid_;
@@ -113,6 +119,8 @@ class Simulate {
std::vector<std::optional<std::pair<mjtNum, mjtNum>>> actuator_ctrlrange_;
std::vector<std::string> actuator_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_;
@@ -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}
};