diff --git a/python/mujoco/experimental/studio/web/index.html b/python/mujoco/experimental/studio/web/index.html new file mode 100644 index 00000000..dff61b65 --- /dev/null +++ b/python/mujoco/experimental/studio/web/index.html @@ -0,0 +1,207 @@ + + + + + + MuJoCo Web Viewer + + + + + + + + + diff --git a/python/mujoco/experimental/studio/web/state_payload.cc b/python/mujoco/experimental/studio/web/state_payload.cc index 608f881b..f5896778 100644 --- a/python/mujoco/experimental/studio/web/state_payload.cc +++ b/python/mujoco/experimental/studio/web/state_payload.cc @@ -46,6 +46,7 @@ void AppendStateBlock(std::vector& buffer, uint32_t tag, } // Serializes the render state (exactly kRenderStateSize bytes) into `ptr`. +// Must copy the same fields in the same order as ParseRenderState. void SerializeRenderStateInto(std::byte* ptr, const mjvCamera& camera, const mjvPerturb& perturb, const mjvOption& vis_options, const mjOption& opt, @@ -78,6 +79,32 @@ void SerializeRenderStateInto(std::byte* ptr, const mjvCamera& camera, } // namespace +// Parses a render state block produced by SerializeStatePayload. +// Must copy the same fields in the same order as SerializeRenderStateInto. +void ParseRenderState(const std::byte* data, RenderStateView* out) { + const std::byte* ptr = data; + + memcpy(&out->camera, ptr, sizeof(mjvCamera)); + ptr += sizeof(mjvCamera); + + memcpy(&out->perturb, ptr, sizeof(mjvPerturb)); + ptr += sizeof(mjvPerturb); + + memcpy(&out->vis_options, ptr, sizeof(mjvOption)); + ptr += sizeof(mjvOption); + + memcpy(&out->opt, ptr, sizeof(mjOption)); + ptr += sizeof(mjOption); + + memcpy(&out->vis, ptr, sizeof(mjVisual)); + ptr += sizeof(mjVisual); + + memcpy(&out->stat, ptr, sizeof(mjStatistic)); + ptr += sizeof(mjStatistic); + + memcpy(out->render_flags, ptr, mjNRNDFLAG); +} + size_t MaxStatePayloadSize(size_t physics_bytes) { return sizeof(StatePayloadHeader) + 3 * sizeof(StateBlockHeader) + (sizeof(int32_t) + physics_bytes) + kRenderStateSize + diff --git a/python/mujoco/experimental/studio/web/state_payload.h b/python/mujoco/experimental/studio/web/state_payload.h index 0ceacf92..c3245e8a 100644 --- a/python/mujoco/experimental/studio/web/state_payload.h +++ b/python/mujoco/experimental/studio/web/state_payload.h @@ -123,6 +123,21 @@ struct StatePayloadView { // with unknown tags are skipped. bool ParseStatePayload(const void* data, size_t size, StatePayloadView* out); +// A decoded render state block. +struct RenderStateView { + mjvCamera camera; + mjvPerturb perturb; + mjvOption vis_options; + mjOption opt; + mjVisual vis; + mjStatistic stat; + uint8_t render_flags[mjNRNDFLAG]; +}; + +// Decodes a render state block produced by the serializer. `data` must hold +// kRenderStateSize bytes (e.g. StatePayloadView::render_state). +void ParseRenderState(const std::byte* data, RenderStateView* out); + } // namespace mujoco::studio #endif // MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_STATE_PAYLOAD_H_ diff --git a/python/mujoco/experimental/studio/web/web_client.cc b/python/mujoco/experimental/studio/web/web_client.cc new file mode 100644 index 00000000..aa909d03 --- /dev/null +++ b/python/mujoco/experimental/studio/web/web_client.cc @@ -0,0 +1,745 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// This file implements the Web Viewer Browser Client. NetImgui is used to +// transmit remote ImGui draw data in a browser and the filament renderer is +// used to render the UI and the 3D scene. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "third_party/SDL2/include/SDL.h" +#include "third_party/SDL2/include/SDL_opengl.h" +#include +#include +#include +#include "experimental/platform/hal/renderer.h" +#include "experimental/platform/hal/window.h" +#include "experimental/platform/sim/model_holder.h" +#include "experimental/platform/ux/interaction.h" +#include +#include "google/logging.h" +#include "state_payload.h" +#include "web_client_local_ui.h" +#include "web_client_remote_ui.h" +#include "web_client_session.h" + +#if !defined(__EMSCRIPTEN__) +#error "web_client.cc is only supported for Emscripten builds" +#endif + +// No-op wrapper for glGetError(), installed via -Wl,--wrap=glGetError. +// Filament's GL backend calls glGetError() hundreds of times per frame +// through GLUtils::checkGLError/assertGLError. In WebGL each call forces a +// synchronous GPU pipeline flush across the JS-WASM bridge (~0.2ms each). +// This stub eliminates that cost entirely. +extern "C" GLenum __wrap_glGetError(void) { return GL_NO_ERROR; } + +using mujoco::studio::DisconnectNotice; +using mujoco::studio::RemoteUi; +using mujoco::studio::RemoteUiState; +using mujoco::studio::RoleWindow; +using mujoco::studio::Session; +using mujoco::studio::SessionRole; +using mujoco::studio::SessionView; +using mujoco::studio::StatePayloadView; + +// How a spectating page drives its camera. The free modes control the local +// camera directly; kSpecCamFollow mirrors the controller's camera from the +// state broadcast. +enum SpectatorCamMode { + kSpecCamTumble = 0, // Orbit around the lookat point with the mouse. + kSpecCamWasd, // Fly camera: WASD/QE moves, mouse drag turns. + kSpecCamFollow, // Follow the controller's camera. +}; + +// Byte-rate telemetry shown in the role window. +struct Telemetry { + double last_rate_time = 0; + uint64_t gui_bytes_per_sec = 0; + uint64_t sim_bytes_per_sec = 0; +}; + +// The implementation of every interface needed by the session and remote UI. +class AppCallbacks final : public RemoteUi::Callbacks, + public Session::Callbacks { + public: + // RemoteUi::Callbacks + uintptr_t UploadTexture(uintptr_t current, const std::byte* rgba, + uint32_t width, uint32_t height) override; + bool GpuReady() override; + + // Session::Callbacks + bool ReadyForPayload() override; + void OnPayload(const StatePayloadView& view) override; + void ConnectRemoteUi() override; + void ShutdownRemoteUi() override; + void SetCameraMode(int mode) override; +}; + +struct App { + std::unique_ptr window; + std::unique_ptr model_holder; + mujoco::platform::Renderer* renderer = nullptr; + mjvPerturb perturb; + mjvCamera camera; + mjvOption vis_options; + + // Spectator camera + int spectator_cam_mode = kSpecCamTumble; + float spectator_cam_speed = 0.001f; // WASD speed; accelerates while held. + + Telemetry telemetry; + + // Main loop frame counters (MainLoopImpl): reconnect pacing for the state + // WebSocket. (The session paces its own /ui claims by time.) + int frame_count = 0; + int last_state_retry_frame = 0; + + // Backend state received from the Python simulation via WebSocket. + std::vector backend_state; + int backend_state_sig = 0; + bool backend_state_dirty = false; + + // User-injected geoms received with the state payload. + std::vector extra_geoms; + + AppCallbacks callbacks; + RemoteUi remote_ui{callbacks}; + Session session{callbacks}; + DisconnectNotice disconnect_notice; + RoleWindow role_window; +}; +App g_app; + +// WebSocket base URL matching the page origin, e.g. "ws://host:8080" or +// "wss://tunnel.example.com" behind an HTTPS tunnel. All viewer WebSockets +// are served as paths (/ui, /state) on the same host and port as the page +// itself, so exposing or tunneling that single port exposes the whole +// viewer. +std::string GetWsBaseUrl() { + char* base = emscripten_run_script_string( + "(window.location.protocol === 'https:' ? 'wss://' : 'ws://') + " + "window.location.host"); + std::string url = base != nullptr ? base : ""; + if (url == "ws://" || url == "wss://" || url.empty()) { + // No host in the page URL (e.g. file://) — assume a local server. + return "ws://localhost:8080"; + } + return url; +} + +// Stable per-tab id sent with every WebSocket connect (?sid=...), letting the +// server tie this page's /ui and /state connections together. It lives in +// sessionStorage so it survives page reloads: on a model change every page +// reloads and the server restarts with the controller slot reserved for the +// previous controller's sid, a reload-stable id is what lets the controller +// keep control. +std::string GetSessionId() { + static const std::string sid = emscripten_run_script_string( + "(function() {" + " Module.session_id = Module.session_id ||" + " sessionStorage.getItem('mjwv_sid') ||" + " (Date.now().toString(36) + Math.random().toString(36).slice(2));" + " sessionStorage.setItem('mjwv_sid', Module.session_id);" + " return Module.session_id;" + "})()"); + return sid; +} + +std::string WsUrl(const char* path) { + return GetWsBaseUrl() + path + "?sid=" + GetSessionId(); +} + +// Returns true if the Filament rendering context is initialized and ready for +// GPU texture uploads. The Renderer object is created in main() and is always +// non-null, but the Filament context is only initialized when Renderer::Init() +// is called from SetupScene after the async model fetch completes. +bool IsFilamentReady() { + return g_app.renderer && g_app.model_holder && g_app.model_holder->ok(); +} + +// Applies a parsed state payload to the app. Called via AppCallbacks::OnPayload +// after Session validates model CRC and payload readiness. +void ApplyStatePayload(const StatePayloadView& view) { + mjModel* model = g_app.model_holder->model(); + + // Physics state. Guard against a size mismatch (e.g. a stale packet from + // before a model change). + if (view.physics != nullptr) { + const size_t expected_bytes = + mj_stateSize(model, view.physics_spec) * sizeof(mjtNum); + if (view.physics_bytes == expected_bytes) { + g_app.backend_state.resize(view.physics_bytes / sizeof(mjtNum)); + memcpy(g_app.backend_state.data(), view.physics, view.physics_bytes); + g_app.backend_state_sig = view.physics_spec; + g_app.backend_state_dirty = true; + } else { + LOG(Warning, "Physics state size mismatch (%zu != %zu); dropping", + view.physics_bytes, expected_bytes); + } + } + + // Render state. The headless viewer owns the camera and perturbation state + // (all input is forwarded to it and handled by the same code as the native + // viewer); the browser just renders them. + if (view.render_state != nullptr) { + mujoco::studio::RenderStateView rs; + mujoco::studio::ParseRenderState(view.render_state, &rs); + + // Spectators in a free camera mode keep their own local camera; everyone + // else (the controller, and spectators in Follow Controller) mirrors the + // controller's camera. + if (g_app.session.Role() != SessionRole::kSpectating || + g_app.spectator_cam_mode == kSpecCamFollow) { + g_app.camera = rs.camera; + } + + g_app.perturb = rs.perturb; + g_app.vis_options = rs.vis_options; + model->opt = rs.opt; + model->vis = rs.vis; + model->stat = rs.stat; + + // Apply render flags to the renderer's scene if available. + if (g_app.renderer) { + memcpy(g_app.renderer->GetRenderFlags(), rs.render_flags, mjNRNDFLAG); + } + } + + // Extra geoms. memcpy since the payload data is not guaranteed to be aligned. + g_app.extra_geoms.resize(view.extra_geom_count); + if (view.extra_geom_count > 0) { + memcpy(g_app.extra_geoms.data(), view.extra_geoms, + view.extra_geom_count * sizeof(mjvGeom)); + } +} + +void SetSpectatorCameraMode(int mode) { + if (mode == g_app.spectator_cam_mode) { + return; + } + g_app.spectator_cam_mode = mode; + if (!g_app.model_holder || !g_app.model_holder->ok()) { + return; + } + const mjModel* model = g_app.model_holder->model(); + if (mode == kSpecCamTumble) { + mujoco::platform::SetCamera(model, &g_app.camera, + mujoco::platform::kTumbleCameraIdx); + } else if (mode == kSpecCamWasd) { + mujoco::platform::SetCamera(model, &g_app.camera, + mujoco::platform::kFreeCameraIdx); + } + // kSpecCamFollow: the next state payload restores the controller's camera. +} + +// Local camera control for a spectating page in a free camera mode. The +// controller's input goes to the headless viewer instead (CaptureAndSendInput), +// which streams its camera back over the state WebSocket. +// TODO(matijak): Share the camera handling code with the studio app (e.g. in +// platform/ux/interaction.cc) instead of duplicating it here. +void HandleSpectatorCameraInput() { + if (g_app.spectator_cam_mode == kSpecCamFollow) { + return; + } + if (!g_app.model_holder || !g_app.model_holder->ok()) { + return; + } + const mjModel* model = g_app.model_holder->model(); + ImGuiIO& io = ImGui::GetIO(); + const bool wasd = g_app.spectator_cam_mode == kSpecCamWasd; + + // The camera can be fixed on entry (SetupScene honours the model's + // vis.global.cameraid, and Follow Controller mirrors whatever camera the + // controller uses), and mjv_moveCamera ignores fixed cameras. Coerce to a + // free camera so the free modes always respond to input. + if (g_app.camera.type == mjCAMERA_FIXED) { + mjv_defaultFreeCamera(model, &g_app.camera); + if (wasd) { + mujoco::platform::SetCamera(model, &g_app.camera, + mujoco::platform::kFreeCameraIdx); + } + } + + if (!io.WantCaptureMouse && io.DisplaySize.x > 0 && io.DisplaySize.y > 0) { + const float mouse_dx = io.MouseDelta.x / io.DisplaySize.x; + const float mouse_dy = io.MouseDelta.y / io.DisplaySize.y; + const bool is_mouse_dragging = + (mouse_dx != 0.0f || mouse_dy != 0.0f) && + (ImGui::IsMouseDown(ImGuiMouseButton_Left) || + ImGui::IsMouseDown(ImGuiMouseButton_Right) || + ImGui::IsMouseDown(ImGuiMouseButton_Middle)); + if (is_mouse_dragging) { + if (ImGui::IsMouseDown(ImGuiMouseButton_Left)) { + if (wasd) { + mjv_moveCamera(model, mjMOUSE_TURN_H, mouse_dx, 0.f, &g_app.camera); + mjv_moveCamera(model, mjMOUSE_TURN_V, 0.f, mouse_dy, &g_app.camera); + } else { + mjv_moveCamera(model, mjMOUSE_ROTATE_H, mouse_dx, 0.f, &g_app.camera); + mjv_moveCamera(model, mjMOUSE_ROTATE_V, 0.f, mouse_dy, &g_app.camera); + } + } else if (ImGui::IsMouseDown(ImGuiMouseButton_Middle) && !wasd) { + mjv_moveCamera(model, mjMOUSE_ZOOM, 0.f, mouse_dy, &g_app.camera); + } + if (ImGui::IsMouseDown(ImGuiMouseButton_Right)) { + mjv_moveCamera(model, io.KeyShift ? mjMOUSE_MOVE_H : mjMOUSE_MOVE_V, + mouse_dx, mouse_dy, &g_app.camera); + } + } + // Mouse scroll zooms towards/away from the lookat point; ignored by the + // user-centered WASD camera which has no lookat point. + const float mouse_scroll = io.MouseWheel / 50.0f; + if (mouse_scroll != 0.0f && !wasd) { + mjv_moveCamera(model, mjMOUSE_ZOOM, 0.f, -mouse_scroll, &g_app.camera); + } + } + + // WASD/QE flying, with the same accelerating speed as the studio app. + if (wasd && !io.WantCaptureKeyboard) { + bool moved = false; + const float speed = g_app.spectator_cam_speed; + if (ImGui::IsKeyDown(ImGuiKey_W)) { + mjv_moveCamera(model, mjMOUSE_MOVE_H_REL, 0, speed, &g_app.camera); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_S)) { + mjv_moveCamera(model, mjMOUSE_MOVE_H_REL, 0, -speed, &g_app.camera); + moved = true; + } + if (ImGui::IsKeyDown(ImGuiKey_A)) { + mjv_moveCamera(model, mjMOUSE_MOVE_H_REL, -speed, 0, &g_app.camera); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_D)) { + mjv_moveCamera(model, mjMOUSE_MOVE_H_REL, speed, 0, &g_app.camera); + moved = true; + } + if (ImGui::IsKeyDown(ImGuiKey_Q)) { + mjv_moveCamera(model, mjMOUSE_MOVE_V_REL, 0, speed, &g_app.camera); + moved = true; + } else if (ImGui::IsKeyDown(ImGuiKey_E)) { + mjv_moveCamera(model, mjMOUSE_MOVE_V_REL, 0, -speed, &g_app.camera); + moved = true; + } + if (moved) { + const float max_speed = io.KeyShift ? 0.1f : 0.01f; + g_app.spectator_cam_speed = + std::min(g_app.spectator_cam_speed + 0.001f, max_speed); + } else { + g_app.spectator_cam_speed = 0.001f; + } + } +} + +// --- AppCallbacks method bodies (see the class declaration above App). --- + +uintptr_t AppCallbacks::UploadTexture(uintptr_t current, const std::byte* rgba, + uint32_t width, uint32_t height) { + return g_app.renderer->UploadImage(current, rgba, width, height, + rgba ? 4 : 0); +} + +bool AppCallbacks::GpuReady() { return IsFilamentReady(); } + +bool AppCallbacks::ReadyForPayload() { + return g_app.model_holder && g_app.model_holder->ok(); +} + +void AppCallbacks::OnPayload(const StatePayloadView& view) { + ApplyStatePayload(view); +} + +void AppCallbacks::ConnectRemoteUi() { g_app.remote_ui.Connect(WsUrl("/ui")); } + +void AppCallbacks::ShutdownRemoteUi() { g_app.remote_ui.Shutdown(); } + +void AppCallbacks::SetCameraMode(int mode) { SetSpectatorCameraMode(mode); } + +void BuildBrowserGui() { + // Refresh the byte-rate telemetry the role window shows. + const double now = ImGui::GetTime(); + if (now - g_app.telemetry.last_rate_time >= 1.0) { + g_app.telemetry.gui_bytes_per_sec = g_app.remote_ui.ConsumeByteCount(); + g_app.telemetry.sim_bytes_per_sec = g_app.session.ConsumeByteCount(); + g_app.telemetry.last_rate_time = now; + } + + const double last_msg = g_app.session.LastMessageTime(); + 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()); + + SessionView view; + g_app.session.FillView(&view); + view.gui_bytes_per_sec = g_app.telemetry.gui_bytes_per_sec; + 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; + // The session is the SessionActions implementation: the role window's + // intents land there directly. + g_app.role_window.Draw(view, g_app.session); +} + +//================================================================================================= +// Main loop — called once per frame. +//================================================================================================= +void MainLoopImpl(); + +void MainLoop() { + // An exception escaping the requestAnimationFrame callback kills the main + // loop silently causing the canvas to freeze on the last rendered frame and + // input capture to stop with only an opaque "Uncaught " in the console. + // Here we catch, log, and stop explicitly instead. + try { + MainLoopImpl(); + } catch (const std::exception& e) { + LOG(Error, "FATAL: uncaught exception in MainLoop: %s", e.what()); + emscripten_cancel_main_loop(); + } catch (...) { + LOG(Error, "FATAL: uncaught non-std exception in MainLoop"); + emscripten_cancel_main_loop(); + } +} + +void MainLoopImpl() { + g_app.frame_count++; + + // Session upkeep (the liveness heartbeat). + g_app.session.Update(); + + // Reconnect the state WebSocket if it dropped. Receiving a payload with a + // different model identity then triggers a page reload (Session). Deliberate + // server closes (session full, inactivity) are transient; retry them too, + // 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() && + 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..."); + g_app.session.Connect(WsUrl("/state")); + } + + // Pass the remote UI (/ui) WebSocket state to the session state machine. + // The session uses this to manage attempts to claim the single controller + // slot, handle retry pacing if a claim is rejected or dropped, and drive role + // transitions (Controlling vs Spectating). + RemoteUiState ui_state = RemoteUiState::kNoSocket; + if (g_app.remote_ui.HasSocket()) { + switch (g_app.remote_ui.ConnectionState()) { + case RemoteUi::ReadyState::kOpen: + ui_state = RemoteUiState::kOpen; + break; + case RemoteUi::ReadyState::kClosed: + case RemoteUi::ReadyState::kError: + ui_state = RemoteUiState::kClosedOrError; + break; + default: + ui_state = RemoteUiState::kConnecting; + break; + } + } + g_app.session.HandleRemoteUiState(ui_state, g_app.remote_ui.CloseCode()); + + // Process incoming UI data BEFORE the ImGui frame: textures and draw + // frames must be ready before we start the local ImGui frame and render. + if (g_app.window) { + g_app.remote_ui.SetMaxClip(static_cast(g_app.window->GetWidth()), + static_cast(g_app.window->GetHeight())); + } + g_app.remote_ui.ReceiveAndProcessCommands(g_app.frame_count); + + // Event loop and ImGui NewFrame via window abstraction. + mujoco::platform::Window::Status status = g_app.window->NewFrame(); + if (status == mujoco::platform::Window::kQuitting) { + // NewFrame() started an ImGui frame; end it before bailing out. + ImGui::EndFrame(); + emscripten_cancel_main_loop(); + return; + } + + // For the controller, all scene interaction (camera orbit/zoom, perturbation, + // picking) is handled by the headless viewer: CaptureAndSendInput() + // forwards this frame's input over NetImgui, and the resulting camera/perturb + // state streams back over the state WebSocket (see ApplyStatePayload). + // Spectators have no input channel; in a free camera mode they drive their + // local camera directly. + g_app.remote_ui.CaptureAndSendInput(); + if (g_app.session.Role() == SessionRole::kSpectating) { + HandleSpectatorCameraInput(); + } + + BuildBrowserGui(); + + // Finalize the local ImGui draw data. ImguiBridge::Update() (called inside + // Filament's Render) will call ImGui::Render() again, but that is a no-op + // once the frame has already been rendered. + ImGui::Render(); + + // Inject remote draw lists into the local ImDrawData so that Filament's + // ImguiBridge renders them alongside the local UI. Remote lists are inserted + // first (background) and local lists are re-added after (foreground), so the + // local UI always renders on top of remote content. Note that + // ImGui::GetDrawData() is only valid after ImGui::Render() and until the next + // call to ImGui::NewFrame(). + ImDrawData* remote_draw_data = g_app.remote_ui.RemoteDrawData(); + if (remote_draw_data && remote_draw_data->Valid && + g_app.session.Role() == SessionRole::kControlling) { + ImDrawData* local_draw_data = ImGui::GetDrawData(); + if (local_draw_data) { + // Save local draw lists. + ImVector local_lists; + local_lists.reserve(local_draw_data->CmdListsCount); + for (int i = 0; i < local_draw_data->CmdListsCount; ++i) { + local_lists.push_back(local_draw_data->CmdLists[i]); + } + + // Clear and rebuild: remote first, then local. + local_draw_data->CmdLists.resize(0); + local_draw_data->CmdListsCount = 0; + local_draw_data->TotalVtxCount = 0; + local_draw_data->TotalIdxCount = 0; + + // Remote draw lists (background). + for (int i = 0; i < remote_draw_data->CmdListsCount; ++i) { + local_draw_data->AddDrawList(remote_draw_data->CmdLists[i]); + } + + // Local draw lists (foreground). + for (int i = 0; i < local_lists.Size; ++i) { + local_draw_data->AddDrawList(local_lists[i]); + } + } + } + + // 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) { + mjModel* model = g_app.model_holder->model(); + mjData* data = g_app.model_holder->data(); + mj_setState(model, data, g_app.backend_state.data(), + g_app.backend_state_sig); + mj_forward(model, data); + g_app.backend_state_dirty = false; + } + + int width = + static_cast(g_app.window->GetWidth() * g_app.window->GetScale()); + int height = + static_cast(g_app.window->GetHeight() * g_app.window->GetScale()); + if (width > 0 && height > 0) { + g_app.renderer->Render(g_app.model_holder->model(), + g_app.model_holder->data(), &g_app.perturb, + &g_app.camera, &g_app.vis_options, width, height, + /*pixels=*/{}, std::span(g_app.extra_geoms)); + } + } + + g_app.window->Present(); +} + +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(); + + mjv_defaultPerturb(&g_app.perturb); + mjv_defaultCamera(&g_app.camera); + mjv_defaultOption(&g_app.vis_options); + + const int model_cam = m->vis.global.cameraid; + if (model_cam >= 0 && model_cam < m->ncam) { + mujoco::platform::SetCamera(m, &g_app.camera, model_cam); + } else { + mjv_defaultFreeCamera(m, &g_app.camera); + } +} + +// Standard CRC-32 (ISO-HDLC, poly 0xEDB88320), matching Python's zlib.crc32 +// so the fetched model's checksum can be compared against the crc the Python +// side stamps into each state payload. +uint32_t Crc32(const uint8_t* data, size_t len) { + uint32_t crc = 0xFFFFFFFFu; + for (size_t i = 0; i < len; ++i) { + crc ^= data[i]; + for (int bit = 0; bit < 8; ++bit) { + crc = (crc >> 1) ^ (0xEDB88320u & (~(crc & 1u) + 1u)); + } + } + return crc ^ 0xFFFFFFFFu; +} + +void OnFetchSuccess(emscripten_fetch_t* fetch) { + LOG(Info, "Fetched model.mjb, size: %llu", fetch->numBytes); + g_app.model_holder = mujoco::platform::ModelHolder::FromBuffer( + std::span( + reinterpret_cast(fetch->data), + static_cast(fetch->numBytes)), + "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"); + } + emscripten_fetch_close(fetch); +} + +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 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"); +} + +// Holds assets (Filament assets and local ImGui fonts) that the page fetches +// and pushes in via the registerAsset() binding before startApp() runs, so the +// resource providers below can resolve them without a filesystem. +class AssetRegistry { + public: + static AssetRegistry& Instance() { + static AssetRegistry instance; + return instance; + } + + void RegisterAsset(std::string filename, std::string contents) { + assets_[std::move(filename)] = std::move(contents); + } + + const std::string& Get(std::string_view filename) const { + filename = filename.substr(filename.find_first_of(':') + 1); + static std::string empty; + auto it = assets_.find(std::string(filename)); + return it != assets_.end() ? it->second : empty; + } + + private: + std::unordered_map assets_; +}; + +// Exposed to JS (see EMSCRIPTEN_BINDINGS): the page calls this once per asset. +void RegisterAsset(std::string filename, std::string contents) { + AssetRegistry::Instance().RegisterAsset(std::move(filename), + std::move(contents)); +} + +// Registers resource providers so that "filament:" and "font:" asset requests +// resolve from the AssetRegistry populated by registerAsset(). +static void RegisterAssetProviders() { + mjpResourceProvider resource_provider; + mjp_defaultResourceProvider(&resource_provider); + + resource_provider.open = [](mjResource* resource) { + AssetRegistry& r = AssetRegistry::Instance(); + return static_cast(r.Get(resource->name).size()); + }; + resource_provider.read = [](mjResource* resource, const void** buffer) { + AssetRegistry& r = AssetRegistry::Instance(); + const std::string& contents = r.Get(resource->name); + *buffer = contents.data(); + return static_cast(contents.size()); + }; + resource_provider.close = [](mjResource* resource) {}; + for (const char* prefix : {"filament", "font"}) { + resource_provider.prefix = prefix; + mjp_registerResourceProvider(&resource_provider); + } +} + +// 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. +void StartApp() { + mujoco::platform::Window::Config config; + config.gfx_mode = mujoco::platform::GraphicsMode::FilamentWebGl; + // Load the Studio UI fonts for the browser's local ImGui (the role window); + // the "font:" resource provider above resolves them from the AssetRegistry. + // (The streamed/remote UI carries its own font atlas separately.) + config.load_fonts = true; + + g_app.window = std::make_unique("MuJoCo Web Viewer", + 1400, 720, config); + ImPlot::CreateContext(); // Needed if the server app uses ImPlot. + + g_app.renderer = new mujoco::platform::Renderer( + g_app.window->GetNativeWindowHandle(), config.gfx_mode); + + 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); +} + +int main(int argc, char** argv) { + RegisterAssetProviders(); + return 0; +} + +EMSCRIPTEN_BINDINGS(web_client_bindings) { + emscripten::function("registerAsset", &RegisterAsset); + emscripten::function("startApp", &StartApp); +} diff --git a/python/mujoco/experimental/studio/web/web_client_local_ui.cc b/python/mujoco/experimental/studio/web/web_client_local_ui.cc new file mode 100644 index 00000000..f5a202c5 --- /dev/null +++ b/python/mujoco/experimental/studio/web/web_client_local_ui.cc @@ -0,0 +1,387 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "web_client_local_ui.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include "experimental/platform/ux/imgui_widgets.h" +#include "google/logging.h" + +namespace mujoco::studio { +namespace { + +const ImVec2 kFullWidth(-FLT_MIN, 0.0f); +const ImVec4 kSpectatingColor(1.0f, 0.75f, 0.2f, 1.0f); +const ImVec4 kControllingColor(0.3f, 0.9f, 0.4f, 1.0f); +const ImVec4 kQueueColor(1.0f, 0.62f, 0.15f, 1.0f); +const ImVec4 kConnectingColor(0.6f, 0.6f, 0.6f, 1.0f); + +// Draws one screen-centered DISCONNECTED window. +void DrawDisconnectWindow(const char* window_id, + std::initializer_list lines) { + const ImGuiIO& io = ImGui::GetIO(); + ImGui::SetNextWindowPos( + ImVec2(io.DisplaySize.x * 0.5f, io.DisplaySize.y * 0.5f), + ImGuiCond_Always, ImVec2(0.5f, 0.5f)); + ImGui::Begin(window_id, nullptr, + ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoResize | + ImGuiWindowFlags_NoMove | ImGuiWindowFlags_AlwaysAutoResize); + platform::CenteredBanner("DISCONNECTED", ImVec4(1.0f, 0.3f, 0.3f, 1.0f)); + for (const char* line : lines) { + platform::CenteredLine(line, nullptr); + } + ImGui::End(); +} + +} // namespace + +void DisconnectNotice::Draw(int server_close_code, + double seconds_since_last_payload, + bool reload_pending) { + if (server_close_code != 0) { + const char* reason = "Disconnected by the viewer."; + if (server_close_code == kWsCloseSessionFull) { + reason = "Session is full: too many viewers connected."; + } else if (server_close_code == kWsCloseInactive) { + reason = "Disconnected after inactivity."; + } + DrawDisconnectWindow("##disconnected_by_server", + {reason, "Retrying; reconnects automatically."}); + } + + // The silence notice keys on state-stream staleness rather than socket + // state: a killed, suspended (Ctrl+Z), or unreachable server all go + // 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. + const bool link_stale = + seconds_since_last_payload > kServerSilenceNoticeSec && + server_close_code == 0 && !reload_pending; + if (!link_stale) { + logged_ = false; + return; + } + if (!logged_) { + LOG(Info, "Showing server-not-reachable notice"); + logged_ = true; + } + DrawDisconnectWindow( + "##disconnected", + {"Viewer server is not reachable: the Python script may", + "have stopped. This page reconnects automatically if the", + "viewer comes back."}); +} + +void RoleWindow::Draw(const SessionView& view, SessionActions& actions) { + ComputeAnchorPositions(); + // While anchored, the window is re-pinned every frame (so it follows + // canvas resizes) for both the collapsed pill and the expanded window + // (same ImGui window, so one call covers whichever Begin runs this + // frame) — except while it is being dragged. + if (anchor_ != kAnchorFree && !dragging_) { + ImGui::SetNextWindowPos(anchor_pos_[anchor_ - kAnchorTopLeft], + ImGuiCond_Always, + kAnchorPivot[anchor_ - kAnchorTopLeft]); + } + + // While someone waits for control, the window background pulses toward + // a dark orange (the queue size itself is shown in the expanded form). + const bool pulse = + view.role == SessionRole::kControlling && view.queue_len > 0; + if (pulse) { + const float phase = + 0.5f + 0.5f * sinf(static_cast(ImGui::GetTime()) * 2.5f); + ImVec4 bg = ImGui::GetStyleColorVec4(ImGuiCol_WindowBg); + const ImVec4 orange(0.45f, 0.22f, 0.02f, bg.w); + const float blend = 0.6f * phase; + bg.x += (orange.x - bg.x) * blend; + bg.y += (orange.y - bg.y) * blend; + bg.z += (orange.z - bg.z) * blend; + ImGui::PushStyleColor(ImGuiCol_WindowBg, bg); + } + + if (mode_ == Mode::kCollapsed) { + DrawCollapsed(view); + } else { + DrawExpanded(view, actions); + } + + if (pulse) { + ImGui::PopStyleColor(); + } +} + +void RoleWindow::ComputeAnchorPositions() { + const ImVec2 canvas = ImGui::GetIO().DisplaySize; + anchor_pos_[0] = ImVec2(kEdgeMargin, kEdgeMargin); + anchor_pos_[1] = ImVec2(canvas.x * 0.5f, kEdgeMargin); + anchor_pos_[2] = ImVec2(canvas.x - kEdgeMargin, kEdgeMargin); + anchor_pos_[3] = ImVec2(kEdgeMargin, canvas.y - kEdgeMargin); + anchor_pos_[4] = ImVec2(canvas.x * 0.5f, canvas.y - kEdgeMargin); + anchor_pos_[5] = ImVec2(canvas.x - kEdgeMargin, canvas.y - kEdgeMargin); +} + +// Tracks a drag of the window and snaps on release: if the window was +// dropped within max(20% of its size, 100px) of where an anchor would +// place it, adopt that anchor; otherwise it floats Free where it was +// dropped. Call between Begin and End of whichever form is visible. +void RoleWindow::UpdateSnap() { + if (ImGui::IsMouseClicked(ImGuiMouseButton_Left) && + ImGui::IsWindowHovered(ImGuiHoveredFlags_AllowWhenBlockedByActiveItem)) { + dragging_ = true; + } else if (dragging_ && !ImGui::IsMouseDown(ImGuiMouseButton_Left)) { + dragging_ = false; + const ImVec2 pos = ImGui::GetWindowPos(); + const ImVec2 size = ImGui::GetWindowSize(); + const float threshold = std::max(0.2f * std::max(size.x, size.y), 100.0f); + int best = -1; + float best_d2 = threshold * threshold; + for (int i = 0; i < 6; ++i) { + // Per-axis distance to where this anchor would place the window, + // clamped to zero on the anchor's off-screen side: a window dropped + // past an edge or corner (mostly off-screen) still snaps to that + // anchor instead of being left stranded outside the canvas. + float dx = pos.x - (anchor_pos_[i].x - kAnchorPivot[i].x * size.x); + if (kAnchorPivot[i].x == 0.0f) { + dx = std::max(dx, 0.0f); // Dropped past the left edge. + } else if (kAnchorPivot[i].x == 1.0f) { + dx = std::min(dx, 0.0f); // Dropped past the right edge. + } + float dy = pos.y - (anchor_pos_[i].y - kAnchorPivot[i].y * size.y); + if (kAnchorPivot[i].y == 0.0f) { + dy = std::max(dy, 0.0f); // Dropped past the top edge. + } else { + dy = std::min(dy, 0.0f); // Dropped past the bottom edge. + } + const float d2 = dx * dx + dy * dy; + if (d2 <= best_d2) { + best_d2 = d2; + best = i; + } + } + anchor_ = + best >= 0 ? static_cast(kAnchorTopLeft + best) : kAnchorFree; + } +} + +// Data rates, shown to both roles (spectators receive no GUI stream, so +// that line reads 0 for them). With more than one viewer, the number in +// parentheses is the host's total outgoing sim bandwidth: every connected +// browser receives the same sim stream. +void RoleWindow::DataRateLines(const SessionView& view) { + ImGui::Text("GUI Data Rate: %" PRIu64 " KiB/s", + static_cast(view.gui_bytes_per_sec / 1024)); + if (view.role == SessionRole::kSpectating) { + ImGui::SetItemTooltip("The UI is not streamed to spectators."); + } else if (!view.have_remote_frame) { + ImGui::SameLine(); + ImGui::TextColored(ImVec4(1.0f, 0.0f, 0.0f, 1.0f), "(Waiting...)"); + } + ImGui::Text("Sim Data Rate: %" PRIu64 " KiB/s", + static_cast(view.sim_bytes_per_sec / 1024)); + if (view.viewers > 1) { + ImGui::SameLine(); + ImGui::Text( + "(%" PRIu64 " KiB/s)", + static_cast(view.sim_bytes_per_sec * view.viewers / 1024)); + ImGui::SetItemTooltip("Total sim data sent across all viewers."); + } +} + +void RoleWindow::DrawCollapsed(const SessionView& view) { + ImGui::PushStyleVar(ImGuiStyleVar_WindowPadding, ImVec2(6.0f, 6.0f)); + ImGui::PushStyleVar(ImGuiStyleVar_WindowMinSize, ImVec2(16.0f, 16.0f)); + + ImGui::Begin("Role", nullptr, + ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoResize | + ImGuiWindowFlags_AlwaysAutoResize); + if (view.role == SessionRole::kSpectating) { + ImGui::TextColored(kSpectatingColor, "SPECTATING"); + } else if (view.role == SessionRole::kControlling) { + ImGui::TextColored(kControllingColor, "CONTROLLING"); + } else { + 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 + // a move (which drops drag capture in ImGui). + if (ImGui::IsWindowFocused() || + ImGui::IsWindowHovered(ImGuiHoveredFlags_AllowWhenBlockedByActiveItem)) { + mode_ = Mode::kExpanded; + } + UpdateSnap(); + ImGui::End(); + + ImGui::PopStyleVar(2); +} + +void RoleWindow::DrawExpanded(const SessionView& view, + SessionActions& actions) { + // No title bar; the window is still movable by dragging empty space + // (ImGui's default when ConfigWindowsMoveFromTitleBarOnly is off). + ImGui::Begin("Role", nullptr, + ImGuiWindowFlags_AlwaysAutoResize | ImGuiWindowFlags_NoTitleBar); + + 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); + ImGui::Separator(); + DataRateLines(view); + } else if (view.role == SessionRole::kSpectating) { + DrawSpectatorContents(view, actions); + } else { + DrawControllerContents(view, actions); + } + + UpdateSnap(); + UpdateCollapse(); + ImGui::End(); +} + +void RoleWindow::DrawSpectatorContents(const SessionView& view, + SessionActions& actions) { + // Control group. + platform::CenteredBanner("SPECTATING", kSpectatingColor); + char queue_text[64]; + if (view.queue_pos > 0) { + snprintf(queue_text, sizeof(queue_text), "Control Queue: Position %d of %d", + view.queue_pos, view.queue_len); + } else if (view.queue_len > 0) { + snprintf(queue_text, sizeof(queue_text), "Control Queue: %d waiting", + view.queue_len); + } else { + snprintf(queue_text, sizeof(queue_text), "Control Queue: (empty)"); + } + platform::CenteredLine(queue_text, nullptr); + if (view.queue_pos == 0) { + if (ImGui::Button("Request control", kFullWidth)) { + actions.RequestControl(); + } + } else { + const float half_width = + (ImGui::GetContentRegionAvail().x - ImGui::GetStyle().ItemSpacing.x) * + 0.5f; + if (ImGui::Button("Leave Queue", ImVec2(half_width, 0.0f))) { + actions.LeaveQueue(); + } + ImGui::SetItemTooltip("Abandons the control request."); + ImGui::SameLine(); + ImGui::PushStyleColor(ImGuiCol_Button, ImVec4(0.65f, 0.15f, 0.15f, 1.0f)); + ImGui::PushStyleColor(ImGuiCol_ButtonHovered, + ImVec4(0.80f, 0.20f, 0.20f, 1.0f)); + ImGui::PushStyleColor(ImGuiCol_ButtonActive, + ImVec4(0.50f, 0.10f, 0.10f, 1.0f)); + if (ImGui::Button("Steal Control", ImVec2(half_width, 0.0f))) { + actions.StealControl(); + } + ImGui::PopStyleColor(3); + ImGui::SetItemTooltip("Takes control immediately, jumping the queue."); + } + + // Data rate group. + ImGui::Separator(); + DataRateLines(view); + + // Camera group. + ImGui::Separator(); + // TODO(matijak): Use the studio camera-selection UI here in future, so + // a spectator can also pick any camera defined in the model. + // Sized to its longest option; the default width would stretch the + // window past the CONTROLLING layout's size. + ImGui::SetNextItemWidth(ImGui::CalcTextSize("Follow Controller").x + + ImGui::GetFrameHeight() * 2.0f); + int cam_mode = view.camera_mode; + if (ImGui::Combo("Camera", &cam_mode, + "Free: tumble\0Free: wasd\0Follow Controller\0")) { + actions.SetCameraMode(cam_mode); + } +} + +void RoleWindow::DrawControllerContents(const SessionView& view, + SessionActions& actions) { + // Control group. + platform::CenteredBanner("CONTROLLING", kControllingColor); + if (view.queue_len > 0) { + char queue_text[64]; + snprintf(queue_text, sizeof(queue_text), ">> Control Queue: %d waiting <<", + view.queue_len); + platform::CenteredLine(queue_text, &kQueueColor); + } else { + platform::CenteredLine("Control Queue: (empty)", nullptr); + } + if (ImGui::Button("Release control", kFullWidth)) { + actions.ReleaseControl(); + } + + // Data rate group. + ImGui::Separator(); + DataRateLines(view); + + // Session settings group. + ImGui::Separator(); + // Local edits win for a grace period (the roster round trip takes a + // few frames and typing takes longer); afterwards the server's value + // is the truth. + if (max_spectators_edit_ < 0 || + ImGui::GetTime() - max_spectators_edit_time_ > 1.5) { + max_spectators_edit_ = view.max_spectators; + } + ImGui::SetNextItemWidth(100.0f); + if (ImGui::InputInt("Max Spectators", &max_spectators_edit_)) { + max_spectators_edit_time_ = ImGui::GetTime(); + max_spectators_edit_ = std::clamp(max_spectators_edit_, 0, 32); + actions.SetMaxSpectators(max_spectators_edit_); + } + if (ImGui::IsItemActive()) { + max_spectators_edit_time_ = ImGui::GetTime(); + } +} + +// Collapse when the mouse leaves the window, with a short grace period. +// A popup (e.g. the camera combo's dropdown) is a separate window, so +// moving the mouse into it unhovers this one; keep the window expanded +// while one of its popups is open, and give brief excursions time to +// come back before snapping shut. On page load the window instead stays +// expanded for a few seconds (the intro); the first hover ends the +// intro, so a quick swipe over the window dismisses it early. +void RoleWindow::UpdateCollapse() { + const bool popup_open = ImGui::IsPopupOpen( + "", ImGuiPopupFlags_AnyPopupId | ImGuiPopupFlags_AnyPopupLevel); + if (popup_open || dragging_ || + ImGui::IsWindowFocused(ImGuiFocusedFlags_RootAndChildWindows) || + ImGui::IsWindowHovered(ImGuiHoveredFlags_AllowWhenBlockedByActiveItem)) { + mode_ = Mode::kExpanded; // The first hover ends the intro. + hover_time_ = ImGui::GetTime(); + } else if (mode_ == Mode::kIntro) { + if (intro_start_ == 0) { + intro_start_ = ImGui::GetTime(); + } else if (ImGui::GetTime() - intro_start_ > kIntroExpandedSec) { + mode_ = Mode::kCollapsed; + } + } else if (ImGui::GetTime() - hover_time_ > kCollapseGraceSec) { + mode_ = Mode::kCollapsed; + } +} + +} // namespace mujoco::studio diff --git a/python/mujoco/experimental/studio/web/web_client_local_ui.h b/python/mujoco/experimental/studio/web/web_client_local_ui.h new file mode 100644 index 00000000..f8aa7a4c --- /dev/null +++ b/python/mujoco/experimental/studio/web/web_client_local_ui.h @@ -0,0 +1,113 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// The web viewer's local (browser-drawn) UI: the role window and the +// DISCONNECTED notices. Deliberately decoupled from the app: it reads a +// SessionView snapshot and reports user intent through the SessionActions +// interface (both defined in web_client_session.h), so drawing depends +// only on ImGui. + +#ifndef MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_LOCAL_UI_H_ +#define MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_LOCAL_UI_H_ + +#include +#include "web_client_session.h" + +namespace mujoco::studio { + +// The screen-centered DISCONNECTED notices: a big red banner over white +// explanation lines, shown for deliberate server closes (session full, +// inactivity) and for server silence. The two never show at the same time: +// a close code suppresses the silence notice. +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); + + private: + // How long the state stream (~60Hz while the Python side is alive) may go + // silent before the notice appears. A model-change restart resumes + // traffic in about a second and should not flash the notice. + static constexpr double kServerSilenceNoticeSec = 3.0; + + // True while the server-not-reachable notice is up (log once per outage). + bool logged_ = false; +}; + +// The role window: a collapsed role pill that expands on hover into the +// session panel (role banner, control queue, data rates, per-role +// settings). Owns its presentation state — mode, intro/hover timers, drag +// and snap anchoring, edit grace. +class RoleWindow { + public: + void Draw(const SessionView& view, SessionActions& actions); + + private: + // Presentation state machine: the window starts in the intro (expanded + // for a few seconds so new users see it exists and learn it expands on + // hover); the first hover ends the intro. + enum class Mode { + kIntro = 0, // Expanded on page load, until the timer or first hover. + kExpanded, // Expanded while hovered (plus a short leave grace). + kCollapsed, // The role pill; expands on hover. + }; + + // Where the window is anchored on the page. It is always draggable; + // releasing a drag near an anchor snaps to it, anywhere else leaves the + // window Free at its dropped position. + enum Anchor { + kAnchorFree = 0, + kAnchorTopLeft, // The anchors must stay contiguous from here on. + kAnchorTopMid, + kAnchorTopRight, + kAnchorBottomLeft, + kAnchorBottomMid, + kAnchorBottomRight, + }; + + static constexpr double kCollapseGraceSec = 0.2; + static constexpr double kIntroExpandedSec = 5.0; + static constexpr float kEdgeMargin = 10.0f; + // Pivots keeping the whole window on-screen kEdgeMargin off the + // anchoring edges; indexed by Anchor - kAnchorTopLeft, like anchor_pos_. + static constexpr ImVec2 kAnchorPivot[6] = { + {0.0f, 0.0f}, {0.5f, 0.0f}, {1.0f, 0.0f}, + {0.0f, 1.0f}, {0.5f, 1.0f}, {1.0f, 1.0f}, + }; + + void ComputeAnchorPositions(); + void UpdateSnap(); + void UpdateCollapse(); + static void DataRateLines(const SessionView& view); + void DrawCollapsed(const SessionView& view); + void DrawExpanded(const SessionView& view, SessionActions& actions); + void DrawSpectatorContents(const SessionView& view, SessionActions& actions); + void DrawControllerContents(const SessionView& view, SessionActions& actions); + + Mode mode_ = Mode::kIntro; + Anchor anchor_ = kAnchorTopMid; + ImVec2 anchor_pos_[6]; // Recomputed each frame (canvas resizes). + double intro_start_ = 0; // 0 until the expanded window first draws. + double hover_time_ = 0; // Feeds the collapse grace period. + bool dragging_ = false; + // Max Spectators edit grace: the local value wins over the roster until + // max_spectators_edit_time_ is old enough (see the InputInt). + int max_spectators_edit_ = -1; + double max_spectators_edit_time_ = -1e9; +}; + +} // namespace mujoco::studio + +#endif // MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_LOCAL_UI_H_ diff --git a/python/mujoco/experimental/studio/web/web_client_remote_ui.cc b/python/mujoco/experimental/studio/web/web_client_remote_ui.cc new file mode 100644 index 00000000..c7049802 --- /dev/null +++ b/python/mujoco/experimental/studio/web/web_client_remote_ui.cc @@ -0,0 +1,771 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "web_client_remote_ui.h" + +#include +#include +#include + +#include "google/logging.h" + +namespace mujoco::studio { + +using namespace NetImgui::Internal; + +namespace { + +// Delta-compress the GUI stream (relayed to the client via CmdInput). +constexpr bool kUseCompression = true; + +void LogUnmappedTexture( + RemoteUi::ClientTextureID client_tex_id, size_t map_size, uint32_t draw_idx, + const std::unordered_map& tex_map) { + VLOG(1, "DrawFrame: UNMAPPED client_tex_id=%lu, map_size=%zu, draw#=%u", + static_cast(client_tex_id), map_size, draw_idx); + std::string keys_str = ""; + for (auto& kv : tex_map) { + keys_str += " " + std::to_string(kv.first) + + "(fil=" + std::to_string(kv.second) + ")"; + } + VLOG(1, " map keys:%s", keys_str.c_str()); +} + +void LogCmdReceived(CmdHeader::eCommands cmd_type, uint32_t cmd_size, + int draw_frames, int textures, + const PendingCom& pending_receive) { + if (cmd_type == CmdHeader::eCommands::Version) { + const CmdVersion* ver = + reinterpret_cast(pending_receive.pCommand); + LOG(Info, + "Received CmdVersion from client: name='%s', version=%d, wchar_size=%d", + ver->mClientName, static_cast(ver->mVersion), ver->mWCharSize); + } else if (cmd_type == CmdHeader::eCommands::DrawFrame) { + VLOG(2, "Received DrawFrame #%d (size=%u)", draw_frames, cmd_size); + } else if (cmd_type == CmdHeader::eCommands::Texture) { + VLOG(2, "Received Texture #%d (size=%u)", textures, cmd_size); + } else if (cmd_type == CmdHeader::eCommands::Background) { + VLOG(2, "Received Background cmd (size=%u)", cmd_size); + } else if (cmd_type != CmdHeader::eCommands::Count && + cmd_type != CmdHeader::eCommands::Clipboard && + cmd_type != CmdHeader::eCommands::Input) { + VLOG(2, "Received UNKNOWN cmd: type=%d, size=%u", + static_cast(cmd_type), cmd_size); + } +} + +} // namespace + +void RemoteUi::Connect(const std::string& url) { + if (socket_) { + Network::Disconnect(socket_); + socket_ = nullptr; + } + // The port argument is unused when a full URL is passed. + LOG(Info, "Connecting to WebSocket at %s", url.c_str()); + socket_ = Network::Connect(url.c_str(), 0); + LOG(Info, "Network::Connect returned: socket=%p", + static_cast(socket_)); +} + +RemoteUi::ReadyState RemoteUi::ConnectionState() const { + return Network::GetReadyState(socket_); +} + +int RemoteUi::CloseCode() const { return Network::GetCloseCode(socket_); } + +void RemoteUi::ReceiveAndProcessCommands(int frame) { + if (!socket_) return; + + const ReadyState state = ConnectionState(); + if (last_state_ != state) { + LOG(Info, "WebSocket status changed: '%s' -> '%s' (frame %d)", + Network::ReadyStateName(last_state_), Network::ReadyStateName(state), + frame); + last_state_ = state; + } + VLOG(1, + "Frame %d: status='%s', handshake=%s, cmds=%d, draws=%d, textures=%d, " + "has_draw_data=%s", + frame, Network::ReadyStateName(state), + handshake_sent_ ? "sent" : "not_sent", total_cmds_received_, + draw_frames_received_, textures_received_, + remote_draw_data_ != nullptr ? "yes" : "no"); + + if (state == ReadyState::kOpen) { + if (!handshake_sent_) { + CmdVersion cmd_version; + StringCopy(cmd_version.mClientName, "MuJoCo Web Viewer"); + LOG(Info, + "Sending CmdVersion handshake: size=%u, type=%d, version=%d, " + "wchar_size=%d, name='%s'", + cmd_version.mSize, static_cast(cmd_version.mType), + static_cast(cmd_version.mVersion), cmd_version.mWCharSize, + cmd_version.mClientName); + + PendingCom pending_send; + pending_send.pCommand = &cmd_version; + pending_send.SizeCurrent = 0; + + int send_attempts = 0; + while (!pending_send.IsDone() && !pending_send.IsError()) { + const size_t before = pending_send.SizeCurrent; + Network::DataSend(socket_, pending_send); + send_attempts++; + if (pending_send.SizeCurrent == before) { + // No progress: the socket already reads OPEN (from emscripten's + // ready state) but the send backend is not ready yet, because the + // open callback that sets mConnected has not run. DataSend then + // returns with neither progress nor error, so this loop would spin + // the browser's main thread forever. Stop and retry on a later + // frame instead (handshake_sent_ stays false below). + break; + } + } + LOG(Info, + "CmdVersion send: done=%s, error=%s, attempts=%d, bytes_sent=%zu", + pending_send.IsDone() ? "true" : "false", + pending_send.IsError() ? "true" : "false", send_attempts, + static_cast(pending_send.SizeCurrent)); + if (pending_send.IsDone()) { + handshake_sent_ = true; + // Fresh connection: the client may resume delta compression against + // a frame from a previous session; ask for an uncompressed keyframe. + request_keyframe_ = true; + } + } + } else { + if (handshake_sent_) { + LOG(Info, "Resetting handshake (status='%s')", + Network::ReadyStateName(state)); + } + handshake_sent_ = false; + } + + // Network receive — drain all available data in one frame. + const bool is_connected = (state == ReadyState::kOpen); + + // If the connection has closed, discard any buffered data immediately + // rather than churning through stale commands for several seconds. + if (was_connected_ && !is_connected) { + LOG(Info, "Connection lost (status='%s'). Discarding buffered data.", + Network::ReadyStateName(state)); + // Reset any in-progress receive. + if (pending_receive_.bAutoFree) + netImguiDeleteSafe(pending_receive_.pCommand); + pending_receive_ = PendingCom(); + was_connected_ = false; + } + if (is_connected) was_connected_ = true; + + int max_commands_per_frame = 64; + int cmds_this_frame = 0; + bool had_pending_data = Network::DataReceivePending(socket_); + if (had_pending_data) { + VLOG(1, "Frame %d: data pending on socket", frame); + } + + while (is_connected && max_commands_per_frame-- > 0) { + if (pending_receive_.IsReady()) { + cmd_pending_read_ = CmdPendingRead(); + pending_receive_.pCommand = &cmd_pending_read_; + pending_receive_.bAutoFree = false; + } + + if (!Network::DataReceivePending(socket_)) break; + + Network::DataReceive(socket_, pending_receive_); + + if (pending_receive_.pCommand->mSize > sizeof(CmdPendingRead) && + pending_receive_.pCommand == &cmd_pending_read_) { + VLOG(2, "Allocating %u bytes for incoming cmd type=%d", + pending_receive_.pCommand->mSize, + static_cast(pending_receive_.pCommand->mType)); + CmdPendingRead* cmd_header = reinterpret_cast( + netImguiSizedNew(pending_receive_.pCommand->mSize)); + *cmd_header = cmd_pending_read_; + pending_receive_.pCommand = cmd_header; + pending_receive_.bAutoFree = true; + } + + if (!pending_receive_.IsDone()) { + if (pending_receive_.IsError()) { + LOG(Error, "Receive ERROR: cmd_size=%u, got=%zu, type=%d", + pending_receive_.pCommand->mSize, + static_cast(pending_receive_.SizeCurrent), + static_cast(pending_receive_.pCommand->mType)); + if (pending_receive_.bAutoFree) + netImguiDeleteSafe(pending_receive_.pCommand); + pending_receive_ = PendingCom(); + } + continue; + } + + // Command fully received — dispatch. + bytes_accum_ += pending_receive_.pCommand->mSize; + cmds_this_frame++; + total_cmds_received_++; + CmdHeader::eCommands cmd_type = pending_receive_.pCommand->mType; + LogCmdReceived(cmd_type, pending_receive_.pCommand->mSize, + draw_frames_received_, textures_received_, pending_receive_); + + if (cmd_type == CmdHeader::eCommands::Count) { + // CmdPendingRead sentinel — skip silently. + } else if (cmd_type == CmdHeader::eCommands::DrawFrame) { + draw_frames_received_++; + ProcessCmdDrawFrame( + reinterpret_cast(pending_receive_.pCommand)); + } else if (cmd_type == CmdHeader::eCommands::Texture) { + textures_received_++; + ProcessCmdTexture( + reinterpret_cast(pending_receive_.pCommand)); + } + + if (pending_receive_.bAutoFree) + netImguiDeleteSafe(pending_receive_.pCommand); + pending_receive_ = PendingCom(); + } + + VLOG(2, "Frame %d: processed %d commands", frame, cmds_this_frame); +} + +void RemoteUi::FlushPendingTextures() { + 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", + static_cast(tex_id), + static_cast(local_tex)); + } + } +} + +void RemoteUi::ProcessCmdTexture(CmdTexture* cmd_texture) { + if (!cmd_texture) return; + + if (!cmd_texture->mpTextureData.IsPointer()) { + cmd_texture->mpTextureData.ToPointer(); + } + + ClientTextureID tex_id = cmd_texture->mTextureClientID; + VLOG( + 1, + "ProcessCmdTexture: client_tex_id=%lu, status=%d, size=%ux%u, format=%d, " + "offset=%u,%u, map_size=%zu", + static_cast(tex_id), + static_cast(cmd_texture->mStatus), + static_cast(cmd_texture->mWidth), + static_cast(cmd_texture->mHeight), + static_cast(cmd_texture->mFormat), + static_cast(cmd_texture->mOffsetX), + static_cast(cmd_texture->mOffsetY), texture_map_.size()); + + uintptr_t& local_tex = texture_map_[tex_id]; + + if (cmd_texture->mStatus == CmdTexture::eType::Destroy) { + if (local_tex != 0 && callbacks_.GpuReady()) { + // Pass nullptr pixels to destroy the GPU texture. + callbacks_.UploadTexture(local_tex, nullptr, 0, 0); + } + local_tex = 0; + texture_cpu_.erase(tex_id); + return; + } + + uint8_t* pixels = cmd_texture->mpTextureData.Get(); + if (!pixels) return; + + // All of width/height/size come off the wire, so compute in size_t (no + // 32-bit overflow) and verify the command actually carries the pixel + // bytes before reading them — a short or corrupt command must not make + // the memcpy/expand below read past the command buffer. + const uint32_t patch_w = cmd_texture->mWidth; + const uint32_t patch_h = cmd_texture->mHeight; + const size_t pixel_count = static_cast(patch_w) * patch_h; + const size_t data_size = cmd_texture->mSize >= sizeof(CmdTexture) + ? cmd_texture->mSize - sizeof(CmdTexture) + : 0; + const size_t expected_rgba = pixel_count * 4; + + // Detect actual format by data size, not format tag (which can be wrong). + const bool is_a8 = (cmd_texture->mFormat == 1) || + (data_size == pixel_count && data_size != expected_rgba); + const size_t needed = is_a8 ? pixel_count : expected_rgba; + if (pixel_count == 0 || data_size < needed) { + LOG(Warning, + "ProcessCmdTexture: tex_id=%lu declares %ux%u but carries only %zu " + "bytes; dropping", + static_cast(tex_id), patch_w, patch_h, data_size); + return; + } + + // Convert incoming pixels to RGBA (Filament requires RGBA). + // For A8 font atlas data, expand each byte to (255, 255, 255, alpha). + std::vector rgba_pixels(pixel_count * 4); + if (is_a8) { + for (size_t p = 0; p < pixel_count; ++p) { + rgba_pixels[p * 4 + 0] = 255; + rgba_pixels[p * 4 + 1] = 255; + rgba_pixels[p * 4 + 2] = 255; + rgba_pixels[p * 4 + 3] = pixels[p]; + } + } else { + memcpy(rgba_pixels.data(), pixels, expected_rgba); + } + + TextureEntry& entry = texture_cpu_[tex_id]; + + if (cmd_texture->mStatus == CmdTexture::eType::Create) { + // Full texture creation — store the CPU-side mirror. + entry.width = patch_w; + entry.height = patch_h; + entry.pixels = std::move(rgba_pixels); + } else { + // Partial update — patch the sub-region into the existing CPU mirror. + // If no CPU mirror exists (e.g. we missed the Create), skip. + if (entry.pixels.empty()) { + LOG(Warning, + "ProcessCmdTexture: partial update for tex_id=%lu " + "but no CPU mirror exists, skipping", + static_cast(tex_id)); + return; + } + const uint32_t off_x = cmd_texture->mOffsetX; + const uint32_t off_y = cmd_texture->mOffsetY; + // The patch rectangle is wire-supplied; reject one that would write past + // the stored mirror (e.g. dims desynced from a stale entry after a + // reconnect) rather than corrupting the heap. + if (static_cast(off_x) + patch_w > entry.width || + static_cast(off_y) + patch_h > entry.height) { + LOG(Warning, + "ProcessCmdTexture: patch %ux%u at (%u,%u) exceeds mirror %ux%u " + "for tex_id=%lu; dropping", + patch_w, patch_h, off_x, off_y, entry.width, entry.height, + static_cast(tex_id)); + return; + } + for (uint32_t row = 0; row < patch_h; ++row) { + size_t dst_offset = + (static_cast(off_y + row) * entry.width + off_x) * 4; + size_t src_offset = static_cast(row) * patch_w * 4; + memcpy(&entry.pixels[dst_offset], &rgba_pixels[src_offset], + static_cast(patch_w) * 4); + } + } + + // 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. + if (callbacks_.GpuReady()) { + local_tex = callbacks_.UploadTexture( + local_tex, reinterpret_cast(entry.pixels.data()), + entry.width, entry.height); + VLOG(1, "uploadGuiImage for client_tex_id=%lu -> filament=%lu", + static_cast(tex_id), + static_cast(local_tex)); + } else { + VLOG(1, "ProcessCmdTexture: buffered tex_id=%lu, deferring GPU upload", + static_cast(tex_id)); + } +} +// Keep in sync with ProcessCmdDrawFrame in the vendored +// netimgui/Code/ServerApp/Source/NetImguiServer_RemoteClient.cpp +void RemoteUi::ProcessCmdDrawFrame(CmdDrawFrame* cmd_draw_frame) { + if (!cmd_draw_frame) return; + + // Take ownership to prevent pending_receive_ from deleting it prematurely. + pending_receive_.bAutoFree = false; + cmd_draw_frame->ToPointers(); + + if (cmd_draw_frame->mCompressed) { + if (last_uncompressed_frame_ != nullptr && + (last_uncompressed_frame_->mFrameIndex + 1) == + cmd_draw_frame->mFrameIndex) { + CmdDrawFrame* uncompressed_frame = DecompressCmdDrawFrame( + last_uncompressed_frame_.get(), cmd_draw_frame); + netImguiDeleteSafe(cmd_draw_frame); + cmd_draw_frame = uncompressed_frame; + } else { + // Missing previous / reference frame data. Ignore this delta-encoded + // drawframe and ask the client for a fresh uncompressed keyframe. + request_keyframe_ = true; + netImguiDeleteSafe(cmd_draw_frame); + return; + } + } + + // Release previous cached frame and store current for the next delta + // decompression. + last_uncompressed_frame_.reset(cmd_draw_frame); + + cmd_draw_frame->ToPointers(); + + RemoteDrawFrame* frame = netImguiNew(); + frame->frame_index = cmd_draw_frame->mFrameIndex; + ImDrawData* draw_data = &frame->draw_data; + draw_data->Valid = true; + draw_data->TotalVtxCount = + static_cast(cmd_draw_frame->mTotalVerticeCount); + draw_data->TotalIdxCount = + static_cast(cmd_draw_frame->mTotalIndiceCount); + + draw_data->DisplayPos.x = cmd_draw_frame->mDisplayArea[0]; + draw_data->DisplayPos.y = cmd_draw_frame->mDisplayArea[1]; + draw_data->DisplaySize.x = + cmd_draw_frame->mDisplayArea[2] - cmd_draw_frame->mDisplayArea[0]; + draw_data->DisplaySize.y = + cmd_draw_frame->mDisplayArea[3] - cmd_draw_frame->mDisplayArea[1]; + draw_data->FramebufferScale = ImGui::GetIO().DisplayFramebufferScale; + + ImDrawList* cmd_list = draw_data->CmdLists[0]; + cmd_list->IdxBuffer.resize(cmd_draw_frame->mTotalIndiceCount); + cmd_list->VtxBuffer.resize(cmd_draw_frame->mTotalVerticeCount); + cmd_list->CmdBuffer.resize(cmd_draw_frame->mTotalDrawCount); + // ImVector::resize() doesn't call constructors. Zero-init to ensure + // TexRef._TexData is NULL, not garbage. + memset(cmd_list->CmdBuffer.Data, 0, + cmd_list->CmdBuffer.Size * sizeof(ImDrawCmd)); + cmd_list->Flags = + ImDrawListFlags_AllowVtxOffset | ImDrawListFlags_AntiAliasedLines | + ImDrawListFlags_AntiAliasedFill | ImDrawListFlags_AntiAliasedLinesUseTex; + + constexpr float kPosRangeMin = static_cast(ImguiVert::kPosRange_Min); + constexpr float kPosRangeMax = static_cast(ImguiVert::kPosRange_Max); + constexpr float kUVRangeMin = static_cast(ImguiVert::kUvRange_Min); + constexpr float kUVRangeMax = static_cast(ImguiVert::kUvRange_Max); + + if (cmd_draw_frame->mTotalDrawCount != 0) { + // WebGL 1.0/2.0 often uses uint16_t for ImDrawIdx. Check for potential + // overflow if the total vertex count exceeds the limit for 16-bit indices. + if (sizeof(ImDrawIdx) == 2 && cmd_draw_frame->mTotalVerticeCount > 65536) { + fprintf(stderr, + "WARNING: NetImgui WASM Viewer received a draw frame with %u " + "vertices. This exceeds the maximum of 65536 for 16-bit " + "ImDrawIdx, potentially causing rendering artifacts due to index " + "wrapping.\n", + cmd_draw_frame->mTotalVerticeCount); + } + uint32_t index_offset(0), vertex_offset(0); + ImDrawIdx* index_dst = &cmd_list->IdxBuffer[0]; + ImDrawVert* vertex_dst = &cmd_list->VtxBuffer[0]; + ImDrawCmd* command_dst = &cmd_list->CmdBuffer[0]; + + for (uint32_t i(0); i < cmd_draw_frame->mDrawGroupCount; ++i) { + const ImguiDrawGroup& draw_group = cmd_draw_frame->mpDrawGroups[i]; + + // Indices + const uint16_t* indices = + reinterpret_cast(draw_group.mpIndices.Get()); + if (draw_group.mBytePerIndex == sizeof(ImDrawIdx)) { + memcpy(index_dst, indices, draw_group.mIndiceCount * sizeof(ImDrawIdx)); + } else { + for (uint32_t index_idx(0); index_idx < draw_group.mIndiceCount; + ++index_idx) { + index_dst[index_idx] = static_cast(indices[index_idx]); + } + } + + // Vertices — unpack quantized positions and UVs. + const ImguiVert* vertex_src = draw_group.mpVertices.Get(); + for (uint32_t vtx_idx(0); vtx_idx < draw_group.mVerticeCount; ++vtx_idx) { + vertex_dst[vtx_idx].pos.x = + (static_cast(vertex_src[vtx_idx].mPos[0]) * + (kPosRangeMax - kPosRangeMin)) / + static_cast(0xFFFF) + + kPosRangeMin + draw_group.mReferenceCoord[0]; + vertex_dst[vtx_idx].pos.y = + (static_cast(vertex_src[vtx_idx].mPos[1]) * + (kPosRangeMax - kPosRangeMin)) / + static_cast(0xFFFF) + + kPosRangeMin + draw_group.mReferenceCoord[1]; + vertex_dst[vtx_idx].uv.x = + (static_cast(vertex_src[vtx_idx].mUV[0]) * + (kUVRangeMax - kUVRangeMin)) / + static_cast(0xFFFF) + + kUVRangeMin; + vertex_dst[vtx_idx].uv.y = + (static_cast(vertex_src[vtx_idx].mUV[1]) * + (kUVRangeMax - kUVRangeMin)) / + static_cast(0xFFFF) + + kUVRangeMin; + vertex_dst[vtx_idx].col = vertex_src[vtx_idx].mColor; + } + + // Draw commands. + // NOTE: WebGL lacks glDrawElementsBaseVertex, so the backend's + // glDrawElements ignores VtxOffset. We bake the offset directly + // into the index values. + const ImguiDraw* draw_src = draw_group.mpDraws.Get(); + for (uint32_t draw_idx(0); draw_idx < draw_group.mDrawCount; ++draw_idx) { + uint32_t vtx_off = draw_src[draw_idx].mVtxOffset + vertex_offset; + uint32_t idx_off = draw_src[draw_idx].mIdxOffset + index_offset; + uint32_t elem_count = draw_src[draw_idx].mIdxCount; + + // Bake vertex offset into index values. + for (uint32_t ei = 0; ei < elem_count; ++ei) { + cmd_list->IdxBuffer[idx_off + ei] += static_cast(vtx_off); + } + + float cx = std::max( + 0.f, std::min(max_clip_[0], draw_src[draw_idx].mClipRect[0])); + float cy = std::max( + 0.f, std::min(max_clip_[1], draw_src[draw_idx].mClipRect[1])); + float cz = std::max( + cx, std::min(max_clip_[0], draw_src[draw_idx].mClipRect[2])); + float cw = std::max( + cy, std::min(max_clip_[1], draw_src[draw_idx].mClipRect[3])); + + command_dst[draw_idx].ClipRect.x = cx; + command_dst[draw_idx].ClipRect.y = cy; + command_dst[draw_idx].ClipRect.z = cz; + command_dst[draw_idx].ClipRect.w = cw; + command_dst[draw_idx].VtxOffset = 0; // Baked into indices + command_dst[draw_idx].IdxOffset = idx_off; + command_dst[draw_idx].ElemCount = elem_count; + command_dst[draw_idx].UserCallback = nullptr; + command_dst[draw_idx].UserCallbackData = nullptr; + + // Map remote ClientTextureID -> local GL handle. + ClientTextureID client_tex_id = draw_src[draw_idx].mClientTexId; + auto it = texture_map_.find(client_tex_id); + if (it != texture_map_.end() && it->second != 0) { + command_dst[draw_idx].TexRef._TexData = nullptr; + command_dst[draw_idx].TexRef._TexID = + static_cast(it->second); + } else { + LogUnmappedTexture(client_tex_id, texture_map_.size(), draw_idx, + texture_map_); + // Skip draw commands with unmapped textures. + command_dst[draw_idx].ElemCount = 0; + } + } + + index_dst += draw_group.mIndiceCount; + vertex_dst += draw_group.mVerticeCount; + command_dst += draw_group.mDrawCount; + index_offset += draw_group.mIndiceCount; + vertex_offset += draw_group.mVerticeCount; + } + } + + // Update internal write pointers to satisfy ImGui's draw list sanity checks. + // AddDrawListToDrawDataEx asserts that _VtxWritePtr/_IdxWritePtr point to the + // end of their respective buffers. Since we populated them via + // resize()+memcpy (bypassing ImGui's PrimReserve API), we must fix up these + // pointers manually. + cmd_list->_VtxWritePtr = cmd_list->VtxBuffer.Data + cmd_list->VtxBuffer.Size; + cmd_list->_IdxWritePtr = cmd_list->IdxBuffer.Data + cmd_list->IdxBuffer.Size; + cmd_list->_VtxCurrentIdx = cmd_list->VtxBuffer.Size; + + if (remote_draw_data_ == nullptr) { + LOG(Info, "First remote draw frame applied (%d cmds, %d vtx)", + draw_data->CmdLists[0]->CmdBuffer.Size, draw_data->TotalVtxCount); + } + remote_draw_data_.reset(frame); +} + +// Keep in sync with CaptureImguiInput in the vendored +// netimgui/Code/ServerApp/Source/NetImguiServer_RemoteClient.cpp +void RemoteUi::CaptureAndSendInput() { + if (!socket_) return; + + if (ConnectionState() != ReadyState::kOpen) return; + + // Capture input from Dear ImGui. + const ImGuiIO& io = ImGui::GetIO(); + + // When the local UI wants to capture mouse or keyboard, suppress the + // corresponding input from being forwarded to the remote client. This + // prevents clicks/keys meant for the local status overlay (or any other + // local window) from leaking through to the UI server. + const bool local_wants_mouse = io.WantCaptureMouse; + const bool local_wants_keyboard = io.WantCaptureKeyboard; + + { + // Only accumulate characters when the local UI is NOT capturing keyboard. + if (!local_wants_keyboard) { + const size_t initial_size = pending_input_chars_.size(); + const size_t added_char = io.InputQueueCharacters.size(); + if (added_char) { + pending_input_chars_.resize(initial_size + added_char); + memcpy(&pending_input_chars_[initial_size], + io.InputQueueCharacters.Data, added_char * sizeof(ImWchar)); + } + } + + // Gate scroll ACCUMULATION on local capture, the totals themselves are + // still sent every frame. This implements the "FREEZING" mentioned in the + // CmdInput wheel comment below. + if (!local_wants_mouse) { + mouse_wheel_pos_[0] += io.MouseWheel; + mouse_wheel_pos_[1] += io.MouseWheelH; + } + } + + CmdInput cmd_input; + cmd_input.mScreenSize[0] = static_cast(io.DisplaySize.x); + cmd_input.mScreenSize[1] = static_cast(io.DisplaySize.y); + // An unstable screen size forces the remote UI to relayout every frame + // (visible as UI flicker), so make changes loud. + if (cmd_input.mScreenSize[0] != last_screen_size_[0] || + cmd_input.mScreenSize[1] != last_screen_size_[1]) { + LOG(Info, "Screen size sent to client changed: %ux%u -> %ux%u", + last_screen_size_[0], last_screen_size_[1], cmd_input.mScreenSize[0], + cmd_input.mScreenSize[1]); + last_screen_size_[0] = cmd_input.mScreenSize[0]; + last_screen_size_[1] = cmd_input.mScreenSize[1]; + } + cmd_input.mFontDPIScaling = 1.f; + cmd_input.mDesiredFps = 60.0f; + cmd_input.mCompressionUse = kUseCompression; + cmd_input.mCompressionSkip = request_keyframe_; + + // NetImgui wheel fields are lifetime running totals, not per-frame deltas: + // the receiving side derives each frame's scroll as the difference between + // the current total and the previous total. So scroll is suppressed during + // local UI capture by FREEZING the totals (accumulation above is gated on + // !local_wants_mouse), while still sending them every frame. Sending 0 + // instead would rewind the receiver's baseline and produce two huge spurious + // deltas: -total when capture starts (wild zoom out as the local window + // expands) and +total when it ends (snap back on collapse). Note the + // deliberate asymmetry with the mouse position below, which IS per-frame + // absolute and can simply be parked while captured. + cmd_input.mMouseWheelVert = mouse_wheel_pos_[0]; + cmd_input.mMouseWheelHoriz = mouse_wheel_pos_[1]; + if (!local_wants_mouse) { + cmd_input.mMousePos[0] = static_cast(io.MousePos.x); + cmd_input.mMousePos[1] = static_cast(io.MousePos.y); + } else { + // Park the mouse off-screen so the remote side doesn't think we're + // hovering over anything. + cmd_input.mMousePos[0] = -1; + cmd_input.mMousePos[1] = -1; + } + + // Mouse button inputs. This static_assert detects when a Dear ImGui update + // changes ImGuiMouseButton, which requires updating NetImgui's enum copy. + static_assert( + static_cast(CmdInput::NetImguiMouseButton::ImGuiMouseButton_COUNT) == + static_cast(ImGuiMouseButton_::ImGuiMouseButton_COUNT), + "Update the NetImgui enum to match the updated Dear ImGui enum"); + cmd_input.mMouseDownMask = 0; + if (!local_wants_mouse) { + cmd_input.mMouseDownMask |= + ImGui::IsMouseDown(ImGuiMouseButton_::ImGuiMouseButton_Left) + ? 1 << CmdInput::ImGuiMouseButton_Left + : 0; + cmd_input.mMouseDownMask |= + ImGui::IsMouseDown(ImGuiMouseButton_::ImGuiMouseButton_Right) + ? 1 << CmdInput::ImGuiMouseButton_Right + : 0; + cmd_input.mMouseDownMask |= + ImGui::IsMouseDown(ImGuiMouseButton_::ImGuiMouseButton_Middle) + ? 1 << CmdInput::ImGuiMouseButton_Middle + : 0; + cmd_input.mMouseDownMask |= + ImGui::IsMouseDown(3) ? 1 << CmdInput::ImGuiMouseButton_Extra1 : 0; + cmd_input.mMouseDownMask |= + ImGui::IsMouseDown(4) ? 1 << CmdInput::ImGuiMouseButton_Extra2 : 0; + } + +// Keyboard inputs. These static_asserts detect when a Dear ImGui update +// changes ImGuiKey, which requires updating NetImgui's enum copy to match. +#define EnumKeynameTest(KEYNAME) \ + static_cast(CmdInput::NetImguiKeys::KEYNAME) == \ + static_cast(ImGuiKey::KEYNAME - ImGuiKey::ImGuiKey_NamedKey_BEGIN), \ + "Update the NetImgui enum to match the updated Dear ImGui enum" + static_assert( + CmdInput::NetImguiKeys::ImGuiKey_COUNT == + (ImGuiKey_NamedKey_END - ImGuiKey_NamedKey_BEGIN), + "Update the NetImgui enum to match the updated Dear ImGui enum"); + static_assert(EnumKeynameTest(ImGuiKey_Tab)); + static_assert(EnumKeynameTest(ImGuiKey_Escape)); + static_assert(EnumKeynameTest(ImGuiKey_RightSuper)); + static_assert(EnumKeynameTest(ImGuiKey_Apostrophe)); + static_assert(EnumKeynameTest(ImGuiKey_Keypad0)); + static_assert(EnumKeynameTest(ImGuiKey_CapsLock)); + static_assert(EnumKeynameTest(ImGuiKey_ReservedForModCtrl)); + static_assert(EnumKeynameTest(ImGuiKey_ReservedForModShift)); + static_assert(EnumKeynameTest(ImGuiKey_ReservedForModAlt)); + static_assert(EnumKeynameTest(ImGuiKey_ReservedForModSuper)); +#undef EnumKeynameTest + + // Save every keydown status to out bitmask — only when local UI is not + // capturing keyboard. + if (!local_wants_keyboard) { + uint64_t value_mask(0); + for (uint32_t i(0); i < ImGuiKey::ImGuiKey_NamedKey_COUNT; ++i) { + value_mask |= + ImGui::IsKeyDown(static_cast(ImGuiKey_NamedKey_BEGIN + i)) + ? 0x0000000000000001ull << (i % 64) + : 0; + if (((i % 64) == 63) || i == (ImGuiKey::ImGuiKey_NamedKey_COUNT - 1)) { + cmd_input.mInputDownMask[i / 64] = value_mask; + value_mask = 0; + } + } + } + // When local UI captures keyboard, mInputDownMask and mInputAnalog stay + // zero-initialized from the CmdInput constructor. + + // Copy waiting characters inputs — only when local UI is not capturing + // keyboard. When captured, pending chars are discarded to avoid buffering + // stale input that would replay when focus returns to remote. + if (!local_wants_keyboard) { + size_t added_key_count = std::min( + ArrayCount(cmd_input.mKeyChars) - cmd_input.mKeyCharCount, + pending_input_chars_.size()); + if (added_key_count) { + memcpy(&cmd_input.mKeyChars[cmd_input.mKeyCharCount], + &pending_input_chars_[0], added_key_count * sizeof(ImWchar)); + cmd_input.mKeyCharCount += static_cast(added_key_count); + size_t char_remain_count = pending_input_chars_.size() - added_key_count; + if (char_remain_count > 0) { + memcpy(&pending_input_chars_[0], &pending_input_chars_[added_key_count], + char_remain_count * sizeof(ImWchar)); + } + pending_input_chars_.resize(char_remain_count); + } + if (cmd_input.mKeyCharCount > 0) { + VLOG(1, "[web_client_remote_ui.cc] Queued %u characters to send\n", + cmd_input.mKeyCharCount); + } + } else { + // Discard any pending characters that were accumulated while local UI + // had keyboard focus. + pending_input_chars_.clear(); + } + + PendingCom pending_send; + pending_send.pCommand = &cmd_input; + pending_send.SizeCurrent = 0; + Network::DataSend(socket_, pending_send); + if (pending_send.IsDone() && cmd_input.mCompressionSkip) { + request_keyframe_ = false; + } +} + +void RemoteUi::Shutdown() { + remote_draw_data_.reset(); + last_uncompressed_frame_.reset(); + if (socket_) { + Network::Disconnect(socket_); + socket_ = nullptr; + } +} + +} // namespace mujoco::studio diff --git a/python/mujoco/experimental/studio/web/web_client_remote_ui.h b/python/mujoco/experimental/studio/web/web_client_remote_ui.h new file mode 100644 index 00000000..e7daea72 --- /dev/null +++ b/python/mujoco/experimental/studio/web/web_client_remote_ui.h @@ -0,0 +1,193 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// The NetImgui client side of the web viewer. Receives the headless Studio UI +// (draw frames/textures) over the /ui WebSocket and replies with browser input. + +#ifndef MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_REMOTE_UI_H_ +#define MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_REMOTE_UI_H_ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include "google/network_status.h" + +namespace mujoco::studio { + +// Deleter for objects allocated by NetImgui. +struct NetImguiDeleter { + template + void operator()(T* ptr) const { + NetImgui::Internal::netImguiDelete(ptr); + } +}; + +// The browser side of the streamed ImGui UI: connects to the UI WebSocket, +// receives NetImgui draw frames and textures from the Python side, assembles +// them into ImGui draw data for rendering, and sends the browser's input back. +class RemoteUi { + public: + using SocketInfo = NetImgui::Internal::Network::SocketInfo; + using ReadyState = NetImgui::Internal::Network::ReadyState; + using ClientTextureID = NetImgui::Internal::ClientTextureID; + + // The renderer-facing callbacks of the link; the app implements them once. + // The link reaches the GPU only through it, so the protocol logic stays + // renderer-agnostic. + class Callbacks { + public: + virtual ~Callbacks() = default; + + // Uploads a full RGBA texture and returns the GPU handle; nullptr + // pixels destroy the texture backing `current`. + virtual uintptr_t UploadTexture(uintptr_t current, const std::byte* rgba, + uint32_t width, uint32_t height) = 0; + + // True once the GPU context can accept texture uploads. + virtual bool GpuReady() = 0; + }; + + explicit RemoteUi(Callbacks& callbacks) : callbacks_(callbacks) {} + ~RemoteUi() { Shutdown(); } + + // (Re)connects to the UI WebSocket. Disconnects any existing socket first. + void Connect(const std::string& url); + + bool HasSocket() const { return socket_ != nullptr; } + + // Current socket state; browser WebSockets connect and close + // asynchronously, so this can differ from HasSocket() (see + // google/network_status.h). + ReadyState ConnectionState() const; + + // The WebSocket close code once the socket has closed, else 0. + int CloseCode() const; + + // Remote clip rects are clamped to this (logical pixels); set each frame + // before ReceiveAndProcessCommands(). + void SetMaxClip(float width, float height) { + max_clip_[0] = width; + max_clip_[1] = height; + } + + // Sends the CmdVersion handshake when the socket is newly open, then drains + // all incoming commands (draw frames, textures) for this frame. + void ReceiveAndProcessCommands(int frame); + + // Captures ImGui input and sends it to the remote client. Must run before + // ImGui::EndFrame(), because EndFrame() clears the io.InputQueueCharacters + // 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(); + + void Shutdown(); + + // The latest assembled remote draw data, or nullptr before the first frame. + ImDrawData* RemoteDrawData() { + return remote_draw_data_ ? &remote_draw_data_->draw_data : nullptr; + } + + // Returns the bytes received since the last call and resets the counter. + uint64_t ConsumeByteCount() { + uint64_t bytes = bytes_accum_; + bytes_accum_ = 0; + return bytes; + } + + void ProcessCmdDrawFrame(NetImgui::Internal::CmdDrawFrame* cmd); + void ProcessCmdTexture(NetImgui::Internal::CmdTexture* cmd); + + private: + // An assembled remote frame: an ImDrawData plus the single command list it + // points at (a plain ImDrawData only references externally-owned lists). + struct RemoteDrawFrame { + RemoteDrawFrame() + : command_list(ImGui::GetCurrentContext() + ? ImGui::GetDrawListSharedData() + : nullptr) { + draw_data.CmdLists.push_back(&command_list); + draw_data.CmdListsCount = 1; + } + + ImDrawData draw_data; + ImDrawList command_list; + uint64_t frame_index = 0; + }; + + Callbacks& callbacks_; + + float max_clip_[2] = {0.0f, 0.0f}; + + // --- Connection-scoped state, reset on every (re)connect. ---------------- + + SocketInfo* socket_ = nullptr; + bool handshake_sent_ = false; + bool was_connected_ = false; + ReadyState last_state_ = ReadyState::kDisconnected; + NetImgui::Internal::PendingCom pending_receive_; + NetImgui::Internal::CmdPendingRead cmd_pending_read_; + + // When true, the next CmdInput asks the client to send one uncompressed + // draw frame (CmdInput::mCompressionSkip). Set on (re)connect and whenever + // a delta frame arrives whose reference frame we don't have — without it, + // a browser joining mid-stream drops every delta-compressed frame forever + // and the remote UI never appears. Mirrors mbCompressionSkipOncePending in + // NetImguiServer RemoteClient. + bool request_keyframe_ = true; + + std::unique_ptr + last_uncompressed_frame_; + + // --- Session state. ------------------------------------------------------- + + std::unique_ptr remote_draw_data_; + std::unordered_map texture_map_; + + // CPU-side texture cache so partial updates from NetImgui can be patched into + // these buffers before re-uploading the full texture (partial updates into + // GPU textures are not supported by Filament). + struct TextureEntry { + std::vector pixels; // Full RGBA pixel data + uint32_t width = 0; + uint32_t height = 0; + }; + std::unordered_map texture_cpu_; + + // Persistent input state for character accumulation across frames. + std::vector pending_input_chars_; + float mouse_wheel_pos_[2] = {0.0f, 0.0f}; + uint16_t last_screen_size_[2] = {0, 0}; + + // --- Telemetry. ----------------------------------------------------------- + + uint64_t bytes_accum_ = 0; + int total_cmds_received_ = 0; + int draw_frames_received_ = 0; + int textures_received_ = 0; +}; + +} // namespace mujoco::studio + +#endif // MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_REMOTE_UI_H_ diff --git a/python/mujoco/experimental/studio/web/web_client_session.cc b/python/mujoco/experimental/studio/web/web_client_session.cc new file mode 100644 index 00000000..9b54319f --- /dev/null +++ b/python/mujoco/experimental/studio/web/web_client_session.cc @@ -0,0 +1,311 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "web_client_session.h" + +#include + +#include +#include + +#include "google/logging.h" + +namespace mujoco::studio { + +// Session-channel messages sent/received as text frames on the state WebSocket. +constexpr char kMsgRequestControl[] = "request_control"; +constexpr char kMsgLeaveQueue[] = "leave_queue"; +constexpr char kMsgForceControl[] = "force_control"; +constexpr char kMsgHeartbeat[] = "heartbeat"; +constexpr char kMsgStateAck[] = "state_ack"; +constexpr char kMsgGrant[] = "grant"; +constexpr char kMsgMaxSpectatorsPrefix[] = "max_spectators="; + +// Liveness heartbeat period. A hidden tab's rendering loop stops, so Update() +// and the heartbeat stop with it. The server kicks spectators (and releases a +// controller with a waiting queue) on silence. +constexpr double kHeartbeatSec = 30.0; + +// Minimum time between /ui claim retries (and between role machine steps). +constexpr double kUiRetrySec = 1.0; + +// After this many consecutive rejected claims the page stops claiming and +// settles into spectating. +constexpr int kMaxUiRejects = 3; + +EM_BOOL Session::OnWsMessage(int event_type, + const EmscriptenWebSocketMessageEvent* event, + void* user_data) { + auto* session = static_cast(user_data); + session->server_close_code_ = 0; // Accepted; hide the disconnect notice. + if (event->isText) { + // Text frames carry session metadata; emscripten null-terminates them. + session->OnSessionText(reinterpret_cast(event->data)); + } else { + session->HandleMessage(event->data, event->numBytes); + // Flow control: the server keeps at most one state payload in flight and + // sends the next (freshest) one only after this ack. Without it, a slow + // link buffers seconds of stale payloads in the socket and the whole viewer + // lags by that queue. + session->SendText(kMsgStateAck); + } + return EM_TRUE; +} + +EM_BOOL Session::OnWsOpen(int event_type, + const EmscriptenWebSocketOpenEvent* event, + void* user_data) { + auto* session = static_cast(user_data); + session->connected_ = true; + // server_close_code_ is NOT cleared here: a rejected connection also fires + // open before the server's closing code arrives. It clears on the first + // received message, which proves the server accepted us. + LOG(Info, "State WebSocket connected"); + return EM_TRUE; +} + +EM_BOOL Session::OnWsError(int event_type, + const EmscriptenWebSocketErrorEvent* event, + void* user_data) { + LOG(Error, "State WebSocket error"); + return EM_TRUE; +} + +EM_BOOL Session::OnWsClose(int event_type, + const EmscriptenWebSocketCloseEvent* event, + void* user_data) { + auto* session = static_cast(user_data); + LOG(Info, "State WebSocket closed (code=%d)", event->code); + session->connected_ = false; + + // Codes 4xxx are deliberate server-side closes (e.g. kWsCloseSessionFull). + // These conditions pass, so the GUI shows a notice while the reconnect loop + // retries at a slower pace. + if (event->code >= 4000 && event->code <= 4999) { + session->server_close_code_ = event->code; + LOG(Info, "Server ended this connection (code=%d); retrying slowly.", + event->code); + } + + // Free the handle; without this, every closed socket (including each failed + // reconnect) leaks its handle and callback registrations in Emscripten's + // socket table. Session (this) is a stable global, so the user_data of any + // already-queued event stays valid; detaching the callbacks first stops them + // firing on the freed handle. + session->CloseSocket(); + + return EM_TRUE; +} + +void Session::CloseSocket() { + if (socket_ <= 0) return; + emscripten_websocket_set_onopen_callback(socket_, nullptr, nullptr); + emscripten_websocket_set_onmessage_callback(socket_, nullptr, nullptr); + emscripten_websocket_set_onerror_callback(socket_, nullptr, nullptr); + emscripten_websocket_set_onclose_callback(socket_, nullptr, nullptr); + emscripten_websocket_delete(socket_); + socket_ = 0; + connected_ = false; +} + +void Session::Connect(const std::string& url) { + // Drop any lingering socket first so a reconnect cannot leak the old one. + CloseSocket(); + EmscriptenWebSocketCreateAttributes attr; + emscripten_websocket_init_create_attributes(&attr); + attr.url = url.c_str(); + attr.protocols = nullptr; + attr.createOnMainThread = EM_TRUE; + + socket_ = emscripten_websocket_new(&attr); + if (socket_ <= 0) { + LOG(Error, "Failed to create state WebSocket"); + return; + } + emscripten_websocket_set_onopen_callback(socket_, this, OnWsOpen); + emscripten_websocket_set_onmessage_callback(socket_, this, OnWsMessage); + emscripten_websocket_set_onerror_callback(socket_, this, OnWsError); + emscripten_websocket_set_onclose_callback(socket_, this, OnWsClose); + LOG(Info, "State WebSocket connecting to %s", url.c_str()); +} + +void Session::SendText(const char* text) { + if (socket_ && connected_) { + emscripten_websocket_send_utf8_text(socket_, text); + } +} + +void Session::SetRole(SessionRole role) { + role_ = role; + // Warning: clang-format splits `!==` into `!= =`, so avoid that syntax! + EM_ASM( + { Module.isSpectator = !!$0; }, + role == SessionRole::kControlling ? 0 : 1); +} + +bool ParseRoster(const char* text, Roster* roster) { + Roster parsed; + char role[16] = {0}; + if (sscanf(text, + "viewers=%d;role=%15[^;];queue_pos=%d;queue_len=%d;" + "max_spectators=%d", + &parsed.viewers, role, &parsed.queue_pos, &parsed.queue_len, + &parsed.max_spectators) != 5) { + return false; + } + parsed.spectator = strcmp(role, "spectator") == 0; + *roster = parsed; + return true; +} + +void Session::OnSessionText(const char* text) { + Roster roster; + if (ParseRoster(text, &roster)) { + if (roster.viewers != roster_.viewers || + roster.queue_pos != roster_.queue_pos || + roster.queue_len != roster_.queue_len) { + LOG(Info, "Session roster: %s", text); + } + roster_ = roster; + // The roster is authoritative about this page's role. Settling on it + // (rather than after several rejected /ui retries) makes the SPECTATING + // banner appear within the first roster (~200ms). Only while claiming with + // /ui closed: an in-flight claim must not be aborted, and an established + // controller is never demoted here (the kWsCloseControllerTaken close path + // handles that). + if (role_ == SessionRole::kClaiming && roster.spectator) { + if (remote_ui_state_ == RemoteUiState::kNoSocket || + remote_ui_state_ == RemoteUiState::kClosedOrError) { + LOG(Info, "Roster says spectator; settling"); + SetRole(SessionRole::kSpectating); + callbacks_.ShutdownRemoteUi(); + } + } + } else if (strcmp(text, kMsgGrant) == 0) { + // The controller slot is reserved for this page. The role flips to + // kControlling when the claim's socket opens. + LOG(Info, "Control granted; claiming the controller slot"); + SetRole(SessionRole::kClaiming); + ui_reject_count_ = 0; + callbacks_.ConnectRemoteUi(); + // Mark the claim in flight now: the roster broadcast right after the grant + // arrives before the next frame refreshes remote_ui_state_, and the settle + // rule above must not shut down the fresh claim. + remote_ui_state_ = RemoteUiState::kConnecting; + } +} + +void Session::FillView(SessionView* view) const { + view->role = role_; + view->viewers = roster_.viewers; + view->queue_pos = roster_.queue_pos; + view->queue_len = roster_.queue_len; + view->max_spectators = roster_.max_spectators; +} + +void Session::Update() { + const double now = emscripten_get_now() / 1000.0; + if (now - last_heartbeat_time_ >= kHeartbeatSec) { + last_heartbeat_time_ = now; + SendText(kMsgHeartbeat); + } +} + +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. + if (state == RemoteUiState::kNoSocket || role_ == SessionRole::kSpectating || + reload_pending_ || server_close_code_ != 0) { + return; + } + const double now = emscripten_get_now() / 1000.0; + if (now - last_ui_retry_time_ < kUiRetrySec) { + return; + } + if (state == RemoteUiState::kOpen) { + ui_reject_count_ = 0; + if (role_ != SessionRole::kControlling) { + SetRole(SessionRole::kControlling); // The claim succeeded. + } + } else if (state == RemoteUiState::kClosedOrError) { + last_ui_retry_time_ = now; + // A kWsCloseControllerTaken close while kControlling means another page + // took the slot (Steal Control): settle instantly. A rejected claim + // (kClaiming) retries a few times first, because a reloading controller + // briefly races its own slot. + const bool ousted = close_code == kWsCloseControllerTaken && + role_ == SessionRole::kControlling; + if (ousted || (close_code == kWsCloseControllerTaken && + ++ui_reject_count_ >= kMaxUiRejects)) { + LOG(Info, "Controller slot taken; spectating"); + SetRole(SessionRole::kSpectating); + // Also drops the last received UI frame: a page forced out of the + // controller slot must not keep showing a frozen Studio UI. + callbacks_.ShutdownRemoteUi(); + } else { + LOG(Info, "UI WebSocket closed; reconnecting..."); + callbacks_.ConnectRemoteUi(); + remote_ui_state_ = RemoteUiState::kConnecting; // As on the grant path. + } + } +} + +void Session::HandleMessage(const uint8_t* data, uint32_t num_bytes) { + last_message_time_ = emscripten_get_now() / 1000.0; + if (reload_pending_ || !callbacks_.ReadyForPayload()) { + return; + } + bytes_accum_ += num_bytes; + + StatePayloadView view; + if (!ParseStatePayload(data, num_bytes, &view)) { + LOG(Error, "Malformed state payload (%u bytes); dropping", num_bytes); + return; + } + + if (!model_crc32_.has_value()) { + model_crc32_ = view.model_crc32; + } 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); }); + return; + } + + callbacks_.OnPayload(view); +} + +void Session::RequestControl() { SendText(kMsgRequestControl); } + +void Session::LeaveQueue() { SendText(kMsgLeaveQueue); } + +void Session::StealControl() { SendText(kMsgForceControl); } + +void Session::ReleaseControl() { + // Become a spectator; the server grants the slot down the queue. + callbacks_.ShutdownRemoteUi(); + SetRole(SessionRole::kSpectating); +} + +void Session::SetCameraMode(int mode) { callbacks_.SetCameraMode(mode); } + +void Session::SetMaxSpectators(int count) { + char msg[48]; + snprintf(msg, sizeof(msg), "%s%d", kMsgMaxSpectatorsPrefix, count); + SendText(msg); +} + +} // namespace mujoco::studio diff --git a/python/mujoco/experimental/studio/web/web_client_session.h b/python/mujoco/experimental/studio/web/web_client_session.h new file mode 100644 index 00000000..312f055b --- /dev/null +++ b/python/mujoco/experimental/studio/web/web_client_session.h @@ -0,0 +1,249 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// The session: this page's relationship with the Python-side viewer. +// +// Owns the /state WebSocket (simulation payloads in, control messages out), +// the session wire protocol (roster, grants, heartbeats, acks), the role +// state machine (claiming -> controlling / spectating), and the +// model-change / page-reload / close-code policies. Everything outside the +// session goes through the Callbacks interface: applying a payload to the +// scene, and driving the remote UI stream on role transitions. The session +// also implements SessionActions, so the role window's intents land here +// directly. + +#ifndef MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_SESSION_H_ +#define MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_SESSION_H_ + +#include + +#include +#include +#include + +#include "state_payload.h" + +namespace mujoco::studio { + +// Our custom WebSocket close codes in range 4xxx. The client shows a notice +// and retries slowly. +constexpr int kWsCloseControllerTaken = 4001; // /ui: another browser controls. +constexpr int kWsCloseSessionFull = 4002; // /state: spectator limit hit. +constexpr int kWsCloseInactive = 4003; // /state: hidden tab kicked. +constexpr int kWsCloseNotController = 4004; // /drop: only controller may + // load models. + +// The page's role in the collaborative session. Every page starts by +// claiming the controller slot; the claim either succeeds (kControlling) +// or the page settles into spectating. A control grant puts a spectator +// back into kClaiming while it reconnects to /ui. +enum class SessionRole { + kClaiming = 0, // /ui claim in flight; the role is not yet resolved. + kControlling, // This page holds the open /ui connection. + kSpectating, // Another page controls; scene + local role window only. +}; + +// Read-only snapshot of the session, passed to the local UI each frame. +struct SessionView { + SessionRole role = SessionRole::kClaiming; + int viewers = 0; + int queue_pos = 0; // 1-based position in the control queue; 0 = unqueued. + int queue_len = 0; + int max_spectators = 0; + uint64_t gui_bytes_per_sec = 0; + uint64_t sim_bytes_per_sec = 0; + bool have_remote_frame = false; + int camera_mode = 0; // [SpectatorCamMode]. +}; + +// User intent reported by the role window. Session implements this; the +// interface is the complete list of effects the local UI can cause. +class SessionActions { + public: + virtual ~SessionActions() = default; + virtual void RequestControl() = 0; + virtual void LeaveQueue() = 0; + virtual void StealControl() = 0; + virtual void ReleaseControl() = 0; + virtual void SetCameraMode(int mode) = 0; // mode is [SpectatorCamMode]. + virtual void SetMaxSpectators(int count) = 0; // Already clamped by the UI. +}; + +// The roster: the server's membership broadcast, sent as a text frame on +// /state whenever the session changes (a viewer joins or leaves, queues +// for control, or control moves). It tells this page how many viewers are +// connected, which role the server currently assigns it, and where it +// stands in the control queue. +struct Roster { + int viewers = 0; + bool spectator = false; // The server's view: true = not the controller. + int queue_pos = 0; // 1-based position in the control queue; 0 = unqueued. + int queue_len = 0; + int max_spectators = 8; // Runtime spectator limit. +}; + +// Parses a roster line; returns false when text is not a roster. +bool ParseRoster(const char* text, Roster* roster); + +// The remote UI stream's connection state, reported to the role state +// machine by the app once per frame. +enum class RemoteUiState { + kNoSocket = 0, // No connection attempt exists. + kConnecting, // In flight (or closing); the machine waits. + kOpen, // The claim succeeded: this page controls. + kClosedOrError, // Rejected or dropped; the machine retries or settles. +}; + +class Session : public SessionActions { + public: + // Everything the session needs from the rest of the application. + class Callbacks { + public: + virtual ~Callbacks() = default; + // Payloads are dropped until this returns true (model loaded). + virtual bool ReadyForPayload() = 0; + // Applies a parsed payload to the application. + virtual void OnPayload(const StatePayloadView& view) = 0; + // Role transition: claim the controller slot. + virtual void ConnectRemoteUi() = 0; + // Role transition: drop the stream when spectating + virtual void ShutdownRemoteUi() = 0; + // Spectator camera mode change. + virtual void SetCameraMode(int mode) = 0; + }; + + explicit Session(Callbacks& callbacks) : callbacks_(callbacks) {} + + void Connect(const std::string& url); + + // Records the CRC32 of the model this page actually loaded. + // Must match zlib.crc32 (used to compute model_crc32 in web_viewer.py). + void SetModelCrc32(uint32_t crc) { model_crc32_ = crc; } + + // True while a connect attempt exists; used to pace reconnects. + // emscripten_websocket_new returns a handle immediately, so this is NOT the + // same as Connected(). + bool HasSocket() const { return socket_ != 0; } + + // 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 + // and retries slowly; the code clears when a connection opens again. + int ServerCloseCode() const { return server_close_code_; } + + // Returns the bytes received since the last call and resets the counter. + uint64_t ConsumeByteCount() { + uint64_t bytes = bytes_accum_; + bytes_accum_ = 0; + return bytes; + } + + // Wall-clock seconds of the last received message, or 0 before the first one. + // Payloads stream at ~60Hz while the Python side is alive, so staleness here + // means the server is gone, even if the socket still looks open (a suspended + // process keeps its sockets established). + double LastMessageTime() const { return last_message_time_; } + + SessionRole Role() const { return role_; } + + // Fills the role and roster fields of the view. + void FillView(SessionView* view) const; + + // Periodic session upkeep (the ~30s liveness heartbeat); call once per frame. + void Update(); + + // Feeds the role state machine the remote UI stream's connection state; call + // once per frame. Owns claim retry pacing, promotion to kControlling when a + // claim opens, instant settling when an open stream closes with + // kWsCloseControllerTaken (ousted by Steal Control), and the retries then + // settle rule for rejected claims. + void HandleRemoteUiState(RemoteUiState state, int close_code); + + // Parses one WebSocket message and applies the model-change/reload policy. + void HandleMessage(const uint8_t* data, uint32_t num_bytes); + + // SessionActions (used to implement the role window UI). + void RequestControl() override; + void LeaveQueue() override; + void StealControl() override; + void ReleaseControl() override; + void SetCameraMode(int mode) override; + void SetMaxSpectators(int count) override; + + private: + // Detaches callbacks and frees socket_ (if any), resetting to disconnected. + void CloseSocket(); + + // Sends a session message (control requests, acks, activity reports) to the + // server as a text frame. Dropped silently while not connected. + void SendText(const char* text); + + // Updates the role and mirrors it into JS (Module.isSpectator), which gates + // controller-only page behavior (model drag-and-drop upload). + void SetRole(SessionRole role); + + // Routes a session text frame: roster updates, control grants. + void OnSessionText(const char* text); + + static EM_BOOL OnWsMessage(int event_type, + const EmscriptenWebSocketMessageEvent* event, + void* user_data); + static EM_BOOL OnWsOpen(int event_type, + const EmscriptenWebSocketOpenEvent* event, + void* user_data); + static EM_BOOL OnWsError(int event_type, + const EmscriptenWebSocketErrorEvent* event, + void* user_data); + static EM_BOOL OnWsClose(int event_type, + const EmscriptenWebSocketCloseEvent* event, + void* user_data); + + Callbacks& callbacks_; + + 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. + std::optional model_crc32_; + bool reload_pending_ = false; + + int server_close_code_ = 0; + + uint64_t bytes_accum_ = 0; + double last_message_time_ = 0; + + // Role state machine and roster. + SessionRole role_ = SessionRole::kClaiming; + Roster roster_; + + // Consecutive rejected /ui claims; the page eventually stops claiming and + // settles into spectating. + int ui_reject_count_ = 0; + RemoteUiState remote_ui_state_ = RemoteUiState::kNoSocket; + double last_ui_retry_time_ = 0; + double last_heartbeat_time_ = 0; +}; + +} // namespace mujoco::studio + +#endif // MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_WEB_CLIENT_SESSION_H_ diff --git a/src/experimental/platform/ux/imgui_widgets.h b/src/experimental/platform/ux/imgui_widgets.h index 29dbd860..a13ed207 100644 --- a/src/experimental/platform/ux/imgui_widgets.h +++ b/src/experimental/platform/ux/imgui_widgets.h @@ -44,20 +44,38 @@ inline float GetStableAvailWidth() { // Walk up past child windows to the actual panel that owns the scrollbar. ImGuiWindow* scroll_owner = window; - while (scroll_owner && - (scroll_owner->Flags & ImGuiWindowFlags_ChildWindow) && + while (scroll_owner && (scroll_owner->Flags & ImGuiWindowFlags_ChildWindow) && scroll_owner->ParentWindow) { scroll_owner = scroll_owner->ParentWindow; } - if (scroll_owner && - !(scroll_owner->Flags & ImGuiWindowFlags_NoScrollbar) && + if (scroll_owner && !(scroll_owner->Flags & ImGuiWindowFlags_NoScrollbar) && !scroll_owner->ScrollbarY) { avail_x -= ImGui::GetStyle().ScrollbarSize; } return avail_x; } +// Draws one line of text horizontally centered in the current window. +inline void CenteredLine(const char* text, const ImVec4* color = nullptr) { + ImGui::SetCursorPosX(ImMax( + 0.0f, (ImGui::GetWindowWidth() - ImGui::CalcTextSize(text).x) * 0.5f)); + if (color != nullptr) { + ImGui::TextColored(*color, "%s", text); + } else { + ImGui::TextUnformatted(text); + } +} + +// Large centered banner text. +inline void CenteredBanner(const char* text, const ImVec4& color) { + ImGui::SetWindowFontScale(1.6f); + ImGui::SetCursorPosX(ImMax( + 0.0f, (ImGui::GetWindowWidth() - ImGui::CalcTextSize(text).x) * 0.5f)); + ImGui::TextColored(color, "%s", text); + ImGui::SetWindowFontScale(1.0f); +} + // FontAwesome icon codes. static constexpr const char ICON_FA_ADJUST[] = "\xEF\x81\x82"; static constexpr const char ICON_FA_ARROWS[] = "\xEF\x81\x87"; @@ -160,9 +178,7 @@ struct ScopedStyle { return *this; } - ScopedStyle& Font(ScopedFont font) { - return Font(static_cast(font)); - } + ScopedStyle& Font(ScopedFont font) { return Font(static_cast(font)); } ScopedStyle& Color(ImGuiCol col, ImColor color) { ImGui::PushStyleColor(col, (ImU32)color); @@ -648,8 +664,8 @@ inline bool ImGui_ColorButtonEx(const char* label, bool active, ImColor color, // Draw border. if (s.FrameBorderSize > 0) { - dl->AddRect(pos, max, ImGui::GetColorU32(ImGuiCol_Border), - s.FrameRounding, corners, s.FrameBorderSize); + dl->AddRect(pos, max, ImGui::GetColorU32(ImGuiCol_Border), s.FrameRounding, + corners, s.FrameBorderSize); } // Draw label centered.