Create mjcPhysics ImageableAPI for visual, non physical visual entities.

PiperOrigin-RevId: 786366034
Change-Id: Ifb72d2daf8f5963bb7007dda1685da94b456ff41

src/experimental/usd/CMakeLists.txt

@@ -174,6 +174,7 @@ add_library(${MJC_PHYSICS_PLUGIN_TARGET_NAME} SHARED)
 target_sources(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PRIVATE
   mjcPhysics/actuator.cpp
   mjcPhysics/collisionAPI.cpp
+  mjcPhysics/imageableAPI.cpp
   mjcPhysics/jointAPI.cpp
   mjcPhysics/keyframe.cpp
   mjcPhysics/materialAPI.cpp

src/experimental/usd/mjcPhysics/generatedSchema.usda

@@ -293,6 +293,16 @@ class "MjcSiteAPI" (
     )
 }
 
+class "MjcImageableAPI" (
+    doc = "API describing attributes for visual entities in MuJoCo."
+)
+{
+    uniform int mjc:group = 0 (
+        displayName = "Group"
+        doc = "Integer MuJoCo group to which the imageable belongs."
+    )
+}
+
 class "MjcCollisionAPI" (
     doc = "API describing a MuJoCo collider."
 )

src/experimental/usd/mjcPhysics/imageableAPI.cpp

@@ -0,0 +1,124 @@
+// Copyright 2025 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
+//
+//     http://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.
+
+#include <mujoco/experimental/usd/mjcPhysics/imageableAPI.h>
+
+#include <pxr/usd/sdf/assetPath.h>
+#include <pxr/usd/sdf/types.h>
+#include <pxr/usd/usd/schemaRegistry.h>
+#include <pxr/usd/usd/typed.h>
+
+PXR_NAMESPACE_OPEN_SCOPE
+
+// Register the schema with the TfType system.
+TF_REGISTRY_FUNCTION(TfType) {
+  TfType::Define<MjcPhysicsImageableAPI, TfType::Bases<UsdAPISchemaBase> >();
+}
+
+/* virtual */
+MjcPhysicsImageableAPI::~MjcPhysicsImageableAPI() {}
+
+/* static */
+MjcPhysicsImageableAPI MjcPhysicsImageableAPI::Get(const UsdStagePtr &stage,
+                                                   const SdfPath &path) {
+  if (!stage) {
+    TF_CODING_ERROR("Invalid stage");
+    return MjcPhysicsImageableAPI();
+  }
+  return MjcPhysicsImageableAPI(stage->GetPrimAtPath(path));
+}
+
+/* virtual */
+UsdSchemaKind MjcPhysicsImageableAPI::_GetSchemaKind() const {
+  return MjcPhysicsImageableAPI::schemaKind;
+}
+
+/* static */
+bool MjcPhysicsImageableAPI::CanApply(const UsdPrim &prim,
+                                      std::string *whyNot) {
+  return prim.CanApplyAPI<MjcPhysicsImageableAPI>(whyNot);
+}
+
+/* static */
+MjcPhysicsImageableAPI MjcPhysicsImageableAPI::Apply(const UsdPrim &prim) {
+  if (prim.ApplyAPI<MjcPhysicsImageableAPI>()) {
+    return MjcPhysicsImageableAPI(prim);
+  }
+  return MjcPhysicsImageableAPI();
+}
+
+/* static */
+const TfType &MjcPhysicsImageableAPI::_GetStaticTfType() {
+  static TfType tfType = TfType::Find<MjcPhysicsImageableAPI>();
+  return tfType;
+}
+
+/* static */
+bool MjcPhysicsImageableAPI::_IsTypedSchema() {
+  static bool isTyped = _GetStaticTfType().IsA<UsdTyped>();
+  return isTyped;
+}
+
+/* virtual */
+const TfType &MjcPhysicsImageableAPI::_GetTfType() const {
+  return _GetStaticTfType();
+}
+
+UsdAttribute MjcPhysicsImageableAPI::GetGroupAttr() const {
+  return GetPrim().GetAttribute(MjcPhysicsTokens->mjcGroup);
+}
+
+UsdAttribute MjcPhysicsImageableAPI::CreateGroupAttr(
+    VtValue const &defaultValue, bool writeSparsely) const {
+  return UsdSchemaBase::_CreateAttr(
+      MjcPhysicsTokens->mjcGroup, SdfValueTypeNames->Int,
+      /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely);
+}
+
+namespace {
+static inline TfTokenVector _ConcatenateAttributeNames(
+    const TfTokenVector &left, const TfTokenVector &right) {
+  TfTokenVector result;
+  result.reserve(left.size() + right.size());
+  result.insert(result.end(), left.begin(), left.end());
+  result.insert(result.end(), right.begin(), right.end());
+  return result;
+}
+}  // namespace
+
+/*static*/
+const TfTokenVector &MjcPhysicsImageableAPI::GetSchemaAttributeNames(
+    bool includeInherited) {
+  static TfTokenVector localNames = {
+      MjcPhysicsTokens->mjcGroup,
+  };
+  static TfTokenVector allNames = _ConcatenateAttributeNames(
+      UsdAPISchemaBase::GetSchemaAttributeNames(true), localNames);
+
+  if (includeInherited)
+    return allNames;
+  else
+    return localNames;
+}
+
+PXR_NAMESPACE_CLOSE_SCOPE
+
+// ===================================================================== //
+// Feel free to add custom code below this line. It will be preserved by
+// the code generator.
+//
+// Just remember to wrap code in the appropriate delimiters:
+// 'PXR_NAMESPACE_OPEN_SCOPE', 'PXR_NAMESPACE_CLOSE_SCOPE'.
+// ===================================================================== //
+// --(BEGIN CUSTOM CODE)--

