From d2efc0e154a27073d438eb546c81ba6556e95d30 Mon Sep 17 00:00:00 2001 From: Google DeepMind Date: Thu, 29 Jan 2026 07:54:50 -0800 Subject: [PATCH] Add object-based accessors for mjData in WASM bindings. PiperOrigin-RevId: 862736608 Change-Id: I0618e9d2af37e377bad1aead1b3fd1a12334df35 --- wasm/README.md | 17 +- wasm/codegen/generated/bindings.cc | 301 +++++++++++++++++++++++------ wasm/codegen/generators/structs.py | 30 +++ wasm/codegen/templates/bindings.cc | 274 ++++++++++++++++++++------ wasm/tests/bindings_test.ts | 230 +++++++++++++++++++++- 5 files changed, 727 insertions(+), 125 deletions(-) diff --git a/wasm/README.md b/wasm/README.md index 68570ee6..b0f9f22f 100644 --- a/wasm/README.md +++ b/wasm/README.md @@ -110,6 +110,15 @@ write your application in C++ and compile it using Emscripten, you may want to copy a subset of the `EMSCRIPTEN_BINDINGS` from `bindings.cc` into your application’s source file. +### Named Access + +The bindings support named access methods, similar to the Python bindings, +allowing convenient access to model and data elements by name or index. For +example, you can access a geometry by name using `model.geom('mygeom')` or a +joint using `data.jnt('myjoint')`. + +For more details and examples of how to use named access, please refer to the [named access tests](tests/bindings_test.ts;l=2010-2487) and [documentation](https://mujoco.readthedocs.io/en/stable/python.html#named-access). + ## Development In order to change the bindings you will need to change the [`bindings.cc`](codegen/generated/bindings.cc) @@ -180,13 +189,7 @@ welcome! These bindings are not yet complete. While the main MuJoCo APIs (`mj_step`, `mj_loadXML`, etc.) are well tested, other APIs (e.g., functions from `mjspec.h`) remain untested in real web applications (though test code for - the `mjspec` bindings does exist). One notable feature not yet supported in - the WASM bindings, which has proved very useful in the Python bindings, is - named access methods — where data distributed across multiple arrays in C can - be conveniently accessed by name, e.g., `model.geom('mygeom')` or - `data.joint('myjoint')`. Currently, this data must be accessed via the - `mj_name2id` function. Adding support for these features is a high priority, - as it affects user code written in JavaScript. + the `mjspec` bindings does exist). 2. **Improve the developer experience.** There is still work to be done to improve the developer experience when diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 175f7414..1c4a7896 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -46,6 +46,7 @@ namespace mujoco::wasm { using emscripten::enum_; using emscripten::function; using emscripten::val; +using emscripten::typed_memory_view; using emscripten::return_value_policy::reference; using emscripten::return_value_policy::take_ownership; @@ -83,6 +84,24 @@ EMSCRIPTEN_DECLARE_VAL_TYPE(StringOrNull); X_ACCESSOR( TUPLE, Tuple, mjOBJ_TUPLE, tuple, ntuple ) \ X_ACCESSOR( KEYFRAME, Keyframe, mjOBJ_KEY, key, nkey ) +// Macro to define accessors for different MuJoCo object types within mjData. +// Each line calls X_ACCESSOR with the following arguments: +// 1. NAME: The object type name in uppercase (e.g., ACTUATOR). +// 2. Name: The object type name in CamelCase (e.g., Actuator). +// 3. OBJTYPE: The corresponding mjOBJ_* enum value (e.g., mjOBJ_ACTUATOR). +// 4. field: The name of the array field in mjData (e.g., actuator). +// 5. nfield: The name of the count field in mjModel (e.g., nu). +#define MJDATA_ACCESSORS \ + X_ACCESSOR( ACTUATOR, Actuator, mjOBJ_ACTUATOR, actuator, nu ) \ + X_ACCESSOR( BODY, Body, mjOBJ_BODY, body, nbody ) \ + X_ACCESSOR( CAMERA, Camera, mjOBJ_CAMERA, cam, ncam ) \ + X_ACCESSOR( GEOM, Geom, mjOBJ_GEOM, geom, ngeom ) \ + X_ACCESSOR( JOINT, Joint, mjOBJ_JOINT, jnt, njnt ) \ + X_ACCESSOR( LIGHT, Light, mjOBJ_LIGHT, light, nlight ) \ + X_ACCESSOR( SENSOR, Sensor, mjOBJ_SENSOR, sensor, nsensor ) \ + X_ACCESSOR( SITE, Site, mjOBJ_SITE, site, nsite ) \ + X_ACCESSOR( TENDON, Tendon, mjOBJ_TENDON, tendon, ntendon ) + // Raises an error if the given val is null or undefined. // A macro is used so that the error contains the name of the variable. // TODO(matijak): Remove this when we can handle strings using UNPACK_STRING? @@ -211,40 +230,79 @@ using mjVisualScale = decltype(::mjVisual::scale); // - dim1: The fixed dimension of the array if not dynamically sized. If "1", // a single element is returned. Otherwise, it's used as the stride // for the typed memory view. -#define X(type, prefix, var, dim0, dim1) \ - auto var() const { \ - if constexpr (std::string_view(#dim0) == "nq") { \ - int start = model_->jnt_qposadr[id_]; \ - int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ - return emscripten::val(emscripten::typed_memory_view(end - start, model_->prefix##var + start)); \ - } else if constexpr (std::string_view(#dim0) == "nv") { \ - int start = model_->jnt_dofadr[id_]; \ - int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ - return emscripten::val(emscripten::typed_memory_view(end - start, model_->prefix##var + start)); \ - } else if constexpr (std::string_view(#dim0) == "nhfielddata") { \ - int start = model_->hfield_adr[id_]; \ - int count = model_->hfield_nrow[id_] * model_->hfield_ncol[id_]; \ - return emscripten::val(emscripten::typed_memory_view(count, model_->hfield_data + start)); \ - } else if constexpr (std::string_view(#dim0) == "ntexdata") { \ - int start = model_->tex_adr[id_]; \ - int count = model_->tex_height[id_] * model_->tex_width[id_] * model_->tex_nchannel[id_]; \ - return emscripten::val(emscripten::typed_memory_view(count, model_->tex_data + start)); \ - } else if constexpr (std::string_view(#dim0) == "nnumericdata") { \ - int start = model_->numeric_adr[id_]; \ - int count = model_->numeric_size[id_]; \ - return emscripten::val(emscripten::typed_memory_view(count, model_->numeric_data + start)); \ - } else if constexpr (std::string_view(#dim0) == "ntupledata") { \ - int start = model_->tuple_adr[id_]; \ - int count = model_->tuple_size[id_]; \ - return emscripten::val(emscripten::typed_memory_view(count, model_->prefix##var + start)); \ - } else { \ - if constexpr (std::string_view(#dim1) == "1") { \ - return model_->prefix##var[id_]; \ - } else { \ - return emscripten::val( \ - emscripten::typed_memory_view(dim1, model_->prefix##var + id_ * dim1)); \ - } \ - } \ +#define X(type, prefix, var, dim0, dim1) \ + emscripten::val get_##var() const { \ + if constexpr (std::string_view(#dim0) == "nq") { \ + int start = model_->jnt_qposadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ + return val(typed_memory_view(end - start, model_->prefix##var + start)); \ + } else if constexpr (std::string_view(#dim0) == "nv") { \ + int start = model_->jnt_dofadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ + return val(typed_memory_view(end - start, model_->prefix##var + start)); \ + } else if constexpr (std::string_view(#dim0) == "nhfielddata") { \ + int start = model_->hfield_adr[id_]; \ + int count = model_->hfield_nrow[id_] * model_->hfield_ncol[id_]; \ + return val(typed_memory_view(count, model_->hfield_data + start)); \ + } else if constexpr (std::string_view(#dim0) == "ntexdata") { \ + int start = model_->tex_adr[id_]; \ + int count = model_->tex_height[id_] * model_->tex_width[id_] * model_->tex_nchannel[id_]; \ + return val(typed_memory_view(count, model_->tex_data + start)); \ + } else if constexpr (std::string_view(#dim0) == "nnumericdata") { \ + int start = model_->numeric_adr[id_]; \ + int count = model_->numeric_size[id_]; \ + return val(typed_memory_view(count, model_->numeric_data + start)); \ + } else if constexpr (std::string_view(#dim0) == "ntupledata") { \ + int start = model_->tuple_adr[id_]; \ + int count = model_->tuple_size[id_]; \ + return val(typed_memory_view(count, model_->prefix##var + start)); \ + } else { \ + if constexpr (std::string_view(#dim1) == "1") { \ + return val(model_->prefix##var[id_]); \ + } else { \ + return val(typed_memory_view(dim1, model_->prefix##var + id_ * dim1)); \ + } \ + } \ + } \ + void set_##var(const emscripten::val& value) { \ + if constexpr (std::string_view(#dim0) == "nq") { \ + int start = model_->jnt_qposadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ + val(typed_memory_view(end - start, model_->prefix##var + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "nv") { \ + int start = model_->jnt_dofadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ + val(typed_memory_view(end - start, model_->prefix##var + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "nhfielddata") { \ + int start = model_->hfield_adr[id_]; \ + int count = model_->hfield_nrow[id_] * model_->hfield_ncol[id_]; \ + val(typed_memory_view(count, model_->hfield_data + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "ntexdata") { \ + int start = model_->tex_adr[id_]; \ + int count = model_->tex_height[id_] * model_->tex_width[id_] * model_->tex_nchannel[id_]; \ + val(typed_memory_view(count, model_->tex_data + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "nnumericdata") { \ + int start = model_->numeric_adr[id_]; \ + int count = model_->numeric_size[id_]; \ + val(typed_memory_view(count, model_->numeric_data + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "ntupledata") { \ + int start = model_->tuple_adr[id_]; \ + int count = model_->tuple_size[id_]; \ + val(typed_memory_view(count, model_->prefix##var + start)) \ + .call("set", value); \ + } else { \ + if constexpr (std::string_view(#dim1) == "1") { \ + model_->prefix##var[id_] = value.as(); \ + } else { \ + val(typed_memory_view(dim1, model_->prefix##var + id_ * dim1)) \ + .call("set", value); \ + } \ + } \ } // Expands to a struct definition for each object type in MJMODEL_ACCESSORS. @@ -254,24 +312,114 @@ using mjVisualScale = decltype(::mjVisual::scale); // - A `name()` method to get the object's name using `mj_id2name`. // - Member functions generated by the `MJMODEL_##NAME` macro, which in turn // uses the `X` macro to define accessors for fields within `mjModel`. -#define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ - struct MjModel##Name##Accessor { \ - MjModel##Name##Accessor(mjModel* model, int id) : model_(model), id_(id) {} \ - \ - int id() const { return id_; } \ - std::string name() const { \ - return mj_id2name(model_, OBJTYPE, id_); \ - } \ - \ - MJMODEL_##NAME \ - \ - private: \ - mjModel* model_; \ - int id_; \ +#define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ + struct MjModel##Name##Accessor { \ + MjModel##Name##Accessor(mjModel* model, int id) : model_(model), id_(id) {} \ + \ + int id() const { return id_; } \ + std::string name() const { \ + const char* name = mj_id2name(model_, OBJTYPE, id_); \ + return name ? name : ""; \ + } \ + \ + MJMODEL_##NAME \ + \ + private: \ + mjModel* model_; \ + int id_; \ }; MJMODEL_ACCESSORS #undef X_ACCESSOR +// Expands to a struct definition for each object type in MJDATA_ACCESSORS. +// Each struct, named `MjData{Name}Accessor`, provides: +// - A constructor taking an `mjData*`, an `mjModel*`, and an integer `id`. +// - An `id()` method to get the object's index. +// - A `name()` method to get the object's name using `mj_id2name`. +// - Member functions generated by the `MJDATA_##NAME` macro, which in turn +// uses the `X` macro to define accessors for fields within `mjData`. +#undef MJ_M +#define MJ_M(n) model_->n +#undef X +#define X(type, prefix, var, dim0, dim1) \ + emscripten::val get_##var() const { \ + if constexpr (std::string_view(#dim0) == "nq") { \ + int start = model_->jnt_qposadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ + return val(typed_memory_view(end - start, data_->prefix##var + start)); \ + } else if constexpr (std::string_view(#dim0) == "nv") { \ + int start = model_->jnt_dofadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ + if constexpr (std::string_view(#dim1) == "1") { \ + return val(typed_memory_view(end - start, data_->prefix##var + start)); \ + } else { \ + return val(typed_memory_view( \ + (end - start) * dim1, data_->prefix##var + start * dim1)); \ + } \ + } else if constexpr (std::string_view(#dim0) == "nsensordata") { \ + int start = model_->sensor_adr[id_]; \ + int count = model_->sensor_dim[id_]; \ + return val(typed_memory_view(count, data_->sensordata + start)); \ + } else { \ + if constexpr (std::string_view(#dim1) == "1") { \ + return val(data_->prefix##var[id_]); \ + } else { \ + return val(typed_memory_view(dim1, data_->prefix##var + id_ * dim1)); \ + } \ + } \ + } \ + void set_##var(const emscripten::val& value) { \ + if constexpr (std::string_view(#dim0) == "nq") { \ + int start = model_->jnt_qposadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ + val(typed_memory_view(end - start, data_->prefix##var + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "nv") { \ + int start = model_->jnt_dofadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ + if constexpr (std::string_view(#dim1) == "1") { \ + val(typed_memory_view(end - start, data_->prefix##var + start)) \ + .call("set", value); \ + } else { \ + val(typed_memory_view((end - start) * dim1, \ + data_->prefix##var + start * dim1)) \ + .call("set", value); \ + } \ + } else if constexpr (std::string_view(#dim0) == "nsensordata") { \ + int start = model_->sensor_adr[id_]; \ + int count = model_->sensor_dim[id_]; \ + val(typed_memory_view(count, data_->sensordata + start)) \ + .call("set", value); \ + } else { \ + if constexpr (std::string_view(#dim1) == "1") { \ + data_->prefix##var[id_] = value.as(); \ + } else { \ + val(typed_memory_view(dim1, data_->prefix##var + id_ * dim1)) \ + .call("set", value); \ + } \ + } \ + } + +#define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ + struct MjData##Name##Accessor { \ + MjData##Name##Accessor(mjData* data, mjModel* model, int id) : data_(data), model_(model), id_(id) {} \ + \ + int id() const { return id_; } \ + std::string name() const { \ + const char* name = mj_id2name(model_, OBJTYPE, id_); \ + return name ? name : ""; \ + } \ + \ + MJDATA_##NAME \ + \ + private: \ + mjData* data_; \ + mjModel* model_; \ + int id_; \ + }; +MJDATA_ACCESSORS +#undef X_ACCESSOR + #undef X #undef MJ_M @@ -6778,6 +6926,28 @@ struct MjData { void set_signature(uint64_t value) { ptr_->signature = value; } + // Generates functions to return accessor classes. + #define X_ACCESSOR(NAME, Name, OBJTYPE, accessor_name, nfield) \ + MjData##Name##Accessor accessor_name(const NumberOrString& val) const { \ + if (val.isString()) { \ + int id = mj_name2id(model, OBJTYPE, val.as().c_str()); \ + if (id == -1) { \ + mju_error("%s", KeyErrorMessage(model, OBJTYPE, model->nfield, val.as(), #accessor_name).c_str()); \ + } \ + return MjData##Name##Accessor(ptr_, model, id); \ + } else if (val.isNumber()) { \ + int id = val.as(); \ + if (id < 0 || id >= model->nfield) { \ + mju_error("%s", IndexErrorMessage(id, model->nfield, #accessor_name).c_str()); \ + } \ + return MjData##Name##Accessor(ptr_, model, id); \ + } else { \ + mju_error(#accessor_name "() argument must be a string or number"); \ + return MjData##Name##Accessor(nullptr, nullptr, 0); \ + } \ + } + MJDATA_ACCESSORS + #undef X_ACCESSOR private: mjData* ptr_; @@ -11048,19 +11218,33 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .value("mjWRAP_CYLINDER", mjWRAP_CYLINDER); // Bindings for the MjModel accessor classes. - #define X(type, prefix, var, dim0, dim1) .property(#var, &Accessor::var) - #define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ - { \ - using Accessor = MjModel##Name##Accessor; \ - emscripten::class_("MjModel" #Name "Accessor") \ - .property("id", &Accessor::id) \ - .property("name", &Accessor::name) \ - MJMODEL_##NAME; \ + #define X(type, prefix, var, dim0, dim1) .property(#var, &Accessor::get_##var, &Accessor::set_##var) + #define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ + { \ + using Accessor = MjModel##Name##Accessor; \ + emscripten::class_("MjModel" #Name "Accessor") \ + .property("id", &Accessor::id) \ + .property("name", &Accessor::name) \ + MJMODEL_##NAME; \ } MJMODEL_ACCESSORS #undef X #undef X_ACCESSOR + // Bindings for the MjData accessor classes. + #define X(type, prefix, var, dim0, dim1) .property(#var, &Accessor::get_##var, &Accessor::set_##var) + #define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ + { \ + using Accessor = MjData##Name##Accessor; \ + emscripten::class_("MjData" #Name "Accessor") \ + .property("id", &Accessor::id) \ + .property("name", &Accessor::name) \ + MJDATA_##NAME; \ + } + MJDATA_ACCESSORS + #undef X + #undef X_ACCESSOR + emscripten::class_("MjContact") .constructor<>() .function("copy", &MjContact::copy, take_ownership()) @@ -11086,6 +11270,11 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { emscripten::class_("MjData") .constructor() .constructor() + // Binds the functions on MjData that return accessors. + #define X_ACCESSOR(NAME, Name, OBJTYPE, field_name, nfield) \ + .function(#field_name, &MjData::field_name) + MJDATA_ACCESSORS + #undef X_ACCESSOR .property("M", &MjData::M) .property("act", &MjData::act) .property("act_dot", &MjData::act_dot) diff --git a/wasm/codegen/generators/structs.py b/wasm/codegen/generators/structs.py index 1dae82cf..f6282693 100644 --- a/wasm/codegen/generators/structs.py +++ b/wasm/codegen/generators/structs.py @@ -413,6 +413,30 @@ def build_struct_header( MJMODEL_ACCESSORS #undef X_ACCESSOR""".lstrip()) + elif w == "MjData": + builder.line(""" + // Generates functions to return accessor classes. + #define X_ACCESSOR(NAME, Name, OBJTYPE, accessor_name, nfield) \\ + MjData##Name##Accessor accessor_name(const NumberOrString& val) const { \\ + if (val.isString()) { \\ + int id = mj_name2id(model, OBJTYPE, val.as().c_str()); \\ + if (id == -1) { \\ + mju_error("%s", KeyErrorMessage(model, OBJTYPE, model->nfield, val.as(), #accessor_name).c_str()); \\ + } \\ + return MjData##Name##Accessor(ptr_, model, id); \\ + } else if (val.isNumber()) { \\ + int id = val.as(); \\ + if (id < 0 || id >= model->nfield) { \\ + mju_error("%s", IndexErrorMessage(id, model->nfield, #accessor_name).c_str()); \\ + } \\ + return MjData##Name##Accessor(ptr_, model, id); \\ + } else { \\ + mju_error(#accessor_name "() argument must be a string or number"); \\ + return MjData##Name##Accessor(nullptr, nullptr, 0); \\ + } \\ + } + MJDATA_ACCESSORS + #undef X_ACCESSOR""".lstrip()) # define private struct members builder.private() @@ -518,6 +542,12 @@ def _build_struct_bindings( if w == "MjData": builder.line(".constructor()") builder.line(".constructor()") + builder.line(""" + // Binds the functions on MjData that return accessors. + #define X_ACCESSOR(NAME, Name, OBJTYPE, field_name, nfield) \\ + .function(#field_name, &MjData::field_name) + MJDATA_ACCESSORS + #undef X_ACCESSOR""".lstrip()) elif w == "MjModel": f1 = common.wrapped_function_name( introspect_functions.FUNCTIONS["mj_loadXML"] diff --git a/wasm/codegen/templates/bindings.cc b/wasm/codegen/templates/bindings.cc index dadcdbc2..8e29a080 100644 --- a/wasm/codegen/templates/bindings.cc +++ b/wasm/codegen/templates/bindings.cc @@ -46,6 +46,7 @@ namespace mujoco::wasm { using emscripten::enum_; using emscripten::function; using emscripten::val; +using emscripten::typed_memory_view; using emscripten::return_value_policy::reference; using emscripten::return_value_policy::take_ownership; @@ -83,6 +84,24 @@ EMSCRIPTEN_DECLARE_VAL_TYPE(StringOrNull); X_ACCESSOR( TUPLE, Tuple, mjOBJ_TUPLE, tuple, ntuple ) \ X_ACCESSOR( KEYFRAME, Keyframe, mjOBJ_KEY, key, nkey ) +// Macro to define accessors for different MuJoCo object types within mjData. +// Each line calls X_ACCESSOR with the following arguments: +// 1. NAME: The object type name in uppercase (e.g., ACTUATOR). +// 2. Name: The object type name in CamelCase (e.g., Actuator). +// 3. OBJTYPE: The corresponding mjOBJ_* enum value (e.g., mjOBJ_ACTUATOR). +// 4. field: The name of the array field in mjData (e.g., actuator). +// 5. nfield: The name of the count field in mjModel (e.g., nu). +#define MJDATA_ACCESSORS \ + X_ACCESSOR( ACTUATOR, Actuator, mjOBJ_ACTUATOR, actuator, nu ) \ + X_ACCESSOR( BODY, Body, mjOBJ_BODY, body, nbody ) \ + X_ACCESSOR( CAMERA, Camera, mjOBJ_CAMERA, cam, ncam ) \ + X_ACCESSOR( GEOM, Geom, mjOBJ_GEOM, geom, ngeom ) \ + X_ACCESSOR( JOINT, Joint, mjOBJ_JOINT, jnt, njnt ) \ + X_ACCESSOR( LIGHT, Light, mjOBJ_LIGHT, light, nlight ) \ + X_ACCESSOR( SENSOR, Sensor, mjOBJ_SENSOR, sensor, nsensor ) \ + X_ACCESSOR( SITE, Site, mjOBJ_SITE, site, nsite ) \ + X_ACCESSOR( TENDON, Tendon, mjOBJ_TENDON, tendon, ntendon ) + // Raises an error if the given val is null or undefined. // A macro is used so that the error contains the name of the variable. // TODO(matijak): Remove this when we can handle strings using UNPACK_STRING? @@ -206,40 +225,79 @@ val get_mjRNDSTRING() { return MakeValArray3(mjRNDSTRING); } // - dim1: The fixed dimension of the array if not dynamically sized. If "1", // a single element is returned. Otherwise, it's used as the stride // for the typed memory view. -#define X(type, prefix, var, dim0, dim1) \ - auto var() const { \ - if constexpr (std::string_view(#dim0) == "nq") { \ - int start = model_->jnt_qposadr[id_]; \ - int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ - return emscripten::val(emscripten::typed_memory_view(end - start, model_->prefix##var + start)); \ - } else if constexpr (std::string_view(#dim0) == "nv") { \ - int start = model_->jnt_dofadr[id_]; \ - int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ - return emscripten::val(emscripten::typed_memory_view(end - start, model_->prefix##var + start)); \ - } else if constexpr (std::string_view(#dim0) == "nhfielddata") { \ - int start = model_->hfield_adr[id_]; \ - int count = model_->hfield_nrow[id_] * model_->hfield_ncol[id_]; \ - return emscripten::val(emscripten::typed_memory_view(count, model_->hfield_data + start)); \ - } else if constexpr (std::string_view(#dim0) == "ntexdata") { \ - int start = model_->tex_adr[id_]; \ - int count = model_->tex_height[id_] * model_->tex_width[id_] * model_->tex_nchannel[id_]; \ - return emscripten::val(emscripten::typed_memory_view(count, model_->tex_data + start)); \ - } else if constexpr (std::string_view(#dim0) == "nnumericdata") { \ - int start = model_->numeric_adr[id_]; \ - int count = model_->numeric_size[id_]; \ - return emscripten::val(emscripten::typed_memory_view(count, model_->numeric_data + start)); \ - } else if constexpr (std::string_view(#dim0) == "ntupledata") { \ - int start = model_->tuple_adr[id_]; \ - int count = model_->tuple_size[id_]; \ - return emscripten::val(emscripten::typed_memory_view(count, model_->prefix##var + start)); \ - } else { \ - if constexpr (std::string_view(#dim1) == "1") { \ - return model_->prefix##var[id_]; \ - } else { \ - return emscripten::val( \ - emscripten::typed_memory_view(dim1, model_->prefix##var + id_ * dim1)); \ - } \ - } \ +#define X(type, prefix, var, dim0, dim1) \ + emscripten::val get_##var() const { \ + if constexpr (std::string_view(#dim0) == "nq") { \ + int start = model_->jnt_qposadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ + return val(typed_memory_view(end - start, model_->prefix##var + start)); \ + } else if constexpr (std::string_view(#dim0) == "nv") { \ + int start = model_->jnt_dofadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ + return val(typed_memory_view(end - start, model_->prefix##var + start)); \ + } else if constexpr (std::string_view(#dim0) == "nhfielddata") { \ + int start = model_->hfield_adr[id_]; \ + int count = model_->hfield_nrow[id_] * model_->hfield_ncol[id_]; \ + return val(typed_memory_view(count, model_->hfield_data + start)); \ + } else if constexpr (std::string_view(#dim0) == "ntexdata") { \ + int start = model_->tex_adr[id_]; \ + int count = model_->tex_height[id_] * model_->tex_width[id_] * model_->tex_nchannel[id_]; \ + return val(typed_memory_view(count, model_->tex_data + start)); \ + } else if constexpr (std::string_view(#dim0) == "nnumericdata") { \ + int start = model_->numeric_adr[id_]; \ + int count = model_->numeric_size[id_]; \ + return val(typed_memory_view(count, model_->numeric_data + start)); \ + } else if constexpr (std::string_view(#dim0) == "ntupledata") { \ + int start = model_->tuple_adr[id_]; \ + int count = model_->tuple_size[id_]; \ + return val(typed_memory_view(count, model_->prefix##var + start)); \ + } else { \ + if constexpr (std::string_view(#dim1) == "1") { \ + return val(model_->prefix##var[id_]); \ + } else { \ + return val(typed_memory_view(dim1, model_->prefix##var + id_ * dim1)); \ + } \ + } \ + } \ + void set_##var(const emscripten::val& value) { \ + if constexpr (std::string_view(#dim0) == "nq") { \ + int start = model_->jnt_qposadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ + val(typed_memory_view(end - start, model_->prefix##var + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "nv") { \ + int start = model_->jnt_dofadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ + val(typed_memory_view(end - start, model_->prefix##var + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "nhfielddata") { \ + int start = model_->hfield_adr[id_]; \ + int count = model_->hfield_nrow[id_] * model_->hfield_ncol[id_]; \ + val(typed_memory_view(count, model_->hfield_data + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "ntexdata") { \ + int start = model_->tex_adr[id_]; \ + int count = model_->tex_height[id_] * model_->tex_width[id_] * model_->tex_nchannel[id_]; \ + val(typed_memory_view(count, model_->tex_data + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "nnumericdata") { \ + int start = model_->numeric_adr[id_]; \ + int count = model_->numeric_size[id_]; \ + val(typed_memory_view(count, model_->numeric_data + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "ntupledata") { \ + int start = model_->tuple_adr[id_]; \ + int count = model_->tuple_size[id_]; \ + val(typed_memory_view(count, model_->prefix##var + start)) \ + .call("set", value); \ + } else { \ + if constexpr (std::string_view(#dim1) == "1") { \ + model_->prefix##var[id_] = value.as(); \ + } else { \ + val(typed_memory_view(dim1, model_->prefix##var + id_ * dim1)) \ + .call("set", value); \ + } \ + } \ } // Expands to a struct definition for each object type in MJMODEL_ACCESSORS. @@ -249,24 +307,114 @@ val get_mjRNDSTRING() { return MakeValArray3(mjRNDSTRING); } // - A `name()` method to get the object's name using `mj_id2name`. // - Member functions generated by the `MJMODEL_##NAME` macro, which in turn // uses the `X` macro to define accessors for fields within `mjModel`. -#define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ - struct MjModel##Name##Accessor { \ - MjModel##Name##Accessor(mjModel* model, int id) : model_(model), id_(id) {} \ - \ - int id() const { return id_; } \ - std::string name() const { \ - return mj_id2name(model_, OBJTYPE, id_); \ - } \ - \ - MJMODEL_##NAME \ - \ - private: \ - mjModel* model_; \ - int id_; \ +#define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ + struct MjModel##Name##Accessor { \ + MjModel##Name##Accessor(mjModel* model, int id) : model_(model), id_(id) {} \ + \ + int id() const { return id_; } \ + std::string name() const { \ + const char* name = mj_id2name(model_, OBJTYPE, id_); \ + return name ? name : ""; \ + } \ + \ + MJMODEL_##NAME \ + \ + private: \ + mjModel* model_; \ + int id_; \ }; MJMODEL_ACCESSORS #undef X_ACCESSOR +// Expands to a struct definition for each object type in MJDATA_ACCESSORS. +// Each struct, named `MjData{Name}Accessor`, provides: +// - A constructor taking an `mjData*`, an `mjModel*`, and an integer `id`. +// - An `id()` method to get the object's index. +// - A `name()` method to get the object's name using `mj_id2name`. +// - Member functions generated by the `MJDATA_##NAME` macro, which in turn +// uses the `X` macro to define accessors for fields within `mjData`. +#undef MJ_M +#define MJ_M(n) model_->n +#undef X +#define X(type, prefix, var, dim0, dim1) \ + emscripten::val get_##var() const { \ + if constexpr (std::string_view(#dim0) == "nq") { \ + int start = model_->jnt_qposadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ + return val(typed_memory_view(end - start, data_->prefix##var + start)); \ + } else if constexpr (std::string_view(#dim0) == "nv") { \ + int start = model_->jnt_dofadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ + if constexpr (std::string_view(#dim1) == "1") { \ + return val(typed_memory_view(end - start, data_->prefix##var + start)); \ + } else { \ + return val(typed_memory_view( \ + (end - start) * dim1, data_->prefix##var + start * dim1)); \ + } \ + } else if constexpr (std::string_view(#dim0) == "nsensordata") { \ + int start = model_->sensor_adr[id_]; \ + int count = model_->sensor_dim[id_]; \ + return val(typed_memory_view(count, data_->sensordata + start)); \ + } else { \ + if constexpr (std::string_view(#dim1) == "1") { \ + return val(data_->prefix##var[id_]); \ + } else { \ + return val(typed_memory_view(dim1, data_->prefix##var + id_ * dim1)); \ + } \ + } \ + } \ + void set_##var(const emscripten::val& value) { \ + if constexpr (std::string_view(#dim0) == "nq") { \ + int start = model_->jnt_qposadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_qposadr[id_ + 1] : model_->nq; \ + val(typed_memory_view(end - start, data_->prefix##var + start)) \ + .call("set", value); \ + } else if constexpr (std::string_view(#dim0) == "nv") { \ + int start = model_->jnt_dofadr[id_]; \ + int end = (id_ < model_->njnt - 1) ? model_->jnt_dofadr[id_ + 1] : model_->nv; \ + if constexpr (std::string_view(#dim1) == "1") { \ + val(typed_memory_view(end - start, data_->prefix##var + start)) \ + .call("set", value); \ + } else { \ + val(typed_memory_view((end - start) * dim1, \ + data_->prefix##var + start * dim1)) \ + .call("set", value); \ + } \ + } else if constexpr (std::string_view(#dim0) == "nsensordata") { \ + int start = model_->sensor_adr[id_]; \ + int count = model_->sensor_dim[id_]; \ + val(typed_memory_view(count, data_->sensordata + start)) \ + .call("set", value); \ + } else { \ + if constexpr (std::string_view(#dim1) == "1") { \ + data_->prefix##var[id_] = value.as(); \ + } else { \ + val(typed_memory_view(dim1, data_->prefix##var + id_ * dim1)) \ + .call("set", value); \ + } \ + } \ + } + +#define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ + struct MjData##Name##Accessor { \ + MjData##Name##Accessor(mjData* data, mjModel* model, int id) : data_(data), model_(model), id_(id) {} \ + \ + int id() const { return id_; } \ + std::string name() const { \ + const char* name = mj_id2name(model_, OBJTYPE, id_); \ + return name ? name : ""; \ + } \ + \ + MJDATA_##NAME \ + \ + private: \ + mjData* data_; \ + mjModel* model_; \ + int id_; \ + }; +MJDATA_ACCESSORS +#undef X_ACCESSOR + #undef X #undef MJ_M @@ -644,19 +792,33 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { // {{ ENUM_BINDINGS }} // Bindings for the MjModel accessor classes. - #define X(type, prefix, var, dim0, dim1) .property(#var, &Accessor::var) - #define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ - { \ - using Accessor = MjModel##Name##Accessor; \ - emscripten::class_("MjModel" #Name "Accessor") \ - .property("id", &Accessor::id) \ - .property("name", &Accessor::name) \ - MJMODEL_##NAME; \ + #define X(type, prefix, var, dim0, dim1) .property(#var, &Accessor::get_##var, &Accessor::set_##var) + #define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ + { \ + using Accessor = MjModel##Name##Accessor; \ + emscripten::class_("MjModel" #Name "Accessor") \ + .property("id", &Accessor::id) \ + .property("name", &Accessor::name) \ + MJMODEL_##NAME; \ } MJMODEL_ACCESSORS #undef X #undef X_ACCESSOR + // Bindings for the MjData accessor classes. + #define X(type, prefix, var, dim0, dim1) .property(#var, &Accessor::get_##var, &Accessor::set_##var) + #define X_ACCESSOR(NAME, Name, OBJTYPE, field, nfield) \ + { \ + using Accessor = MjData##Name##Accessor; \ + emscripten::class_("MjData" #Name "Accessor") \ + .property("id", &Accessor::id) \ + .property("name", &Accessor::name) \ + MJDATA_##NAME; \ + } + MJDATA_ACCESSORS + #undef X + #undef X_ACCESSOR + // {{ STRUCTS_BINDINGS }} emscripten::class_("MjVFS") diff --git a/wasm/tests/bindings_test.ts b/wasm/tests/bindings_test.ts index 3b524964..a7f30eb9 100644 --- a/wasm/tests/bindings_test.ts +++ b/wasm/tests/bindings_test.ts @@ -45,17 +45,17 @@ const TEST_XML = ` - + - + - - + + @@ -900,7 +900,7 @@ describe('MuJoCo WASM Bindings', () => { // Corresponds to bindings_test.py:test_can_read_array it('should read an array from the model', () => { const expected = - new Float64Array([0, 0, 0, 0, 0, 0.1, 0, 0, 0, 0, 0, 0, 42, 0, 42]); + new Float64Array([0, 0, 0, 0, 0, 0.1, 0, 0, 1, 2, 0, 1, 42, 0, 42]); const bodyPos = new Float64Array(model!.body_pos); expectArraysEqual(bodyPos, expected); }); @@ -1881,7 +1881,7 @@ describe('MuJoCo WASM Bindings', () => { model!.geom('badgeom'); }) .toThrowError( - `MuJoCo Error: Invalid name 'badgeom' for geom. Valid names: \['mybox', 'myplane'\]`); + `MuJoCo Error: Invalid name 'badgeom' for geom. Valid names: ['mocap-sphere', 'mybox', 'myplane']`); }); // Corresponds to @@ -2160,6 +2160,223 @@ describe('MuJoCo WASM Bindings', () => { }); }); + describe('MjData named access', () => { + it('should support named access for MjData', () => { + mujoco.mj_forward(model!, data!); + + // Actuator + expect(model!.nu).toBe(1); + const actuator = data!.actuator(0); + expect(actuator.name).toBe('myactuator'); + expect(actuator.velocity).toBe(0); + expect(actuator.force).toBe(0); + expectArraysClose( + actuator.moment, new Float64Array([1, 0, 0, 0, 0, 0, 0, 0, 0, 0])); + + // Body + expect(model!.nbody).toBe(5); + expect(data!.body('mybox').name).toBe('mybox'); + expectArraysClose( + data!.body('mybox').xpos, new Float64Array([0, 0, 0.1])); + expect(data!.body('myhinge-body').name).toBe('myhinge-body'); + expectArraysClose( + data!.body('myhinge-body').xpos, new Float64Array([0, 0, 1])); + expect(data!.body('myball-body').name).toBe('myball-body'); + expectArraysClose( + data!.body('myball-body').xpos, new Float64Array([2, 0, 1])); + expect(data!.body('mocap-body').name).toBe('mocap-body'); + expectArraysClose( + data!.body('mocap-body').xpos, new Float64Array([42, 0, 42])); + + // Geom + expect(model!.ngeom).toBe(3); + expect(data!.geom('myplane').name).toBe('myplane'); + expectArraysClose( + data!.geom('myplane').xpos, new Float64Array([0, 0, 0])); + expectArraysClose( + data!.geom('myplane').xmat, + new Float64Array([1, 0, 0, 0, 1, 0, 0, 0, 1])); + + expect(data!.geom('mybox').name).toBe('mybox'); + expectArraysClose( + data!.geom('mybox').xpos, new Float64Array([0, 0, 0.1])); + expectArraysClose( + data!.geom('mybox').xmat, + new Float64Array([1, 0, 0, 0, 1, 0, 0, 0, 1])); + + expect(data!.geom('mocap-sphere').name).toBe('mocap-sphere'); + expectArraysClose( + data!.geom('mocap-sphere').xpos, new Float64Array([42, 0, 42])); + expectArraysClose( + data!.geom('mocap-sphere').xmat, + new Float64Array([1, 0, 0, 0, 1, 0, 0, 0, 1])); + + // Joint + expect(model!.njnt).toBe(3); + const freeJnt = data!.jnt('myfree'); + expect(freeJnt.name).toBe('myfree'); + expectArraysClose(freeJnt.xanchor, new Float64Array([0, 0, 0.1])); + expectArraysClose(freeJnt.xaxis, new Float64Array([0, 0, 1])); + + const hingeJnt = data!.jnt('myhinge'); + expect(hingeJnt.name).toBe('myhinge'); + expectArraysClose(hingeJnt.xanchor, new Float64Array([0, 0, 1])); + expectArraysClose(hingeJnt.xaxis, new Float64Array([0, 1, 0])); + + const ballJnt = data!.jnt('myball'); + expect(ballJnt.name).toBe('myball'); + expectArraysClose(ballJnt.xanchor, new Float64Array([2, 0, 1])); + expectArraysClose(ballJnt.xaxis, new Float64Array([0, 0, 1])); + + // Sensor + expect(model!.nsensor).toBe(2); + const jointvelSensor = data!.sensor('myjointvel'); + expect(jointvelSensor.name).toBe('myjointvel'); + expectArraysClose(jointvelSensor.data, new Float64Array([0])); + + const accelSensor = data!.sensor('myaccelerometer'); + expect(accelSensor.name).toBe('myaccelerometer'); + expectArraysClose(accelSensor.data, new Float64Array([0, 0, 9.81])); + + // Site + expect(model!.nsite).toBe(1); + const site = data!.site('mysite'); + expect(site.name).toBe('mysite'); + expectArraysClose(site.xpos, new Float64Array([0, 0, 0])); + expectArraysClose( + site.xmat, new Float64Array([1, 0, 0, 0, 1, 0, 0, 0, 1])); + }); + + // Corresponds to bindings_test.py:test_indexer_name_id + it('should support named and indexed access for MjData geoms', () => { + const xml = ` + + + + + + + + + + `; + const tempXmlFilename = '/tmp/geom_idx.xml'; + writeXMLFile(tempXmlFilename, xml); + const model = mujoco.MjModel.mj_loadXML(tempXmlFilename); + const data = new mujoco.MjData(model); + try { + assertExists(model); + assertExists(data); + mujoco.mj_forward(model, data); + + expect(data.geom('mygeom').id).toBe(0); + expect(data.geom('myothergeom').id).toBe(3); + expect(data.geom(0).name).toBe('mygeom'); + expect(data.geom(1).name).toBe(''); + expect(data.geom(2).name).toBe(''); + expect(data.geom(3).name).toBe('myothergeom'); + expect(data.geom(4).name).toBe(''); + + expect(data.geom(0).xpos[2]).toBeCloseTo(1); + expect(data.geom(1).xpos[2]).toBeCloseTo(2); + expect(data.geom(2).xpos[2]).toBeCloseTo(3); + expect(data.geom(3).xpos[2]).toBeCloseTo(4); + expect(data.geom(4).xpos[2]).toBeCloseTo(5); + } finally { + model?.delete(); + data?.delete(); + unlinkXMLFile(tempXmlFilename); + } + }); + + // Corresponds to bindings_test.py:test_named_indexing_invalid_names_in_data + it('should throw an error for invalid geom names in data', () => { + expect(() => { + data!.geom('badgeom'); + }) + .toThrowError( + `MuJoCo Error: Invalid name 'badgeom' for geom. Valid names: ['mocap-sphere', 'mybox', 'myplane']`); + }); + + // Corresponds to bindings_test.py:test_named_indexing_invalid_index_in_data + it('should throw an error for invalid geom indices in data', () => { + expect(() => { + data!.geom(3); + }) + .toThrowError( + 'MuJoCo Error: Invalid index 3 for geom. Valid indices from 0 to 2'); + expect(() => { + data!.geom(-1); + }) + .toThrowError( + 'MuJoCo Error: Invalid index -1 for geom. Valid indices from 0 to 2'); + }); + + // Corresponds to bindings_test.py:test_named_indexing_actuator_ctrl + it('should support named access for Actuator ctrl', () => { + const actuatorId = mujoco.mj_name2id( + model!, mujoco.mjtObj.mjOBJ_ACTUATOR.value, 'myactuator'); + assertExists(actuatorId); + + const actuator = data!.actuator('myactuator'); + expect(actuator).toEqual(data!.actuator(actuatorId)); + expect(actuator.ctrl).toEqual(data!.actuator(actuatorId)!.ctrl); + + // Test that the indexer is returning a view into the underlying + data!.ctrl[actuatorId] = 5; + expect(actuator.ctrl).toBe(5); + actuator.ctrl = 7; + expect(data!.ctrl[actuatorId]).toBe(7); + }); + + // Corresponds to bindings_test.py:test_named_indexing_ragged_qpos + it('should support named access for Joint qpos', () => { + const balljointId = + mujoco.mj_name2id(model!, mujoco.mjtObj.mjOBJ_JOINT.value, 'myball'); + assertExists(balljointId); + + const ballJnt = data!.jnt('myball'); + expect(ballJnt).toEqual(data!.jnt(balljointId)); + expect(ballJnt.qpos).toEqual(data!.jnt(balljointId)!.qpos); + + // Test that the indexer is returning a view into the underlying struct. + const qposFromIndexer = ballJnt.qpos; + const qposIdx = model!.jnt_qposadr[balljointId]; + data!.qpos.set([4, 5, 6, 7], qposIdx); + expectArraysEqual(qposFromIndexer, new Float64Array([4, 5, 6, 7])); + ballJnt.qpos.set(new Float64Array([9, 8, 7, 6])) + expectArraysEqual( + data!.qpos.slice(qposIdx, qposIdx + 4), + new Float64Array([9, 8, 7, 6])); + }); + + // Corresponds to bindings_test.py:test_named_indexing_ragged2d_cdof + it('should support named access for Joint cdof', () => { + const freejointId = + mujoco.mj_name2id(model!, mujoco.mjtObj.mjOBJ_JOINT.value, 'myfree'); + assertExists(freejointId); + mujoco.mj_forward(model!, data!); + + const freeJnt = data!.jnt('myfree'); + expect(freeJnt).toEqual(data!.jnt(freejointId)); + expect(freeJnt.cdof).toEqual(data!.jnt(freejointId)!.cdof); + expect(freeJnt.cdof.length).toBe(36); // 6x6 + + // Test that the indexer is returning a view into the underlying + const cdofFromIndexer = freeJnt.cdof; + const dofIdx = model!.jnt_dofadr[freejointId]; + const testArray = new Float64Array(36); + for (let i = 0; i < 36; i++) { + testArray[i] = i; + } + data!.cdof.set(testArray); + expectArraysEqual(freeJnt.cdof, testArray); + const expectedCdof = new Float64Array(36).fill(42); + freeJnt.cdof.set(expectedCdof); + expectArraysEqual(data!.cdof.slice(dofIdx, dofIdx + 36), expectedCdof); + }); + }); + it('should save model to buffer', () => { assertExists(model); const modelSize = mujoco.mj_sizeModel(model); @@ -2279,4 +2496,5 @@ describe('MuJoCo WASM Bindings', () => { unlinkXMLFile(mjbFilename); } }); + });