diff --git a/python/mujoco/experimental/studio/web/state_payload.cc b/python/mujoco/experimental/studio/web/state_payload.cc new file mode 100644 index 00000000..608f881b --- /dev/null +++ b/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 +#include +#include +#include + +#include + +namespace mujoco::studio { +namespace { + +// Appends raw bytes to the payload buffer. +void AppendBytes(std::vector& buffer, const void* data, + size_t size) { + const std::byte* bytes = static_cast(data); + buffer.insert(buffer.end(), bytes, bytes + size); +} + +// Appends a complete [u32 tag][u32 size][payload] block. +void AppendStateBlock(std::vector& buffer, uint32_t tag, + const void* data, size_t size) { + StateBlockHeader block_header{tag, static_cast(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& 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(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 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& render_flags, + const mjvGeom* extra_geoms, size_t extra_geom_count) { + extra_geom_count = + extra_geom_count > kMaxExtraGeoms ? kMaxExtraGeoms : extra_geom_count; + std::vector 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(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(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(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 diff --git a/python/mujoco/experimental/studio/web/state_payload.h b/python/mujoco/experimental/studio/web/state_payload.h new file mode 100644 index 00000000..0ceacf92 --- /dev/null +++ b/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 +#include +#include + +#include + +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 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& 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_ diff --git a/python/mujoco/experimental/studio/web/state_payload_py.cc b/python/mujoco/experimental/studio/web/state_payload_py.cc new file mode 100644 index 00000000..9be00030 --- /dev/null +++ b/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 +#include +#include +#include + +#include +#include "state_payload.h" +#include "structs.h" +#include +#include + +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& render_flags, + const std::vector& extra_geoms) { + std::vector 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 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(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; +}