From 37762e3f7008926fbdf339fc1601e0a398096eed Mon Sep 17 00:00:00 2001 From: Google DeepMind Date: Tue, 13 Jan 2026 08:35:51 -0800 Subject: [PATCH] Add support for `vfs` assets in `mj_compile` function to MuJoCo WASM. Regarding mj_compile VFS argument: I was encountering a memory error when using `std::optional` with `mj_compile` (because vfs should be an optional argument of the function). This is likely because `MjVFS` doesn't have copy/move constructors, which are required for types used with `std::optional`, causing issues with memory management when std::optional copies the MjVFS object. To solve this and to make vfs an optional argument to `mj_compile`, I replaced `mj_compile_wrapper` with two overloaded functions. This avoids using `std::optional` and resolves the memory issue, while still allowing mj_compile to be called with or without the vfs argument from JavaScript. But what if we instead add those copy/move constructors so it works with `std::optional`?: While making `MjVFS` copyable or movable would theoretically allow it to work with std::optional, it's not a feasible approach in this case. The MjVFS struct manages file data in memory, allocating buffers when files are added and freeing them when `mj_deleteVFS` is called in its destructor. A copy would require a deep copy of the entire virtual filesystem, but MuJoCo does not provide a public API function (like a "mj_copyVFS") to do this. Implementing a deep copy manually would be complex and fragile. A shallow copy would result in multiple MjVFS objects pointing to the same memory, leading to double-free errors when their destructors run. Given these constraints, MjVFS cannot safely be made copyable. The approach of providing two overloaded C++ functions is better. PiperOrigin-RevId: 855740541 Change-Id: Ia08b6db8868ad245f07145db3fb0322cb9f737cb --- wasm/codegen/generated/bindings.cc | 50 ++++++++++++++++++++++++---- wasm/codegen/generators/constants.py | 33 +++++++++--------- wasm/codegen/generators/functions.py | 5 +-- wasm/codegen/templates/bindings.cc | 47 ++++++++++++++++++++++++-- wasm/tests/bindings_test.ts | 50 ++++++++++++++++++++++++++-- 5 files changed, 157 insertions(+), 28 deletions(-) diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index ed60e6b5..6a31d3e8 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -31,9 +31,9 @@ #include #include +#include #include #include -#include #include "engine/engine_util_errmem.h" #include "wasm/unpack.h" @@ -7638,6 +7638,27 @@ struct MjvScene { std::vector camera; }; +struct MjVFS { + MjVFS() : ptr_(new mjVFS) { mj_defaultVFS(ptr_); } + ~MjVFS() { + mj_deleteVFS(ptr_); + } + void AddBuffer(const std::string& name, const emscripten::val& buffer) { + std::vector vec = emscripten::vecFromJSArray(buffer); + int result = mj_addBufferVFS(ptr_, name.c_str(), vec.data(), vec.size()); + if (result != 0) { + mju_error("Could not add buffer to VFS: %d", result); + } + } + void DeleteFile(const std::string& filename) { + mj_deleteFileVFS(ptr_, filename.c_str()); + } + mjVFS* get() const { return ptr_; } + + private: + mjVFS* ptr_; +}; + MjModel::MjModel(mjModel* ptr) : ptr_(ptr), opt(&ptr->opt), vis(&ptr->vis), stat(&ptr->stat) {} @@ -7833,7 +7854,7 @@ std::unique_ptr parseXMLString_wrapper(const std::string &xml) { return std::unique_ptr(new MjSpec(ptr)); } -std::unique_ptr mj_compile_wrapper(const MjSpec& spec) { +std::unique_ptr mj_compile_wrapper_1(const MjSpec& spec) { mjSpec* spec_ptr = spec.get(); mjModel* model = mj_compile(spec_ptr, nullptr); if (!model || mjs_isWarning(spec_ptr)) { @@ -7842,6 +7863,16 @@ std::unique_ptr mj_compile_wrapper(const MjSpec& spec) { return std::unique_ptr(new MjModel(model)); } +std::unique_ptr mj_compile_wrapper_2(const MjSpec& spec, const MjVFS& vfs) { + mjSpec* spec_ptr = spec.get(); + mjVFS* vfs_ptr = vfs.get(); + mjModel* model = mj_compile(spec_ptr, vfs_ptr); + if (!model || mjs_isWarning(spec_ptr)) { + mju_error("%s", mjs_getError(spec_ptr)); + } + return std::unique_ptr(new MjModel(model)); +} + void error_wrapper(const String& msg) { mju_error("%s\n", msg.as().data()); } int mj_saveLastXML_wrapper(const String& filename, const MjModel& m) { @@ -12334,6 +12365,11 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { emscripten::register_optional(); emscripten::register_optional(); + emscripten::class_("MjVFS") + .constructor<>() + .function("addBuffer", &MjVFS::AddBuffer) + .function("deleteFile", &MjVFS::DeleteFile); + function("mj_Euler", &mj_Euler_wrapper); function("mj_RungeKutta", &mj_RungeKutta_wrapper); function("mj_addContact", &mj_addContact_wrapper); @@ -12348,7 +12384,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { function("mj_comPos", &mj_comPos_wrapper); function("mj_comVel", &mj_comVel_wrapper); function("mj_compareFwdInv", &mj_compareFwdInv_wrapper); - function("mj_compile", &mj_compile_wrapper); function("mj_constraintUpdate", &mj_constraintUpdate_wrapper); function("mj_contactForce", &mj_contactForce_wrapper); function("mj_copyBack", &mj_copyBack_wrapper); @@ -12427,13 +12462,11 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { function("mj_resetDataKeyframe", &mj_resetDataKeyframe_wrapper); function("mj_rne", &mj_rne_wrapper); function("mj_rnePostConstraint", &mj_rnePostConstraint_wrapper); - function("mj_saveLastXML", &mj_saveLastXML_wrapper); function("mj_sensorAcc", &mj_sensorAcc_wrapper); function("mj_sensorPos", &mj_sensorPos_wrapper); function("mj_sensorVel", &mj_sensorVel_wrapper); function("mj_setConst", &mj_setConst_wrapper); function("mj_setKeyframe", &mj_setKeyframe_wrapper); - function("mj_setLengthRange", &mj_setLengthRange_wrapper); function("mj_setState", &mj_setState_wrapper); function("mj_setTotalmass", &mj_setTotalmass_wrapper); function("mj_sizeModel", &mj_sizeModel_wrapper); @@ -12710,9 +12743,14 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { function("mjv_updateCamera", &mjv_updateCamera_wrapper); function("mjv_updateScene", &mjv_updateScene_wrapper); function("mjv_updateSkin", &mjv_updateSkin_wrapper); - function("parseXMLString", &parseXMLString_wrapper, take_ownership()); function("error", &error_wrapper); + function("mj_saveLastXML", &mj_saveLastXML_wrapper); + function("mj_setLengthRange", &mj_setLengthRange_wrapper); + // mj_compile is bound using two overloads to handle the optional MjVFS argument, + // as using std::optional caused memory errors due to missing copy/move constructors. + function("mj_compile", emscripten::select_overload(const MjSpec&)>(&mj_compile_wrapper_1)); + function("mj_compile", emscripten::select_overload(const MjSpec&, const MjVFS&)>(&mj_compile_wrapper_2)); emscripten::class_>("FloatBuffer") .constructor() diff --git a/wasm/codegen/generators/constants.py b/wasm/codegen/generators/constants.py index 8bd3d71f..0fe8b6e2 100644 --- a/wasm/codegen/generators/constants.py +++ b/wasm/codegen/generators/constants.py @@ -64,12 +64,16 @@ _SKIPPED_PLUGIN_FUNCTIONS: tuple[str, ...] = ( # Functions that are bound as class methods _SKIPPED_CLASS_METHODS: tuple[str, ...] = ( # go/keep-sorted start + "mj_addBufferVFS", "mj_copyData", "mj_copyModel", "mj_copySpec", + "mj_defaultVFS", "mj_deleteData", + "mj_deleteFileVFS", "mj_deleteModel", "mj_deleteSpec", + "mj_deleteVFS", "mj_loadXML", "mj_makeData", "mj_makeSpec", @@ -113,11 +117,9 @@ _SKIPPED_ASSET_CACHE_FUNCTIONS: tuple[str, ...] = ( # Omitted Virtual Filesystem (VFS) functions _SKIPPED_VFS_FUNCTIONS: tuple[str, ...] = ( # go/keep-sorted start - "mj_addBufferVFS", "mj_addFileVFS", - "mj_defaultVFS", - "mj_deleteFileVFS", - "mj_deleteVFS", + "mj_mountVFS", + "mj_unmountVFS", # go/keep-sorted end ) @@ -187,17 +189,6 @@ _SKIPPED_UTILITY_FUNCTIONS: tuple[str, ...] = ( # go/keep-sorted end ) -# Functions that require special wrappers. -# These functions are not bound automatically but are written by hand instead. -MANUAL_WRAPPER_FUNCTIONS: tuple[str, ...] = ( - # go/keep-sorted start - "mj_compile", - "mj_saveLastXML", - "mj_setLengthRange", - "mju_error", - # go/keep-sorted end -) - # List of functions that should be skipped during the code generation process. SKIPPED_FUNCTIONS: tuple[str, ...] = ( _SKIPPED_CLASS_METHODS @@ -232,12 +223,24 @@ SKIPPED_STRUCTS: tuple[str, ...] = ( # go/keep-sorted end ) +# Functions that require special wrappers. +# These functions are not bound automatically but are written by hand instead. +MANUAL_WRAPPER_FUNCTIONS: tuple[str, ...] = ( + # go/keep-sorted start + "mj_compile", + "mj_saveLastXML", + "mj_setLengthRange", + "mju_error", + # go/keep-sorted end +) + # Structs for which header generation is done manually. # mjvScene is included here because buffer sizes need to be calculated based on # introspect doc strings, which was considered a brittle unreliable solution in # the past. MANUAL_STRUCTS_HEADERS: tuple[str, ...] = ( "mjvScene", + "mjVFS", ) # Structs for which source code generation is done manually. MANUAL_STRUCTS_SOURCES: tuple[str, ...] = ( diff --git a/wasm/codegen/generators/functions.py b/wasm/codegen/generators/functions.py index 49104568..bc0f73d1 100644 --- a/wasm/codegen/generators/functions.py +++ b/wasm/codegen/generators/functions.py @@ -272,8 +272,9 @@ def generate( function_bindings = [] for func in sorted(functions, key=lambda f: f.name): - w = common.wrapped_function_name(func) - function_bindings.append(f'function("{func.name}", &{w});') + if func.name not in constants.MANUAL_WRAPPER_FUNCTIONS: + w = common.wrapped_function_name(func) + function_bindings.append(f'function("{func.name}", &{w});') bindings_content = "\n".join(function_bindings) return [ diff --git a/wasm/codegen/templates/bindings.cc b/wasm/codegen/templates/bindings.cc index dfae0925..5e040fc7 100644 --- a/wasm/codegen/templates/bindings.cc +++ b/wasm/codegen/templates/bindings.cc @@ -31,9 +31,9 @@ #include #include +#include #include #include -#include #include "engine/engine_util_errmem.h" #include "wasm/unpack.h" @@ -212,6 +212,27 @@ struct MjvScene { std::vector camera; }; +struct MjVFS { + MjVFS() : ptr_(new mjVFS) { mj_defaultVFS(ptr_); } + ~MjVFS() { + mj_deleteVFS(ptr_); + } + void AddBuffer(const std::string& name, const emscripten::val& buffer) { + std::vector vec = emscripten::vecFromJSArray(buffer); + int result = mj_addBufferVFS(ptr_, name.c_str(), vec.data(), vec.size()); + if (result != 0) { + mju_error("Could not add buffer to VFS: %d", result); + } + } + void DeleteFile(const std::string& filename) { + mj_deleteFileVFS(ptr_, filename.c_str()); + } + mjVFS* get() const { return ptr_; } + + private: + mjVFS* ptr_; +}; + MjModel::MjModel(mjModel* ptr) : ptr_(ptr), opt(&ptr->opt), vis(&ptr->vis), stat(&ptr->stat) {} @@ -407,7 +428,7 @@ std::unique_ptr parseXMLString_wrapper(const std::string &xml) { return std::unique_ptr(new MjSpec(ptr)); } -std::unique_ptr mj_compile_wrapper(const MjSpec& spec) { +std::unique_ptr mj_compile_wrapper_1(const MjSpec& spec) { mjSpec* spec_ptr = spec.get(); mjModel* model = mj_compile(spec_ptr, nullptr); if (!model || mjs_isWarning(spec_ptr)) { @@ -416,6 +437,16 @@ std::unique_ptr mj_compile_wrapper(const MjSpec& spec) { return std::unique_ptr(new MjModel(model)); } +std::unique_ptr mj_compile_wrapper_2(const MjSpec& spec, const MjVFS& vfs) { + mjSpec* spec_ptr = spec.get(); + mjVFS* vfs_ptr = vfs.get(); + mjModel* model = mj_compile(spec_ptr, vfs_ptr); + if (!model || mjs_isWarning(spec_ptr)) { + mju_error("%s", mjs_getError(spec_ptr)); + } + return std::unique_ptr(new MjModel(model)); +} + void error_wrapper(const String& msg) { mju_error("%s\n", msg.as().data()); } int mj_saveLastXML_wrapper(const String& filename, const MjModel& m) { @@ -444,10 +475,20 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { // {{ STRUCTS_BINDINGS }} - // {{ FUNCTION_BINDINGS }} + emscripten::class_("MjVFS") + .constructor<>() + .function("addBuffer", &MjVFS::AddBuffer) + .function("deleteFile", &MjVFS::DeleteFile); + // {{ FUNCTION_BINDINGS }} function("parseXMLString", &parseXMLString_wrapper, take_ownership()); function("error", &error_wrapper); + function("mj_saveLastXML", &mj_saveLastXML_wrapper); + function("mj_setLengthRange", &mj_setLengthRange_wrapper); + // mj_compile is bound using two overloads to handle the optional MjVFS argument, + // as using std::optional caused memory errors due to missing copy/move constructors. + function("mj_compile", emscripten::select_overload(const MjSpec&)>(&mj_compile_wrapper_1)); + function("mj_compile", emscripten::select_overload(const MjSpec&, const MjVFS&)>(&mj_compile_wrapper_2)); emscripten::class_>("FloatBuffer") .constructor() diff --git a/wasm/tests/bindings_test.ts b/wasm/tests/bindings_test.ts index 10aac900..f0c8a7a3 100644 --- a/wasm/tests/bindings_test.ts +++ b/wasm/tests/bindings_test.ts @@ -17,7 +17,7 @@ import 'jasmine'; import {MainModule, MjContact, MjContactVec, MjData, MjLROpt, MjModel, MjOption, MjsGeom, MjSolverStat, MjSpec, MjStatistic, MjTimerStat, MjvCamera, MjvFigure, MjvGeom, MjvGLCamera, MjvLight, MjvOption, MjvPerturb, MjvScene, -MjWarningStat} from '../dist/mujoco_wasm.js'; +MjWarningStat, MjVFS} from '../dist/mujoco_wasm.js'; import loadMujoco from '../dist/mujoco_wasm.js' @@ -1803,7 +1803,7 @@ describe('MuJoCo WASM Bindings', () => { } }); - it('should compile a spec from XML string', () => { + it('should compile a spec from XML string with no assets', () => { let spec = null; let model = null; try { @@ -1826,4 +1826,50 @@ describe('MuJoCo WASM Bindings', () => { } } }); + + it('should compile a spec from XML with .obj asset', () => { + const xml = ` + + + + + + + + `; + + const cube1 = ` + v -1 -1 1 + v 1 -1 1 + v -1 1 1 + v 1 1 1 + v -1 1 -1 + v 1 1 -1 + v -1 -1 -1 + v 1 -1 -1`; + + let spec: MjSpec|null = null; + let model: MjModel|null = null; + let vfs: MjVFS|null = null; + try { + spec = mujoco.parseXMLString(xml); + assertExists(spec); + + vfs = new mujoco.MjVFS(); + vfs.addBuffer('cube.obj', new TextEncoder().encode(cube1)); + assertExists(vfs); + + model = mujoco.mj_compile(spec, vfs); + assertExists(model); + expect(model.nmesh).toBe(1); + + const meshId = + mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MESH.value, 'cube'); + expect(meshId).toBeGreaterThanOrEqual(0); + } finally { + spec?.delete(); + vfs?.delete(); + model?.delete(); + } + }); });