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

Copybara import of the project:

Browse Source 
--
c9acc0a6f677951db34f2b607bd60fc18a43f72b by Levi Burner <leviburner@gmail.com>:

Fix race condition in Python viewers set_X methods

--
b89ab8f7a0f628ee7a02ebfbac2a97f113112076 by Levi Burner <leviburner@gmail.com>:

use std::swap to replace some copying

--
02625bed0cffbfc9ea4537f74689f0bbcb8bf44c by Levi Burner <leviburner@gmail.com>:

fix whitespace

COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/2613 from aftersomemath:simulate-set-race 02625bed0cffbfc9ea4537f74689f0bbcb8bf44c
PiperOrigin-RevId: 766646168
Change-Id: I1e1dc16afbfb1e69fb958d54bd1075cc9aed3569
This commit is contained in:
Levi Burner
2025-06-03 07:00:24 -07:00
committed by Copybara-Service
parent fe81373ffe
commit 3eb31f56cd
3 changed files with 91 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/simulate.cc
+65 -27
View File
@@ -96,7 +96,6 @@ class SimulateWrapper {
void Destroy() {
if (simulate_) {
ClearImages();
delete simulate_;
simulate_ = nullptr;
destroyed_.store(1);
@@ -140,40 +139,71 @@ class SimulateWrapper {
void SetFigures(
const std::vector<std::pair<mjrRect, py::object>>& viewports_figures) {
// Pairs of [viewport, figure], where viewport corresponds to the location
// of the figure on the viewer window.
std::vector<std::pair<mjrRect, mjvFigure>> user_figures;
for (const auto& [viewport, figure] : viewports_figures) {
mjvFigure casted_figure = *figure.cast<MjvFigureWrapper&>().get();
user_figures.push_back(std::make_pair(viewport, casted_figure));
// TODO: replace with atomic wait when we migrate to C++20
while (simulate_ && simulate_->newfigurerequest.load() != 0) {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
// Set them all at once to prevent figure flickering.
simulate_->user_figures_ = user_figures;
// Pairs of [viewport, figure], where viewport corresponds to the location
// of the figure on the viewer window.
for (const auto& [viewport, figure] : viewports_figures) {
mjvFigure casted_figure = *figure.cast<MjvFigureWrapper&>().get();
simulate_->user_figures_new_.push_back(std::make_pair(viewport, casted_figure));
}
int value = 0;
simulate_->newfigurerequest.compare_exchange_strong(value, 1);
}
void ClearFigures() { simulate_->user_figures_.clear(); }
void ClearFigures() {
// TODO: replace with atomic wait when we migrate to C++20
while (simulate_ && simulate_->newfigurerequest.load() != 0) {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
simulate_->user_figures_new_.clear();
int value = 0;
simulate_->newfigurerequest.compare_exchange_strong(value, 1);
}
void SetTexts(
const std::vector<std::tuple<int, int, std::string, std::string>>&
texts) {
// Collection of [font, gridpos, text1, text2] tuples for overlay text
std::vector<std::tuple<int, int, std::string, std::string>> user_texts;
for (const auto& [font, gridpos, text1, text2] : texts) {
user_texts.push_back(std::make_tuple(font, gridpos, text1, text2));
// TODO: replace with atomic wait when we migrate to C++20
while (simulate_ && simulate_->newtextrequest.load() != 0) {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
// Set them all at once to prevent text flickering.
simulate_->user_texts_ = user_texts;
// Collection of [font, gridpos, text1, text2] tuples for overlay text
for (const auto& [font, gridpos, text1, text2] : texts) {
simulate_->user_texts_new_.push_back(std::make_tuple(font, gridpos, text1, text2));
}
int value = 0;
simulate_->newtextrequest.compare_exchange_strong(value, 1);
}
void ClearTexts() { simulate_->user_texts_.clear(); }
void ClearTexts() {
// TODO: replace with atomic wait when we migrate to C++20
while (simulate_ && simulate_->newtextrequest.load() != 0) {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
simulate_->user_texts_new_.clear();
int value = 0;
simulate_->newtextrequest.compare_exchange_strong(value, 1);
}
void SetImages(
const std::vector<std::tuple<mjrRect, pybind11::array&>> viewports_images
) {
// Clear previous images to prevent memory leaks
ClearImages();
// TODO: replace with atomic wait when we migrate to C++20
while (simulate_ && simulate_->newimagerequest.load() != 0) {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
for (const auto& [viewport, image] : viewports_images) {
auto buf = image.request();
@@ -192,20 +222,28 @@ class SimulateWrapper {
size_t width = buf.shape[1];
size_t size = height * width * 3;
// Make a copy of the image data to prevent flickering
unsigned char* image_copy = new unsigned char[size];
std::memcpy(image_copy, buf.ptr, size);
// Make a copy of the image data since Python is
// not required to keep it
std::unique_ptr<unsigned char[]> image_copy(new unsigned char[size]());
std::memcpy(image_copy.get(), buf.ptr, size);
simulate_->user_images_.push_back(std::make_tuple(viewport, image_copy));
simulate_->user_images_new_.push_back(std::make_tuple(viewport, std::move(image_copy)));
}
int value = 0;
simulate_->newimagerequest.compare_exchange_strong(value, 1);
}
void ClearImages() {
// Free memory for each image before clearing the vector
for (const auto& [viewport, image_ptr] : simulate_->user_images_) {
delete[] image_ptr;
// TODO: replace with atomic wait when we migrate to C++20
while (simulate_ && simulate_->newimagerequest.load() != 0) {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
simulate_->user_images_.clear();
simulate_->user_images_new_.clear();
int value = 0;
simulate_->newimagerequest.compare_exchange_strong(value, 1);
}
private: