Add StatePayload and serialization helpers for web viewer

PiperOrigin-RevId: 956400340
Change-Id: Ibe049afe10eca85e5fb6d718a78863a641b864bb

python/mujoco/experimental/studio/web/state_payload.cc

@@ -0,0 +1,171 @@
+// 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.
+
+// Wire-format implementation for the web viewer's state payload (see
+// state_payload.h). Compiled into BOTH sides of the wire: the state_payload
+// pybind module (serializer, via state_payload_py.cc) and the wasm
+// web_client (parser, via web_client_session.cc). It must therefore stay
+// free of python- or browser-specific dependencies.
+
+#include "state_payload.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <vector>
+
+#include <mujoco/mujoco.h>
+
+namespace mujoco::studio {
+namespace {
+
+// Appends raw bytes to the payload buffer.
+void AppendBytes(std::vector<std::byte>& buffer, const void* data,
+                 size_t size) {
+  const std::byte* bytes = static_cast<const std::byte*>(data);
+  buffer.insert(buffer.end(), bytes, bytes + size);
+}
+
+// Appends a complete [u32 tag][u32 size][payload] block.
+void AppendStateBlock(std::vector<std::byte>& buffer, uint32_t tag,
+                      const void* data, size_t size) {
+  StateBlockHeader block_header{tag, static_cast<uint32_t>(size)};
+  AppendBytes(buffer, &block_header, sizeof(block_header));
+  AppendBytes(buffer, data, size);
+}
+
+// Serializes the render state (exactly kRenderStateSize bytes) into `ptr`.
+void SerializeRenderStateInto(std::byte* ptr, const mjvCamera& camera,
+                              const mjvPerturb& perturb,
+                              const mjvOption& vis_options, const mjOption& opt,
+                              const mjVisual& vis, const mjStatistic& stat,
+                              const std::vector<uint8_t>& render_flags) {
+  memcpy(ptr, &camera, sizeof(mjvCamera));
+  ptr += sizeof(mjvCamera);
+
+  memcpy(ptr, &perturb, sizeof(mjvPerturb));
+  ptr += sizeof(mjvPerturb);
+
+  memcpy(ptr, &vis_options, sizeof(mjvOption));
+  ptr += sizeof(mjvOption);
+
+  memcpy(ptr, &opt, sizeof(mjOption));
+  ptr += sizeof(mjOption);
+
+  memcpy(ptr, &vis, sizeof(mjVisual));
+  ptr += sizeof(mjVisual);
+
+  memcpy(ptr, &stat, sizeof(mjStatistic));
+  ptr += sizeof(mjStatistic);
+
+  // Pack render flags (mjNRNDFLAG bytes).
+  memset(ptr, 0, mjNRNDFLAG);
+  for (size_t i = 0; i < mjNRNDFLAG && i < render_flags.size(); ++i) {
+    ptr[i] = static_cast<std::byte>(render_flags[i]);
+  }
+}
+
+}  // namespace
+
+size_t MaxStatePayloadSize(size_t physics_bytes) {
+  return sizeof(StatePayloadHeader) + 3 * sizeof(StateBlockHeader) +
+         (sizeof(int32_t) + physics_bytes) + kRenderStateSize +
+         kMaxExtraGeoms * sizeof(mjvGeom);
+}
+
+std::vector<std::byte> SerializeStatePayload(
+    uint32_t model_crc32, int32_t physics_spec, const void* physics,
+    size_t physics_bytes, const mjvCamera& camera, const mjvPerturb& perturb,
+    const mjvOption& vis_options, const mjOption& opt, const mjVisual& vis,
+    const mjStatistic& stat, const std::vector<uint8_t>& render_flags,
+    const mjvGeom* extra_geoms, size_t extra_geom_count) {
+  extra_geom_count =
+      extra_geom_count > kMaxExtraGeoms ? kMaxExtraGeoms : extra_geom_count;
+  std::vector<std::byte> buffer;
+  buffer.reserve(MaxStatePayloadSize(physics_bytes));
+
+  StatePayloadHeader header;
+  header.nblocks = extra_geom_count > 0 ? 3 : 2;
+  header.model_crc32 = model_crc32;
+  AppendBytes(buffer, &header, sizeof(header));
+
+  // Physics state: [i32 spec][mjtNum values...].
+  StateBlockHeader physics_header{
+      kTagPhysicsState, static_cast<uint32_t>(sizeof(int32_t) + physics_bytes)};
+  AppendBytes(buffer, &physics_header, sizeof(physics_header));
+  AppendBytes(buffer, &physics_spec, sizeof(int32_t));
+  AppendBytes(buffer, physics, physics_bytes);
+
+  // Render state, serialized into place.
+  StateBlockHeader render_header{kTagRenderState,
+                                 static_cast<uint32_t>(kRenderStateSize)};
+  AppendBytes(buffer, &render_header, sizeof(render_header));
+  const size_t render_offset = buffer.size();
+  buffer.resize(render_offset + kRenderStateSize);
+  SerializeRenderStateInto(buffer.data() + render_offset, camera, perturb,
+                           vis_options, opt, vis, stat, render_flags);
+
+  // Extra geoms (only when present).
+  if (extra_geom_count > 0) {
+    AppendStateBlock(buffer, kTagExtraGeoms, extra_geoms,
+                     extra_geom_count * sizeof(mjvGeom));
+  }
+
+  return buffer;
+}
+
+bool ParseStatePayload(const void* data, size_t size, StatePayloadView* out) {
+  const std::byte* bytes = static_cast<const std::byte*>(data);
+  if (size < sizeof(StatePayloadHeader)) return false;
+
+  StatePayloadHeader header;
+  memcpy(&header, bytes, sizeof(header));
+  if (header.magic != kStatePayloadMagic) return false;
+  if (header.version != kStatePayloadVersion) return false;
+  out->model_crc32 = header.model_crc32;
+
+  size_t offset = sizeof(StatePayloadHeader);
+  for (uint16_t i = 0; i < header.nblocks; ++i) {
+    if (offset + sizeof(StateBlockHeader) > size) return false;
+    StateBlockHeader block;
+    memcpy(&block, bytes + offset, sizeof(block));
+    offset += sizeof(StateBlockHeader);
+    if (offset + block.size > size) return false;
+    const std::byte* payload = bytes + offset;
+
+    switch (block.tag) {
+      case kTagPhysicsState:
+        if (block.size < sizeof(int32_t)) return false;
+        memcpy(&out->physics_spec, payload, sizeof(int32_t));
+        out->physics = payload + sizeof(int32_t);
+        out->physics_bytes = block.size - sizeof(int32_t);
+        break;
+      case kTagRenderState:
+        if (block.size != kRenderStateSize) return false;
+        out->render_state = payload;
+        break;
+      case kTagExtraGeoms:
+        if (block.size % sizeof(mjvGeom) != 0) return false;
+        out->extra_geoms = payload;
+        out->extra_geom_count = block.size / sizeof(mjvGeom);
+        break;
+      default:
+        break;  // Unknown tag: skip.
+    }
+    offset += block.size;
+  }
+  return true;
+}
+
+}  // namespace mujoco::studio

