Matija Kecman
c2e95b4161
PiperOrigin-RevId: 956400340 Change-Id: Ibe049afe10eca85e5fb6d718a78863a641b864bb
73 lines
2.9 KiB
C++
73 lines
2.9 KiB
C++
// 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;
|
|
}
|