src/experimental/usd/mjcPhysics/plugInfo.json

@@ -23,6 +23,16 @@
             ],
             "schemaKind": "singleApplyAPI"
           },
+          "MjcPhysicsImageableAPI": {
+            "alias": {
+              "UsdSchemaBase": "MjcImageableAPI"
+            },
+            "autoGenerated": true,
+            "bases": [
+              "UsdAPISchemaBase"
+            ],
+            "schemaKind": "singleApplyAPI"
+          },
           "MjcPhysicsJointAPI": {
             "alias": {
               "UsdSchemaBase": "MjcJointAPI"

src/experimental/usd/mjcPhysics/schema.usda

@@ -649,6 +649,25 @@ class "MjcSiteAPI"
     )
 }
 
+class "MjcImageableAPI"
+(
+    customData = {
+        string className = "ImageableAPI"
+    }
+    doc = """API describing attributes for visual entities in MuJoCo."""
+
+    inherits = </APISchemaBase>
+)
+{
+    uniform int mjc:group = 0 (
+      customData = {
+          string apiName = "Group"
+      }
+      displayName = "Group"
+      doc = """Integer MuJoCo group to which the imageable belongs."""
+    )
+}
+
 class "MjcCollisionAPI"
 (
     customData = {

src/experimental/usd/mjcPhysics/tokens.cpp

@@ -177,6 +177,7 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
       user("user", TfToken::Immortal),
       MjcActuator("MjcActuator", TfToken::Immortal),
       MjcCollisionAPI("MjcCollisionAPI", TfToken::Immortal),
+      MjcImageableAPI("MjcImageableAPI", TfToken::Immortal),
       MjcJointAPI("MjcJointAPI", TfToken::Immortal),
       MjcKeyframe("MjcKeyframe", TfToken::Immortal),
       MjcMaterialAPI("MjcMaterialAPI", TfToken::Immortal),
@@ -336,6 +337,7 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
                  user,
                  MjcActuator,
                  MjcCollisionAPI,
+                 MjcImageableAPI,
                  MjcJointAPI,
                  MjcKeyframe,
                  MjcMaterialAPI,

src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc

@@ -656,7 +656,6 @@ class ModelWriter {
       create_flag_attr(token, flag, false);
     }
 
-
     // Compiler attributes
     WriteUniformAttribute(physics_scene_path, pxr::SdfValueTypeNames->Bool,
                           MjcPhysicsTokens->mjcCompilerAutoLimits,
@@ -684,7 +683,8 @@ class ModelWriter {
 
     WriteUniformAttribute(physics_scene_path, pxr::SdfValueTypeNames->Token,
                           MjcPhysicsTokens->mjcCompilerAngle,
-                          spec_->compiler.degree ? MjcPhysicsTokens->degree:MjcPhysicsTokens->radian);
+                          spec_->compiler.degree ? MjcPhysicsTokens->degree
+                                                 : MjcPhysicsTokens->radian);
 
     WriteUniformAttribute(physics_scene_path, pxr::SdfValueTypeNames->Bool,
                           MjcPhysicsTokens->mjcCompilerFitAABB,
@@ -1621,6 +1621,13 @@ class ModelWriter {
     WriteUniformAttribute(geom_path, pxr::SdfValueTypeNames->Int,
                           MjcPhysicsTokens->mjcGroup, geom->group);
 
+    if (model_->geom_contype[geom_id] == 0 &&
+                           model_->geom_conaffinity[geom_id] == 0) {
+      // If the geom is purely visual, apply the imageable API.
+      ApplyApiSchema(data_, geom_path,
+                     MjcPhysicsTokens->MjcImageableAPI);
+    }
+
     // Apply the physics schemas if we are writing physics and the
     // geom participates in collisions.
     if (write_physics_ && (model_->geom_contype[geom_id] != 0 ||
@@ -1715,6 +1722,11 @@ class ModelWriter {
         SetAttributeDefault(data_, approximation_attr,
                             pxr::UsdPhysicsTokens->convexHull);
       }
+    } else {
+      // Currently imageable only has a group API. But since it's the same
+      // naming in MjcCollisionsAPI we've already set it earlier in this
+      // function.
+      ApplyApiSchema(data_, geom_path, MjcPhysicsTokens->MjcImageableAPI);
     }
 
     mjsDefault *spec_default = mjs_getDefault(geom->element);

src/experimental/usd/usd_to_mjspec.cc

@@ -21,6 +21,7 @@
 
 #include <mujoco/experimental/usd/mjcPhysics/actuator.h>
 #include <mujoco/experimental/usd/mjcPhysics/collisionAPI.h>
+#include <mujoco/experimental/usd/mjcPhysics/imageableAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/jointAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/keyframe.h>
 #include <mujoco/experimental/usd/mjcPhysics/materialAPI.h>
@@ -1335,6 +1336,7 @@ void ParseUsdGeomGprim(mjSpec* spec, const pxr::UsdPrim& gprim,
   if (!MaybeParseGeomPrimitive(gprim, geom, caches.xform_cache)) {
     ParseUsdMesh(spec, gprim, geom, caches.xform_cache);
   }
+
   pxr::UsdShadeMaterial bound_material =
       pxr::UsdShadeMaterialBindingAPI(gprim).ComputeBoundMaterial(
           &caches.bindings_cache, &caches.collection_query_cache);
@@ -1350,6 +1352,14 @@ void ParseUsdGeomGprim(mjSpec* spec, const pxr::UsdPrim& gprim,
     }
     mjs_setString(geom->material, mjs_getName(material->element)->c_str());
   }
+
+  if (gprim.HasAPI<pxr::MjcPhysicsImageableAPI>()) {
+    auto imageable_api = pxr::MjcPhysicsImageableAPI(gprim);
+    auto group_attr = imageable_api.GetGroupAttr();
+    if (group_attr.HasAuthoredValue()) {
+      group_attr.Get(&geom->group);
+    }
+  }
 }
 
 void ParseUsdPhysicsCollider(mjSpec* spec,
@@ -1389,9 +1399,10 @@ void ParseUsdPhysicsCollider(mjSpec* spec,
     }
   }
 
-  // Parse the Mass API after the physics material APIs since the density attribute
-  // from the Mass API is supposed to override the Material API density attribute.
-  // See https://openusd.org/dev/api/usd_physics_page_front.html
+  // Parse the Mass API after the physics material APIs since the density
+  // attribute from the Mass API is supposed to override the Material API
+  // density attribute. See
+  // https://openusd.org/dev/api/usd_physics_page_front.html
   if (prim.HasAPI<pxr::UsdPhysicsMassAPI>()) {
     ParseUsdPhysicsMassAPIForGeom(geom, pxr::UsdPhysicsMassAPI(prim));
   }
@@ -1672,8 +1683,9 @@ void PopulateSpecFromTree(pxr::UsdStageRefPtr stage, mjSpec* spec,
           : stage->GetPrimAtPath(current_node->body_path);
 
   for (const auto& gprim_path : current_node->visual_gprims) {
-    ParseUsdGeomGprim(spec, stage->GetPrimAtPath(gprim_path),
-                      body_prim_for_xform, current_mj_body, caches);
+    auto gprim = stage->GetPrimAtPath(gprim_path);
+    ParseUsdGeomGprim(spec, gprim, body_prim_for_xform, current_mj_body,
+                      caches);
   }
 
   for (const auto& collider_path : current_node->colliders) {