python/mujoco/experimental/studio/web/state_payload.h

@@ -0,0 +1,128 @@
+// 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 defines the serialization format for the web viewer's browser
+// client render payload containing the data needed so that the browser can
+// render the scene using the following call:
+//
+// Render(model, data, perturb, camera, vis_options, width, height, extra_geoms)
+//
+// The arguments come from the Python process:
+//
+//   * model       : fetched once over HTTP as /model.mjb; its runtime-mutable
+//                   parts (opt/vis/stat) are re-sent in the render state block.
+//   * data        : streamed as the physics state vector (mjSTATE_INTEGRATION);
+//                   the browser recomputes the rest via mj_setState/mj_forward.
+//   * width/height: the browser canvas size.
+//   * extra_geoms : optional variable-size kTagExtraGeoms block.
+//   * ...         : the rest of the arguments are sent as a fixed-size block
+//
+// The payload (SerializeStatePayload) is a sequence of tagged blocks:
+//
+//   [StatePayloadHeader][u32 tag][u32 size][payload]...
+//
+// The payload is serialized by Python, sent over the /state WebSocket, and
+// parsed by the browser.
+//
+// TODO(matijak): Try shrinking the physics block: float32 (or quantized) values
+// instead of doubles, and/or delta-encoding against the client's last-acked
+// payload. The /state ack (web_server.py) tells the server which snapshot each
+// client last applied, which is the baseline that delta compression needs. For
+// 100humanoids.xml the payload is ~181 KB of doubles and dominates slow links.
+
+#ifndef MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_STATE_PAYLOAD_H_
+#define MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_STATE_PAYLOAD_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+#include <mujoco/mujoco.h>
+
+namespace mujoco::studio {
+
+// "MJWS" as little-endian bytes. This magic constant identifies the
+// StateServer WebSocket payload header and helps detect malformed or
+// misrouted messages.
+constexpr uint32_t kStatePayloadMagic =
+    'M' | ('J' << 8) | ('W' << 16) | ('S' << 24);
+constexpr uint16_t kStatePayloadVersion = 1;
+
+struct StatePayloadHeader {
+  uint32_t magic = kStatePayloadMagic;
+  uint16_t version = kStatePayloadVersion;
+  uint16_t nblocks = 0;
+  // CRC32 of the model's MJB bytes. When this changes, the browser must
+  // refetch /model.mjb before applying any further state.
+  uint32_t model_crc32 = 0;
+};
+static_assert(sizeof(StatePayloadHeader) == 12);
+
+// Block tags. Readers must skip unknown tags.
+enum StateBlockTag : uint32_t {
+  kTagPhysicsState = 1,  // [i32 mjtState spec signature][mjtNum values...]
+  kTagRenderState = 2,   // fixed-size block of kRenderStateSize bytes
+  kTagExtraGeoms = 3,    // n x mjvGeom (n = size / sizeof(mjvGeom))
+};
+
+struct StateBlockHeader {
+  uint32_t tag = 0;
+  uint32_t size = 0;
+};
+static_assert(sizeof(StateBlockHeader) == 8);
+
+// Fixed byte size of the render state block appended after physics state.
+// These are plain C structs of int/float/double members whose total size is
+// fixed, independent of the model and generally negligible compared to the size
+// of the physics state
+constexpr size_t kRenderStateSize =
+    sizeof(mjvCamera) + sizeof(mjvPerturb) + sizeof(mjvOption) +
+    sizeof(mjOption) + sizeof(mjVisual) + sizeof(mjStatistic) + mjNRNDFLAG;
+
+// Maximum number of extra geoms serialized per frame. Bounds the shared
+// memory buffer the StateServer allocates; WebViewer truncates longer lists.
+constexpr uint32_t kMaxExtraGeoms = 1024;
+
+// Upper bound of a serialized payload, used to size the StateServer's shared
+// memory buffer. `physics_bytes` is mj_stateSize(...) * sizeof(mjtNum).
+size_t MaxStatePayloadSize(size_t physics_bytes);
+
+// Serialize the complete state payload sent over the state WebSocket.
+std::vector<std::byte> SerializeStatePayload(
+    uint32_t model_crc32, int32_t physics_spec, const void* physics,
+    size_t physics_bytes, const mjvCamera& camera, const mjvPerturb& perturb,
+    const mjvOption& vis_options, const mjOption& opt, const mjVisual& vis,
+    const mjStatistic& stat, const std::vector<uint8_t>& render_flags,
+    const mjvGeom* extra_geoms, size_t extra_geom_count);
+
+// Parsed view into a serialized payload. Pointers alias the input buffer and
+// are NOT guaranteed to be aligned; so you must memcpy the data out before use.
+struct StatePayloadView {
+  uint32_t model_crc32 = 0;
+  int32_t physics_spec = 0;
+  const std::byte* physics = nullptr;
+  size_t physics_bytes = 0;
+  const std::byte* render_state = nullptr;  // kRenderStateSize bytes when non-null
+  const std::byte* extra_geoms = nullptr;   // extra_geom_count * sizeof(mjvGeom)
+  size_t extra_geom_count = 0;
+};
+
+// Parses a payload produced by SerializeStatePayload. Returns false if the
+// buffer is malformed (bad magic/version or out-of-bounds block). Blocks
+// with unknown tags are skipped.
+bool ParseStatePayload(const void* data, size_t size, StatePayloadView* out);
+
+}  // namespace mujoco::studio
+
+#endif  // MUJOCO_PYTHON_EXPERIMENTAL_STUDIO_WEB_STATE_PAYLOAD_H_

