diff --git a/python/mujoco/experimental/studio/web/index.html b/python/mujoco/experimental/studio/web/index.html index 2d7a5856..1f12606c 100644 --- a/python/mujoco/experimental/studio/web/index.html +++ b/python/mujoco/experimental/studio/web/index.html @@ -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"; diff --git a/python/mujoco/experimental/studio/web/state_payload.h b/python/mujoco/experimental/studio/web/state_payload.h index c3245e8a..0165cfa8 100644 --- a/python/mujoco/experimental/studio/web/state_payload.h +++ b/python/mujoco/experimental/studio/web/state_payload.h @@ -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); diff --git a/python/mujoco/experimental/studio/web/web_client.cc b/python/mujoco/experimental/studio/web/web_client.cc index 360094f7..21146716 100644 --- a/python/mujoco/experimental/studio/web/web_client.cc +++ b/python/mujoco/experimental/studio/web/web_client.cc @@ -31,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -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 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( - reinterpret_cast(fetch->data), - static_cast(fetch->numBytes)), + std::span(reinterpret_cast(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(fetch->data), - static_cast(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(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(ptr); +} + +// Frees a buffer previously returned by AllocModelBuffer. +void FreeModelBuffer(uintptr_t ptr_val) { + if (ptr_val) { + delete[] reinterpret_cast(ptr_val); + } +} + +// Parses an MJB model from a buffer. +void ParseModelBuffer(uintptr_t ptr_val, size_t size) { + if (ParseModelBufferImpl(reinterpret_cast(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); } diff --git a/python/mujoco/experimental/studio/web/web_client_local_ui.cc b/python/mujoco/experimental/studio/web/web_client_local_ui.cc index f5a202c5..197451c6 100644 --- a/python/mujoco/experimental/studio/web/web_client_local_ui.cc +++ b/python/mujoco/experimental/studio/web/web_client_local_ui.cc @@ -25,6 +25,7 @@ #include #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(view.bytes_downloaded) / + static_cast(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) { diff --git a/python/mujoco/experimental/studio/web/web_client_local_ui.h b/python/mujoco/experimental/studio/web/web_client_local_ui.h index f8aa7a4c..7b1882f9 100644 --- a/python/mujoco/experimental/studio/web/web_client_local_ui.h +++ b/python/mujoco/experimental/studio/web/web_client_local_ui.h @@ -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 diff --git a/python/mujoco/experimental/studio/web/web_client_remote_ui.cc b/python/mujoco/experimental/studio/web/web_client_remote_ui.cc index 944b1d7c..36224107 100644 --- a/python/mujoco/experimental/studio/web/web_client_remote_ui.cc +++ b/python/mujoco/experimental/studio/web/web_client_remote_ui.cc @@ -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(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(tex_id), static_cast(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(entry.pixels.data()), diff --git a/python/mujoco/experimental/studio/web/web_client_remote_ui.h b/python/mujoco/experimental/studio/web/web_client_remote_ui.h index e7daea72..0c9e06ac 100644 --- a/python/mujoco/experimental/studio/web/web_client_remote_ui.h +++ b/python/mujoco/experimental/studio/web/web_client_remote_ui.h @@ -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(); diff --git a/python/mujoco/experimental/studio/web/web_client_session.cc b/python/mujoco/experimental/studio/web/web_client_session.cc index 67d6447d..517cf2eb 100644 --- a/python/mujoco/experimental/studio/web/web_client_session.cc +++ b/python/mujoco/experimental/studio/web/web_client_session.cc @@ -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; } diff --git a/python/mujoco/experimental/studio/web/web_client_session.h b/python/mujoco/experimental/studio/web/web_client_session.h index 312f055b..c66e8dab 100644 --- a/python/mujoco/experimental/studio/web/web_client_session.h +++ b/python/mujoco/experimental/studio/web/web_client_session.h @@ -28,6 +28,7 @@ #include +#include #include #include #include @@ -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 model_crc32_; - bool reload_pending_ = false; int server_close_code_ = 0; diff --git a/python/mujoco/experimental/studio/web/web_server.py b/python/mujoco/experimental/studio/web/web_server.py index d7bc9f56..3266af87 100644 --- a/python/mujoco/experimental/studio/web/web_server.py +++ b/python/mujoco/experimental/studio/web/web_server.py @@ -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": } + # 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 diff --git a/python/mujoco/experimental/studio/web_viewer.py b/python/mujoco/experimental/studio/web_viewer.py index 73736d35..ac47c995 100644 --- a/python/mujoco/experimental/studio/web_viewer.py +++ b/python/mujoco/experimental/studio/web_viewer.py @@ -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)