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

MuJoCo Web Viewer: Implement parallel chunked model downloading and in-place model reloading.

Browse Source 
This change introduces parallel chunked downloading for large model files (.mjb) directly into WASM linear memory, enables model loading without full page reloads and reduces memory overhead.

Key changes:
- Implement a chunked model endpoint in the Python web server to support range requests.
- Add parallel chunked fetching in the frontend with retry logic and a single-fetch fallback.
- Increase initial WASM memory to 3 GB to accommodate large models and prevent heap fragmentation.
- Support in-place model reloading in the C++ client, including texture cache invalidation when the Filament context is recreated.
- Display a model download progress bar and model parsing/loading banner to the UI.
- Fixes model drag and drop (caused by typo in sessionId, corrected to session_id).

PiperOrigin-RevId: 960249236
Change-Id: Icac89e6a4ca099882aaf9b111744c1b6ab0cc6c1
This commit is contained in:
Matija Kecman
2026-08-06 05:51:01 -07:00
committed by Copybara-Service
parent 7fd2061f5c
commit 84950fa371
11 changed files with 406 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/experimental/studio/web/index.html
+169 -9
View File
@@ -45,6 +45,160 @@ limitations under the License.
const WS_CLOSE_INACTIVE = 4003;
const WS_CLOSE_NOT_CONTROLLER = 4004;
// Parallel chunked model download.
//
// Fetches /model in up to PARALLEL concurrent 64 MiB requests
// and streams the chunks directly into a WASM linear heap buffer
// allocated via allocModelBuffer(totalSize). Each chunk is retried up
// to MAX_RETRIES times on transient errors (proxy resets, HTTP2
// protocol errors, etc.). Returns { ptr, size } on success, or null
// on failure.
async function fetchModelChunked() {
const CHUNK = 64 * 1024 * 1024; // 64 MiB per request
const PARALLEL = 6; // concurrent fetches
const MAX_RETRIES = 3; // per-chunk retry limit
const RETRY_DELAY = 500; // ms between retries
let ptr = 0;
try {
const totalSizeResp = await fetch("/model?total_bytes");
if (!totalSizeResp.ok) return null;
const totalSizeData = await totalSizeResp.json();
const totalSize = totalSizeData.total_bytes;
if (!totalSize || totalSize <= 0) return null;
ptr = Module.allocModelBuffer(totalSize);
if (!ptr) {
console.error("[model] failed to allocate WASM buffer of size:", totalSize);
return null;
}
const chunks = [];
for (let offset = 0; offset < totalSize; offset += CHUNK) {
chunks.push({ offset: offset, size: Math.min(CHUNK, totalSize - offset) });
}
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(0, totalSize, 0);
}
let bytesDownloaded = 0;
let maxRetry = 0;
// Bounded-parallelism worker pool with per-chunk retries.
let nextIndex = 0;
let errors = 0;
async function worker() {
while (nextIndex < chunks.length) {
const index = nextIndex++;
const chunk = chunks[index];
const url = "/model?offset_bytes=" + chunk.offset + "&size_bytes=" + chunk.size;
let ok = false;
for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
try {
if (attempt > 0) {
if (attempt > maxRetry) maxRetry = attempt;
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(bytesDownloaded, totalSize, maxRetry);
}
console.log("[model] chunk", index, "retry", attempt);
await new Promise((r) => setTimeout(r, RETRY_DELAY * attempt));
}
const resp = await fetch(url);
if (!resp.ok) continue;
const buffer = await resp.arrayBuffer();
HEAPU8.set(new Uint8Array(buffer), ptr + chunk.offset);
bytesDownloaded += chunk.size;
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(bytesDownloaded, totalSize, maxRetry);
}
ok = true;
break;
} catch (e) {
console.warn("[model] chunk", index, "attempt", attempt, e);
}
}
if (!ok) {
console.error("[model] chunk", index, "failed after", MAX_RETRIES, "retries");
errors++;
}
}
}
const workers = [];
for (let i = 0; i < Math.min(PARALLEL, chunks.length); i++) {
workers.push(worker());
}
await Promise.all(workers);
if (errors > 0) return null; // finally frees ptr
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(totalSize, totalSize, 0);
}
// Allow ~4 frames (60ms) to elapse so the 100% progress bar paints
// before C++ synchronous parsing blocks the thread.
await new Promise((resolve) => setTimeout(resolve, 60));
const result = { ptr: ptr, size: totalSize };
ptr = 0; // transfer ownership to caller
return result;
} catch (e) {
console.error("[model] chunked fetch failed:", e);
return null; // finally frees ptr
} finally {
if (ptr) Module.freeModelBuffer(ptr);
}
}
// Single-fetch fallback if chunked download using fetchModelChunked() fails.
async function fetchModelSingle() {
let ptr = 0;
try {
const resp = await fetch("/model");
if (!resp.ok) return null;
const buffer = await resp.arrayBuffer();
const totalSize = buffer.byteLength;
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(0, totalSize, 0);
}
ptr = Module.allocModelBuffer(totalSize);
if (!ptr) return null;
HEAPU8.set(new Uint8Array(buffer), ptr);
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(totalSize, totalSize, 0);
}
// Allow ~4 frames (60ms) to elapse so the 100% progress bar paints
// before C++ synchronous parsing blocks the thread.
await new Promise((resolve) => setTimeout(resolve, 60));
const result = { ptr: ptr, size: totalSize };
ptr = 0; // transfer ownership to caller
return result;
} catch (e) {
console.error("[model] single fetch failed:", e);
return null; // finally frees ptr
} finally {
if (ptr) Module.freeModelBuffer(ptr);
}
}
// Called from C++ (via EM_ASM) when the Python side swaps the model.
async function reloadModel() {
// Retry the chunked download up to 3 times with increasing delay.
// Falling back to a single-fetch for large models hits the same
// proxy errors, so retrying chunks is the only viable path.
for (let attempt = 0; attempt < 3; attempt++) {
if (attempt > 0) {
console.log("[model] retrying chunked download, attempt", attempt + 1);
await new Promise((r) => setTimeout(r, 1000 * attempt));
}
const modelRes = await fetchModelChunked();
if (modelRes) {
Module.parseModelBuffer(modelRes.ptr, modelRes.size);
Module.freeModelBuffer(modelRes.ptr);
return;
}
}
console.error("[model] all chunked download attempts failed; trying single-fetch");
const modelRes = await fetchModelSingle();
if (modelRes) {
Module.parseModelBuffer(modelRes.ptr, modelRes.size);
Module.freeModelBuffer(modelRes.ptr);
}
}
var Module = {
canvas: (() => {
const canvas = document.getElementById("canvas");
@@ -92,20 +246,26 @@ limitations under the License.
];
const assetPromises = assetsToPrefetch.map(async (filename) => {
try {
const response = await fetch("assets/" + filename);
if (!response.ok) {
console.error(`Failed to fetch asset ${filename}: ${response.statusText}`);
for (let attempt = 0; attempt < 3; attempt++) {
try {
if (attempt > 0) {
await new Promise((r) => setTimeout(r, 500 * attempt));
}
const response = await fetch("assets/" + filename);
if (!response.ok) continue;
const buffer = await response.arrayBuffer();
Module.registerAsset(filename, new Uint8Array(buffer));
return;
} catch (error) {
console.warn(`Attempt ${attempt + 1} failed for ${filename}:`, error);
}
const buffer = await response.arrayBuffer();
Module.registerAsset(filename, new Uint8Array(buffer));
} catch (error) {
console.error(`Error prefetching asset ${filename}:`, error);
}
console.error(`Error prefetching asset ${filename} after 3 attempts`);
});
Promise.all(assetPromises).then(() => {
Module.startApp();
reloadModel();
});
},
};
@@ -186,7 +346,7 @@ limitations under the License.
const proto = location.protocol === "https:" ? "wss://" : "ws://";
// The session id tells the server which page is dropping; it only
// accepts drops from the controller.
const sid = (window.Module && Module.sessionId) || "";
const sid = (window.Module && Module.session_id) || "";
const dropUrl = proto + location.host + "/drop?sid=" + encodeURIComponent(sid);
const ws = new WebSocket(dropUrl);
ws.binaryType = "arraybuffer";
python/mujoco/experimental/studio/web/state_payload.h
+2 -2
View File
@@ -20,7 +20,7 @@
//
// The arguments come from the Python process:
//
// * model : fetched once over HTTP as /model.mjb; its runtime-mutable
// * model : fetched once over HTTP as /model; its runtime-mutable
// parts (opt/vis/stat) are re-sent in the render state block.
// * data : streamed as the physics state vector (mjSTATE_INTEGRATION);
// the browser recomputes the rest via mj_setState/mj_forward.
@@ -64,7 +64,7 @@ struct StatePayloadHeader {
uint16_t version = kStatePayloadVersion;
uint16_t nblocks = 0;
// CRC32 of the model's MJB bytes. When this changes, the browser must
// refetch /model.mjb before applying any further state.
// refetch /model before applying any further state.
uint32_t model_crc32 = 0;
};
static_assert(sizeof(StatePayloadHeader) == 12);
python/mujoco/experimental/studio/web/web_client.cc
+112 -50
View File
@@ -31,6 +31,7 @@
#include <fstream>
#include <initializer_list>
#include <memory>
#include <new>
#include <span>
#include <string>
#include <string_view>
@@ -89,6 +90,13 @@ struct Telemetry {
uint64_t sim_bytes_per_sec = 0;
};
struct ModelDownloadStatus {
bool is_downloading = true;
size_t bytes_downloaded = 0;
size_t total_bytes = 0;
int retry_count = 0;
};
// The implementation of every interface needed by the session and remote UI.
class AppCallbacks final : public RemoteUi::Callbacks,
public Session::Callbacks {
@@ -101,6 +109,7 @@ class AppCallbacks final : public RemoteUi::Callbacks,
// Session::Callbacks
bool ReadyForPayload() override;
void OnPayload(const StatePayloadView& view) override;
void OnModelChanged() override;
void ConnectRemoteUi() override;
void ShutdownRemoteUi() override;
void SetCameraMode(int mode) override;
@@ -133,6 +142,8 @@ struct App {
// User-injected geoms received with the state payload.
std::vector<mjvGeom> extra_geoms;
ModelDownloadStatus download_status;
AppCallbacks callbacks;
RemoteUi remote_ui{callbacks};
Session session{callbacks};
@@ -377,6 +388,13 @@ void AppCallbacks::OnPayload(const StatePayloadView& view) {
ApplyStatePayload(view);
}
void AppCallbacks::OnModelChanged() {
// Mark downloading immediately so the main loop stops applying state updates
// to the old model while the new one downloads asynchronously.
g_app.download_status.is_downloading = true;
EM_ASM({ reloadModel(); });
}
void AppCallbacks::ConnectRemoteUi() { g_app.remote_ui.Connect(WsUrl("/ui")); }
void AppCallbacks::ShutdownRemoteUi() { g_app.remote_ui.Shutdown(); }
@@ -396,7 +414,7 @@ void BuildBrowserGui() {
const double stale_sec =
last_msg > 0 ? emscripten_get_now() / 1000.0 - last_msg : -1.0;
g_app.disconnect_notice.Draw(g_app.session.ServerCloseCode(), stale_sec,
g_app.session.ReloadPending());
g_app.download_status.is_downloading);
SessionView view;
g_app.session.FillView(&view);
@@ -404,6 +422,10 @@ void BuildBrowserGui() {
view.sim_bytes_per_sec = g_app.telemetry.sim_bytes_per_sec;
view.have_remote_frame = g_app.remote_ui.RemoteDrawData() != nullptr;
view.camera_mode = g_app.spectator_cam_mode;
view.is_downloading = g_app.download_status.is_downloading;
view.bytes_downloaded = g_app.download_status.bytes_downloaded;
view.total_bytes = g_app.download_status.total_bytes;
view.retry_count = g_app.download_status.retry_count;
// The session is the SessionActions implementation: the role window's
// intents land there directly.
g_app.role_window.Draw(view, g_app.session);
@@ -442,8 +464,8 @@ void MainLoopImpl() {
// using a gentler pace.
const int state_retry_interval =
g_app.session.ServerCloseCode() != 0 ? 300 : 60;
if (!g_app.session.HasSocket() && !g_app.session.ReloadPending() &&
g_app.model_holder && g_app.model_holder->ok() &&
if (!g_app.session.HasSocket() && g_app.model_holder &&
g_app.model_holder->ok() && !g_app.download_status.is_downloading &&
g_app.frame_count - g_app.last_state_retry_frame > state_retry_interval) {
g_app.last_state_retry_frame = g_app.frame_count;
LOG(Info, "State WebSocket down; reconnecting...");
@@ -545,7 +567,7 @@ void MainLoopImpl() {
// Render the scene and all ImGui UI (local and remote).
if (g_app.renderer && g_app.model_holder && g_app.model_holder->ok()) {
// Apply backend state if available.
if (g_app.backend_state_dirty) {
if (g_app.backend_state_dirty && !g_app.download_status.is_downloading) {
mjModel* model = g_app.model_holder->model();
mjData* data = g_app.model_holder->data();
mj_setState(model, data, g_app.backend_state.data(),
@@ -572,9 +594,9 @@ void MainLoopImpl() {
void SetupScene(const mjModel* m) {
g_app.renderer->Init(m);
// Upload any textures (e.g. the font atlas) that were buffered before the
// Filament context was available.
g_app.remote_ui.FlushPendingTextures();
// Invalidate all texture caches (ImGui font atlas + NetImgui streamed
// textures) and re-upload them on the new Filament context.
g_app.remote_ui.UpdateTextures();
mjv_defaultPerturb(&g_app.perturb);
mjv_defaultCamera(&g_app.camera);
@@ -602,52 +624,69 @@ uint32_t Crc32(const uint8_t* data, size_t len) {
return crc ^ 0xFFFFFFFFu;
}
void OnFetchSuccess(emscripten_fetch_t* fetch) {
LOG(Info, "Fetched model.mjb, size: %llu", fetch->numBytes);
bool ParseModelBufferImpl(const char* data, size_t size) {
LOG(Info, "Fetched model.mjb, size: %zu", size);
g_app.model_holder = mujoco::platform::ModelHolder::FromBuffer(
std::span<const std::byte>(
reinterpret_cast<const std::byte*>(fetch->data),
static_cast<size_t>(fetch->numBytes)),
std::span<const std::byte>(reinterpret_cast<const std::byte*>(data),
size),
"application/mjb", "model.mjb");
if (g_app.model_holder && g_app.model_holder->ok()) {
LOG(Info, "Model loaded successfully!");
mj_forward(g_app.model_holder->model(), g_app.model_holder->data());
if (g_app.window) {
SetupScene(g_app.model_holder->model());
}
// Baseline the state link on the model we just fetched, so a model swap
// that raced this fetch is caught by the very first payload.
g_app.session.SetModelCrc32(
Crc32(reinterpret_cast<const uint8_t*>(fetch->data),
static_cast<size_t>(fetch->numBytes)));
// Connect the state WebSocket to receive simulation state from Python.
g_app.session.Connect(WsUrl("/state"));
} else {
LOG(Error, "Failed to load model: %s",
g_app.model_holder ? g_app.model_holder->error().data()
: "Unknown error");
Crc32(reinterpret_cast<const uint8_t*>(data), size));
LOG(Info, "Model parsed successfully!");
return true;
}
emscripten_fetch_close(fetch);
LOG(Error, "Failed to load model: %s",
g_app.model_holder ? g_app.model_holder->error().data()
: "Unknown error");
return false;
}
void StartModelFetch();
void OnFetchError(emscripten_fetch_t* fetch) {
LOG(Error, "Failed to fetch model.mjb, status: %d; retrying", fetch->status);
emscripten_fetch_close(fetch);
// The most likely cause is the Python side restarting its server after a
// model change, so retry rather than leaving the page permanently blank.
emscripten_async_call([](void*) { StartModelFetch(); }, nullptr, 1000);
void UpdateModelDownloadProgress(size_t bytes_downloaded, size_t total_bytes,
int retry_count) {
g_app.download_status.is_downloading = true;
g_app.download_status.bytes_downloaded = bytes_downloaded;
g_app.download_status.total_bytes = total_bytes;
g_app.download_status.retry_count = retry_count;
}
void StartModelFetch() {
emscripten_fetch_attr_t attr;
emscripten_fetch_attr_init(&attr);
strcpy(attr.requestMethod, "GET");
attr.attributes = EMSCRIPTEN_FETCH_LOAD_TO_MEMORY;
attr.onsuccess = OnFetchSuccess;
attr.onerror = OnFetchError;
emscripten_fetch(&attr, "/model.mjb");
void FinishModelLoad() {
mj_forward(g_app.model_holder->model(), g_app.model_holder->data());
if (g_app.window) {
SetupScene(g_app.model_holder->model());
}
// Connect the state WebSocket to receive simulation state from Python.
if (!g_app.session.HasSocket()) {
g_app.session.Connect(WsUrl("/state"));
}
g_app.download_status.is_downloading = false;
LOG(Info, "Model loaded and scene initialized successfully!");
}
// Allocates a buffer in WASM linear memory for zero-copy streaming chunk
// downloads. The caller owns the returned pointer and must free it via
// FreeModelBuffer when done.
uintptr_t AllocModelBuffer(size_t size) {
char* ptr = new (std::nothrow) char[size];
if (!ptr) {
LOG(Error, "Failed to allocate WASM model buffer of size %zu", size);
return 0;
}
return reinterpret_cast<uintptr_t>(ptr);
}
// Frees a buffer previously returned by AllocModelBuffer.
void FreeModelBuffer(uintptr_t ptr_val) {
if (ptr_val) {
delete[] reinterpret_cast<char*>(ptr_val);
}
}
// Parses an MJB model from a buffer.
void ParseModelBuffer(uintptr_t ptr_val, size_t size) {
if (ParseModelBufferImpl(reinterpret_cast<char*>(ptr_val), size)) {
FinishModelLoad();
}
}
// Holds assets (Filament assets and local ImGui fonts) that the page fetches
@@ -705,9 +744,7 @@ static void RegisterAssetProviders() {
}
// Starts the viewer once the page has registered every asset. Exposed to JS
// and called from index.html after the fetches complete. Note that calling
// emscripten_set_main_loop with simulate_infinite_loop=1 never returns, so
// there is no post-loop cleanup.
// and called from index.html after the fetches complete.
void StartApp() {
mujoco::platform::Window::Config config;
config.gfx_mode = mujoco::platform::GraphicsMode::FilamentWebGl;
@@ -723,15 +760,20 @@ void StartApp() {
g_app.renderer = new mujoco::platform::Renderer(
g_app.window->GetNativeWindowHandle(), config.gfx_mode);
// Initialize an empty dummy scene so Filament and ImGui are ready to render
// the "DOWNLOADING..." progress bar while /model downloads asynchronously
g_app.model_holder = mujoco::platform::ModelHolder::FromSpec(mj_makeSpec());
if (g_app.model_holder && g_app.model_holder->ok()) {
SetupScene(g_app.model_holder->model());
}
NetImgui::Internal::Network::Startup();
SDL_StartTextInput();
// The UI stream is a path on the page's own host and port.
g_app.remote_ui.Connect(WsUrl("/ui"));
StartModelFetch();
emscripten_set_main_loop(MainLoop, 0, 1);
emscripten_set_main_loop(MainLoop, 0, 0);
}
int main(int argc, char** argv) {
@@ -740,6 +782,26 @@ int main(int argc, char** argv) {
}
EMSCRIPTEN_BINDINGS(web_client_bindings) {
// Registers staged static files fetched by index.html into the C++ in-memory
// AssetRegistry before startApp() runs.
emscripten::function("registerAsset", &RegisterAsset);
// Allocates a buffer in WASM linear memory for zero-copy streaming model
// chunk downloads. The caller must free it via freeModelBuffer when done.
emscripten::function("allocModelBuffer", &AllocModelBuffer);
// Frees a buffer previously allocated by allocModelBuffer.
emscripten::function("freeModelBuffer", &FreeModelBuffer);
// Parses a completed MJB model buffer and reinitializes the Filament scene.
emscripten::function("parseModelBuffer", &ParseModelBuffer);
// Initializes the Filament window, ImGui context, WebSocket connections, and
// starts the Emscripten main simulation loop.
emscripten::function("startApp", &StartApp);
// Updates the model download progress in C++ so the ImGui role window can
// display a real-time progress bar and status while downloading.
emscripten::function("updateModelDownloadProgress",
&UpdateModelDownloadProgress);
}
python/mujoco/experimental/studio/web/web_client_local_ui.cc
+37 -7
View File
@@ -25,6 +25,7 @@
#include <imgui.h>
#include "experimental/platform/ux/imgui_widgets.h"
#include "google/logging.h"
#include "web_client_session.h"
namespace mujoco::studio {
namespace {
@@ -56,7 +57,7 @@ void DrawDisconnectWindow(const char* window_id,
void DisconnectNotice::Draw(int server_close_code,
double seconds_since_last_payload,
bool reload_pending) {
bool is_downloading) {
if (server_close_code != 0) {
const char* reason = "Disconnected by the viewer.";
if (server_close_code == kWsCloseSessionFull) {
@@ -73,10 +74,10 @@ void DisconnectNotice::Draw(int server_close_code,
// silent, but only a killed one closes its sockets. Clears itself when
// traffic resumes. Suppressed before the first payload (negative
// staleness), after a deliberate server-side close (its own notice
// above), and during model-swap reloads.
// above), and while a new model is downloading.
const bool link_stale =
seconds_since_last_payload > kServerSilenceNoticeSec &&
server_close_code == 0 && !reload_pending;
server_close_code == 0 && !is_downloading;
if (!link_stale) {
logged_ = false;
return;
@@ -215,12 +216,17 @@ void RoleWindow::DrawCollapsed(const SessionView& view) {
ImGui::Begin("Role", nullptr,
ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoResize |
ImGuiWindowFlags_AlwaysAutoResize);
if (view.role == SessionRole::kSpectating) {
if (view.is_downloading) {
bool done_downloading =
view.total_bytes > 0 && view.bytes_downloaded >= view.total_bytes;
ImGui::TextColored(kConnectingColor,
done_downloading ? "PARSING..." : "DOWNLOADING...");
} else if (view.role == SessionRole::kSpectating) {
ImGui::TextColored(kSpectatingColor, "SPECTATING");
} else if (view.role == SessionRole::kControlling) {
ImGui::TextColored(kControllingColor, "CONTROLLING");
} else {
ImGui::TextColored(kConnectingColor, "CONNECTING");
ImGui::TextColored(kConnectingColor, "CONNECTING...");
}
// Keep the window expanded while focused or being dragged, so dragging out
// of the collapsed bounds doesn't instantly collapse it midway through
@@ -242,10 +248,34 @@ void RoleWindow::DrawExpanded(const SessionView& view,
ImGui::Begin("Role", nullptr,
ImGuiWindowFlags_AlwaysAutoResize | ImGuiWindowFlags_NoTitleBar);
if (view.role == SessionRole::kClaiming) {
if (view.is_downloading) {
bool done_downloading =
view.total_bytes > 0 && view.bytes_downloaded >= view.total_bytes;
platform::CenteredBanner(done_downloading ? "PARSING..." : "DOWNLOADING...",
kConnectingColor);
if (!done_downloading) {
float progress = view.total_bytes > 0
? static_cast<float>(view.bytes_downloaded) /
static_cast<float>(view.total_bytes)
: 0.0f;
char buf[128];
if (view.total_bytes > 0) {
snprintf(buf, sizeof(buf), "%.1f / %.1f MB (%.0f%%)",
view.bytes_downloaded / (1024.0 * 1024.0),
view.total_bytes / (1024.0 * 1024.0), progress * 100.0f);
} else {
snprintf(buf, sizeof(buf), "Connecting...");
}
ImGui::ProgressBar(progress, ImVec2(-1, 0), buf);
if (view.retry_count > 0) {
ImGui::TextColored(ImVec4(1.0f, 0.8f, 0.2f, 1.0f),
"Retrying chunk (%d)...", view.retry_count);
}
}
} else if (view.role == SessionRole::kClaiming) {
// The first roster or /ui claim outcome resolves this within a few
// hundred milliseconds of page load.
platform::CenteredBanner("CONNECTING", kConnectingColor);
platform::CenteredBanner("CONNECTING...", kConnectingColor);
ImGui::Separator();
DataRateLines(view);
} else if (view.role == SessionRole::kSpectating) {
python/mujoco/experimental/studio/web/web_client_local_ui.h
+1 -1
View File
@@ -34,7 +34,7 @@ class DisconnectNotice {
public:
// seconds_since_last_payload is negative before the first payload.
void Draw(int server_close_code, double seconds_since_last_payload,
bool reload_pending);
bool is_downloading);
private:
// How long the state stream (~60Hz while the Python side is alive) may go
python/mujoco/experimental/studio/web/web_client_remote_ui.cc
+15 -3
View File
@@ -241,14 +241,26 @@ void RemoteUi::ReceiveAndProcessCommands(int frame) {
VLOG(2, "Frame %d: processed %d commands", frame, cmds_this_frame);
}
void RemoteUi::FlushPendingTextures() {
void RemoteUi::UpdateTextures() {
// Clear cached GPU texture IDs so all textures are re-uploaded.
for (auto& [tex_id, local_tex] : texture_map_) {
local_tex = 0;
}
// Reset the ImGui font atlas so it gets re-created on the new context.
if (ImGui::GetCurrentContext() && ImGui::GetIO().Fonts &&
ImGui::GetIO().Fonts->TexData) {
ImGui::GetIO().Fonts->TexData->SetStatus(ImTextureStatus_WantCreate);
}
// Re-upload all CPU-buffered textures (font atlas, streamed UI images).
for (auto& [tex_id, entry] : texture_cpu_) {
uintptr_t& local_tex = texture_map_[tex_id];
if (local_tex == 0 && !entry.pixels.empty()) {
local_tex = callbacks_.UploadTexture(
local_tex, reinterpret_cast<const std::byte*>(entry.pixels.data()),
entry.width, entry.height);
LOG(Info, "FlushPendingTextures: uploaded tex_id=%lu -> filament=%lu",
LOG(Info, "UpdateTextures: uploaded tex_id=%lu -> filament=%lu",
static_cast<unsigned long>(tex_id),
static_cast<unsigned long>(local_tex));
}
@@ -370,7 +382,7 @@ void RemoteUi::ProcessCmdTexture(CmdTexture* cmd_texture) {
// Upload the full CPU-side texture to the GPU if the context is ready.
// If it isn't yet (model still loading), the texture stays in texture_cpu_
// and will be flushed by FlushPendingTextures() later.
// and will be flushed by UpdateTextures() later.
if (callbacks_.GpuReady()) {
local_tex = callbacks_.UploadTexture(
local_tex, reinterpret_cast<const std::byte*>(entry.pixels.data()),
python/mujoco/experimental/studio/web/web_client_remote_ui.h
+3 -3
View File
@@ -98,9 +98,9 @@ class RemoteUi {
// queue this reads (note that ImGui::Render() calls EndFrame() implicitly).
void CaptureAndSendInput();
// Uploads CPU-buffered textures (e.g. the font atlas) that arrived before
// the GPU context became available.
void FlushPendingTextures();
// Invalidates cached GPU texture IDs and re-uploads all CPU-buffered
// textures (font atlas, streamed UI images) after a Filament context reset.
void UpdateTextures();
void Shutdown();
python/mujoco/experimental/studio/web/web_client_session.cc
+7 -5
View File
@@ -226,9 +226,9 @@ void Session::Update() {
void Session::HandleRemoteUiState(RemoteUiState state, int close_code) {
remote_ui_state_ = state;
// No stream to manage, nothing left to claim (settled spectator), or the
// session itself is down or reloading — the /state policies rule then.
// session itself is down — the /state policies rule then.
if (state == RemoteUiState::kNoSocket || role_ == SessionRole::kSpectating ||
reload_pending_ || server_close_code_ != 0) {
server_close_code_ != 0) {
return;
}
const double now = emscripten_get_now() / 1000.0;
@@ -265,7 +265,7 @@ void Session::HandleRemoteUiState(RemoteUiState state, int close_code) {
void Session::HandleMessage(const uint8_t* data, uint32_t num_bytes) {
last_message_time_ = emscripten_get_now() / 1000.0;
if (reload_pending_ || !callbacks_.ReadyForPayload()) {
if (!callbacks_.ReadyForPayload()) {
return;
}
bytes_accum_ += num_bytes;
@@ -281,8 +281,10 @@ void Session::HandleMessage(const uint8_t* data, uint32_t num_bytes) {
} else if (view.model_crc32 != *model_crc32_) {
LOG(Info, "Model changed on the Python side (ident %u -> %u); reloading",
*model_crc32_, view.model_crc32);
reload_pending_ = true;
EM_ASM({ setTimeout(function() { location.reload(); }, 0); });
// Set the new CRC so subsequent payloads don't re-trigger while the new
// model is being fetched.
model_crc32_ = view.model_crc32;
callbacks_.OnModelChanged();
return;
}
python/mujoco/experimental/studio/web/web_client_session.h
+10 -8
View File
@@ -28,6 +28,7 @@
#include <emscripten/websocket.h>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <string>
@@ -65,6 +66,10 @@ struct SessionView {
uint64_t sim_bytes_per_sec = 0;
bool have_remote_frame = false;
int camera_mode = 0; // [SpectatorCamMode].
bool is_downloading = true;
size_t bytes_downloaded = 0;
size_t total_bytes = 0;
int retry_count = 0;
};
// User intent reported by the role window. Session implements this; the
@@ -115,6 +120,8 @@ class Session : public SessionActions {
virtual bool ReadyForPayload() = 0;
// Applies a parsed payload to the application.
virtual void OnPayload(const StatePayloadView& view) = 0;
// Server swapped models; fetch and load the new one in-place.
virtual void OnModelChanged() = 0;
// Role transition: claim the controller slot.
virtual void ConnectRemoteUi() = 0;
// Role transition: drop the stream when spectating
@@ -139,10 +146,6 @@ class Session : public SessionActions {
// True only while the WebSocket is actually open.
bool Connected() const { return connected_; }
// True once a payload with a new model has scheduled a page reload; all
// traffic is dropped from then on.
bool ReloadPending() const { return reload_pending_; }
// The close code from the server deliberately ending this connection (codes
// 4000-4999, e.g. kWsCloseSessionFull), else 0. Such conditions are transient
// (a slot frees up, the user returns to the tab), so the page shows a notice
@@ -221,11 +224,10 @@ class Session : public SessionActions {
EMSCRIPTEN_WEBSOCKET_T socket_ = 0;
bool connected_ = false;
// CRC32 of the model this page loaded. When the payload's crc changes, the
// Python side has swapped models so reload the page; this refetches
// /model.mjb and reconnects everything.
// CRC32 of the model this page loaded. When the payload's CRC changes, the
// Python side has swapped models so we trigger an in-place reload that
// refetches /model and reinitializes the scene.
std::optional<uint32_t> model_crc32_;
bool reload_pending_ = false;
int server_close_code_ = 0;
python/mujoco/experimental/studio/web/web_server.py
+47 -7
View File
@@ -16,7 +16,7 @@
This server runs in a child process with one asyncio loop on a single public
port:
* Plain HTTP GET serves static files (index.html, WASM, assets), /model.mjb.
* Plain HTTP GET serves static files (index.html, WASM, assets), /model.
* WebSocket /ui serves the bridge to the headless NetImgui client, which
connects over loopback TCP (see headless_ui.cc).
* WebSocket /state serves the latest-wins state payload broadcast at ~60Hz
@@ -40,6 +40,7 @@ import asyncio
import ctypes
import datetime
import enum
import json
import logging
import multiprocessing
import multiprocessing.queues
@@ -50,6 +51,7 @@ import socket
import struct
import sys
import threading
import urllib.parse
from typing import Any, Awaitable, Callable, Optional, cast
from websockets.asyncio.server import serve
@@ -123,7 +125,7 @@ _NETIMGUI_CMD_VERSION_SIZE = 120
# reconnects within ~1s in the normal case; this only bounds the pathological
# one where it never appears (e.g., the headless UI failed to start), so a stuck
# controller cannot lock every other browser out forever.
_UI_TCP_WAIT_SEC = 15.0
_UI_TCP_WAIT_SEC = 30.0
# Content types for the static files the HTTP handler serves.
_CONTENT_TYPES = {
@@ -446,7 +448,7 @@ def _run_server(
def _serve_http(path: str) -> Response:
"""Builds the HTTP response for a non-WebSocket GET request."""
if path == "/model.mjb":
if path == "/model":
if not mjb_data:
return Response(404, "Not Found", Headers(), b"no model\n")
# The model changes on hot-swap; never serve a cached copy.
@@ -947,9 +949,44 @@ def _run_server(
connection: ServerConnection, request: Request
) -> Optional[Response]:
del connection
path = request.path.split("?")[0]
# Split path and query string.
if "?" in request.path:
path, query_string = request.path.split("?", 1)
else:
path, query_string = request.path, ""
if path in ("/ui", "/state", "/drop"):
return None # Proceed with the WebSocket handshake.
# Chunked model endpoint: the client fetches /model in parallel chunks
#
# GET /model?total_bytes -> {"total_bytes": <n>}
# GET /model?offset_bytes=X&size_bytes=Y -> bytes [X, X+Y)
#
# Full model endpoint: the client fetches /model in a single request
#
# GET /model -> full model bytes
if path == "/model" and mjb_data and query_string:
params = urllib.parse.parse_qs(query_string)
if "total_bytes" in params or query_string == "total_bytes":
body = json.dumps({"total_bytes": len(mjb_data)}).encode()
headers = _http_headers(
"application/json", len(body), cacheable=False
)
return Response(200, "OK", headers, body)
try:
offset = int(params.get("offset_bytes", [0])[0])
size = int(params.get("size_bytes", [0])[0])
except (ValueError, TypeError):
return Response(400, "Bad Request", Headers(), b"bad params\n")
if size <= 0 or offset < 0 or offset >= len(mjb_data):
return Response(400, "Bad Request", Headers(), b"bad range\n")
chunk = mjb_data[offset : min(offset + size, len(mjb_data))]
headers = _http_headers(
"application/octet-stream", len(chunk), cacheable=False
)
return Response(200, "OK", headers, chunk)
return _serve_http(path)
def process_response(
@@ -1089,7 +1126,10 @@ def _run_server(
process_request=process_request,
process_response=process_response,
compression=None,
close_timeout=1.0,
ping_interval=None,
ping_timeout=None,
open_timeout=None,
close_timeout=None,
# Big enough for model files uploaded via /drop.
max_size=2**26,
)
@@ -1097,7 +1137,7 @@ def _run_server(
tcp_port = tcp_sock.getsockname()[1]
logger.debug(
"[Http] Serving on http://%s:%d "
"(/, /model.mjb, /ui, /state; NetImgui TCP on 127.0.0.1:%d)",
"(/, /model, /ui, /state; NetImgui TCP on 127.0.0.1:%d)",
http_host,
http_port,
tcp_port,
@@ -1147,7 +1187,7 @@ class WebServer:
http_sock: The listening socket for HTTP and WebSocket traffic.
tcp_sock: The listening socket for NetImgui traffic.
static_files_dir: The directory containing the static files to serve.
mjb_data: The model data to serve from /model.mjb.
mjb_data: The model data to serve from /model.
max_payload_size: The maximum size of the state payload.
drop_queue: The queue to put dropped files onto.
controller_sid_shared: The shared value containing the controller's
python/mujoco/experimental/studio/web_viewer.py
+3 -3
View File
@@ -174,7 +174,7 @@ class WebViewer(viewer_protocol.Viewer):
extra_geoms: List of extra geoms. Internal list is created if None.
host: Public interface the server binds to. The default "::" accepts both
IPv6 and IPv4 connections (IPv4-only where IPv6 is unavailable).
http_port: The single public port: page, WASM, /model.mjb, and the /ui and
http_port: The single public port: page, WASM, /model, and the /ui and
/state WebSocket paths. None falls back to config.http_port; 0 picks the
first free port starting at 8080, so several viewers can run side by
side.
@@ -256,13 +256,13 @@ class WebViewer(viewer_protocol.Viewer):
"""Starts (or restarts) the web server, serving the current model."""
self._stop_servers()
# Serialize the compiled model to MJB bytes (served as /model.mjb).
# Serialize the compiled model to MJB bytes (served as /model).
buffer = np.empty(mujoco.mj_sizeModel(self.model), np.uint8)
mujoco.mj_saveModel(self.model, None, buffer)
mjb_data = buffer.tobytes()
# Identity of the served model, included in every state payload. When it
# changes, the browser refetches /model.mjb by reloading the page.
# changes, the browser refetches /model by reloading the page.
self._model_crc32 = zlib.crc32(mjb_data)
state_sig = int(mujoco.mjtState.mjSTATE_INTEGRATION)