python/mujoco/experimental/studio/web/state_payload_py.cc

@@ -0,0 +1,72 @@
+// 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.
+
+// Python bindings for the state payload wire format (state_payload.h).
+//
+// WebViewer serializes the /state WebSocket payload with this module each
+// frame; the browser parses it with the same header (web_client_session).
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <mujoco/mujoco.h>
+#include "state_payload.h"
+#include "structs.h"
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+// Serialize the complete state WebSocket payload (see state_payload.h):
+// physics state, render state and extra geoms as tagged blocks.
+static py::bytes SerializeStatePayload(
+    uint32_t model_crc32, int physics_spec, const py::bytes& physics_state,
+    const mujoco::python::MjvCameraWrapper& camera,
+    const mujoco::python::MjvPerturbWrapper& perturb,
+    const mujoco::python::MjvOptionWrapper& vis_options,
+    const mujoco::python::MjModelWrapper& model,
+    const std::vector<uint8_t>& render_flags,
+    const std::vector<mujoco::python::MjvGeomWrapper>& extra_geoms) {
+  std::vector<mjvGeom> geoms;
+  geoms.reserve(extra_geoms.size());
+  for (const mujoco::python::MjvGeomWrapper& geom_wrapper : extra_geoms) {
+    if (geom_wrapper.get()) {
+      geoms.push_back(*geom_wrapper.get());
+    }
+  }
+
+  std::string physics = physics_state;
+  const std::vector<std::byte> buffer = mujoco::studio::SerializeStatePayload(
+      model_crc32, physics_spec, physics.data(), physics.size(), *camera.get(),
+      *perturb.get(), *vis_options.get(), model.get()->opt, model.get()->vis,
+      model.get()->stat, render_flags, geoms.data(), geoms.size());
+  return py::bytes(reinterpret_cast<const char*>(buffer.data()), buffer.size());
+}
+
+// Upper bound of a serialized payload for a model whose physics state is
+// `physics_bytes` long. Used to size the StateServer's shared memory.
+static size_t MaxStatePayloadSize(size_t physics_bytes) {
+  return mujoco::studio::MaxStatePayloadSize(physics_bytes);
+}
+
+PYBIND11_MODULE(state_payload, m, pybind11::mod_gil_not_used()) {
+  py::module_::import("mujoco._structs");
+  m.doc() = "MuJoCo web viewer state payload serialization";
+
+  m.def("serialize_state_payload", &SerializeStatePayload);
+  m.def("max_state_payload_size", &MaxStatePayloadSize);
+  m.attr("MAX_EXTRA_GEOMS") = mujoco::studio::kMaxExtraGeoms;
+}